Shade/packages/shade-storage-encrypted/src/crypto/row-codec.ts

/**
 * Row codec — bridges between the StorageProvider's typed values and the
 * AEAD-sealed BLOB that lives in an `_enc` table.
 *
 * The codec is deliberately shared between SQLite and Postgres backends so
 * the wire format (and AAD binding) is identical across them. A backup
 * created with one backend can be re-keyed and restored under the other.
 */

import type {
  IdentityKeyPair,
  OneTimePreKey,
  PersistedStreamState,
  RetiredIdentity,
  SessionState,
  SignedPreKey,
} from '@shade/core';
import {
  serializeIdentityKeyPair, deserializeIdentityKeyPair,
  serializeOneTimePreKey, deserializeOneTimePreKey,
  serializeSessionState, deserializeSessionState,
  serializeSignedPreKey, deserializeSignedPreKey,
  toBase64, fromBase64,
} from '@shade/core';
import { aeadOpen, aeadSeal } from './aead.js';
import { buildAad, deriveNonce } from './kdf.js';
import type { KeyManager } from './key-manager.js';

const TEXT_ENCODER = new TextEncoder();
const TEXT_DECODER = new TextDecoder();

/** Logical column identifiers — used for fieldKey + AAD binding. */
export const COL = {
  identity: 'identity',
  config: 'config',
  signedPrekey: 'signed_prekey',
  oneTimePrekey: 'one_time_prekey',
  session: 'session',
  trustedIdentity: 'trusted_identity',
  retiredIdentity: 'retired_identity',
  streamSensitive: 'stream_sensitive',
} as const;

/** Logical table identifiers — used for fieldKey + AAD binding. */
export const TBL = {
  identity: 'identity',
  config: 'config',
  signedPrekeys: 'signed_prekeys',
  oneTimePrekeys: 'one_time_prekeys',
  sessions: 'sessions',
  trustedIdentities: 'trusted_identities',
  retiredIdentities: 'retired_identities',
  streamState: 'stream_state',
} as const;

/** Encrypt an arbitrary string payload bound to (table, column, pk). */
export async function sealString(
  km: KeyManager,
  table: string,
  column: string,
  pk: string,
  plaintext: string,
): Promise<Uint8Array> {
  const key = km.fieldKey(table, column);
  const nonce = deriveNonce(key, table, pk);
  const aad = buildAad(table, column, pk);
  return aeadSeal(key, nonce, TEXT_ENCODER.encode(plaintext), aad);
}

/** Decrypt a blob into a string, reconstructing AAD from (table, column, pk). */
export async function openString(
  km: KeyManager,
  table: string,
  column: string,
  pk: string,
  blob: Uint8Array,
): Promise<string> {
  const key = km.fieldKey(table, column);
  const expectedNonce = deriveNonce(key, table, pk);
  const aad = buildAad(table, column, pk);
  const pt = await aeadOpen(key, blob, aad, expectedNonce);
  return TEXT_DECODER.decode(pt);
}

/** Encrypt arbitrary bytes payload. */
export async function sealBytes(
  km: KeyManager,
  table: string,
  column: string,
  pk: string,
  plaintext: Uint8Array,
): Promise<Uint8Array> {
  const key = km.fieldKey(table, column);
  const nonce = deriveNonce(key, table, pk);
  const aad = buildAad(table, column, pk);
  return aeadSeal(key, nonce, plaintext, aad);
}

/** Decrypt arbitrary bytes payload. */
export async function openBytes(
  km: KeyManager,
  table: string,
  column: string,
  pk: string,
  blob: Uint8Array,
): Promise<Uint8Array> {
  const key = km.fieldKey(table, column);
  const expectedNonce = deriveNonce(key, table, pk);
  const aad = buildAad(table, column, pk);
  return aeadOpen(key, blob, aad, expectedNonce);
}

// ─── Typed encoders for each StorageProvider entity ──────────────────────

export async function sealIdentity(km: KeyManager, kp: IdentityKeyPair): Promise<Uint8Array> {
  return sealString(km, TBL.identity, COL.identity, '1', serializeIdentityKeyPair(kp));
}

export async function openIdentity(km: KeyManager, blob: Uint8Array): Promise<IdentityKeyPair> {
  return deserializeIdentityKeyPair(await openString(km, TBL.identity, COL.identity, '1', blob));
}

export async function sealConfig(km: KeyManager, key: string, value: string): Promise<Uint8Array> {
  return sealString(km, TBL.config, COL.config, key, value);
}

export async function openConfig(km: KeyManager, key: string, blob: Uint8Array): Promise<string> {
  return openString(km, TBL.config, COL.config, key, blob);
}

export async function sealSignedPreKey(km: KeyManager, k: SignedPreKey): Promise<Uint8Array> {
  return sealString(km, TBL.signedPrekeys, COL.signedPrekey, String(k.keyId), serializeSignedPreKey(k));
}

export async function openSignedPreKey(km: KeyManager, keyId: number, blob: Uint8Array): Promise<SignedPreKey> {
  return deserializeSignedPreKey(await openString(km, TBL.signedPrekeys, COL.signedPrekey, String(keyId), blob));
}

export async function sealOneTimePreKey(km: KeyManager, k: OneTimePreKey): Promise<Uint8Array> {
  return sealString(km, TBL.oneTimePrekeys, COL.oneTimePrekey, String(k.keyId), serializeOneTimePreKey(k));
}

export async function openOneTimePreKey(km: KeyManager, keyId: number, blob: Uint8Array): Promise<OneTimePreKey> {
  return deserializeOneTimePreKey(await openString(km, TBL.oneTimePrekeys, COL.oneTimePrekey, String(keyId), blob));
}

export async function sealSession(km: KeyManager, address: string, state: SessionState): Promise<Uint8Array> {
  return sealString(km, TBL.sessions, COL.session, address, serializeSessionState(state));
}

export async function openSession(km: KeyManager, address: string, blob: Uint8Array): Promise<SessionState> {
  return deserializeSessionState(await openString(km, TBL.sessions, COL.session, address, blob));
}

export async function sealTrust(km: KeyManager, address: string, identityKey: Uint8Array): Promise<Uint8Array> {
  return sealString(km, TBL.trustedIdentities, COL.trustedIdentity, address, toBase64(identityKey));
}

export async function openTrust(km: KeyManager, address: string, blob: Uint8Array): Promise<Uint8Array> {
  return fromBase64(await openString(km, TBL.trustedIdentities, COL.trustedIdentity, address, blob));
}

/**
 * Retired identities are append-only with auto-incremented row IDs. We bind
 * AAD on (retiredAt as string) since retired_at is a unique-enough natural
 * key for this row; collisions are practically impossible (ms timestamp).
 */
export async function sealRetired(km: KeyManager, ri: RetiredIdentity): Promise<Uint8Array> {
  const pk = String(ri.retiredAt);
  return sealString(km, TBL.retiredIdentities, COL.retiredIdentity, pk, serializeIdentityKeyPair(ri.keyPair));
}

export async function openRetired(km: KeyManager, retiredAt: number, blob: Uint8Array): Promise<RetiredIdentity> {
  const pk = String(retiredAt);
  const json = await openString(km, TBL.retiredIdentities, COL.retiredIdentity, pk, blob);
  return { keyPair: deserializeIdentityKeyPair(json), retiredAt };
}

/**
 * Stream-state sensitive bundle. Plaintext fields (peer, status, dir,
 * timestamps, streamId) stay in their own columns so the storage backend
 * can run efficient queries; everything else lives in this encrypted blob.
 */
interface StreamSensitiveBundle {
  metadataJson: string;
  partitionJson: string;
  laneStateJson: string;
  ioDescriptorJson: string;
  secretEnc: string;       // base64
  secretNonce: string;     // base64
  overallHashState?: string;
}

function packStreamSensitive(s: PersistedStreamState): StreamSensitiveBundle {
  const out: StreamSensitiveBundle = {
    metadataJson: s.metadataJson,
    partitionJson: s.partitionJson,
    laneStateJson: s.laneStateJson,
    ioDescriptorJson: s.ioDescriptorJson,
    secretEnc: toBase64(s.secretEnc),
    secretNonce: toBase64(s.secretNonce),
  };
  if (s.overallHashState !== undefined) out.overallHashState = s.overallHashState;
  return out;
}

export async function sealStreamSensitive(km: KeyManager, s: PersistedStreamState): Promise<Uint8Array> {
  return sealString(km, TBL.streamState, COL.streamSensitive, s.streamId, JSON.stringify(packStreamSensitive(s)));
}

export async function openStreamSensitive(
  km: KeyManager,
  streamId: string,
  blob: Uint8Array,
): Promise<Pick<PersistedStreamState, 'metadataJson' | 'partitionJson' | 'laneStateJson' | 'ioDescriptorJson' | 'secretEnc' | 'secretNonce' | 'overallHashState'>> {
  const json = await openString(km, TBL.streamState, COL.streamSensitive, streamId, blob);
  const b = JSON.parse(json) as StreamSensitiveBundle;
  const out = {
    metadataJson: b.metadataJson,
    partitionJson: b.partitionJson,
    laneStateJson: b.laneStateJson,
    ioDescriptorJson: b.ioDescriptorJson,
    secretEnc: fromBase64(b.secretEnc),
    secretNonce: fromBase64(b.secretNonce),
  } as Pick<PersistedStreamState, 'metadataJson' | 'partitionJson' | 'laneStateJson' | 'ioDescriptorJson' | 'secretEnc' | 'secretNonce' | 'overallHashState'>;
  if (b.overallHashState !== undefined) (out as { overallHashState?: string }).overallHashState = b.overallHashState;
  return out;
}