Shade/packages/shade-storage-encrypted/src/storage/encrypted-postgres.ts

import type { Sql } from 'postgres';
import postgres from 'postgres';
import type {
  IdentityKeyPair,
  OneTimePreKey,
  PeerVerification,
  PeerVerificationSource,
  PersistedStreamState,
  RetiredIdentity,
  SessionState,
  SignedPreKey,
  StorageProvider,
} from '@shade/core';
import { constantTimeEqual, toBase64 } from '@shade/core';
import { KeyManager } from '../crypto/key-manager.js';
import {
  openConfig, openIdentity, openOneTimePreKey, openRetired, openSession,
  openSignedPreKey, openStreamSensitive, openTrust, sealConfig, sealIdentity,
  sealOneTimePreKey, sealRetired, sealSession, sealSignedPreKey,
  sealStreamSensitive, sealTrust,
} from '../crypto/row-codec.js';

/**
 * PostgreSQL-backed StorageProvider with at-rest encryption (V3.2).
 *
 * Tables prefixed `shade_*_enc` to allow side-by-side migration with the
 * unencrypted `shade_*` tables. Sensitive payloads are sealed with
 * AES-256-GCM bound to (table, column, pk) AAD.
 */
export class EncryptedPostgresStorage implements StorageProvider {
  private constructor(
    private readonly sql: Sql,
    private readonly km: KeyManager,
    private readonly ownsConnection: boolean,
  ) {}

  /** Create from connection string (owns connection). */
  static async create(connectionString: string, km: KeyManager): Promise<EncryptedPostgresStorage> {
    const sql = postgres(connectionString);
    const store = new EncryptedPostgresStorage(sql, km, true);
    await ensureEncryptedClientTables(sql);
    await store.assertKeyMatchesOrPersistFingerprint();
    return store;
  }

  /** Wrap an existing postgres-js Sql client (caller owns it). */
  static async fromClient(sql: Sql, km: KeyManager): Promise<EncryptedPostgresStorage> {
    const store = new EncryptedPostgresStorage(sql, km, false);
    await ensureEncryptedClientTables(sql);
    await store.assertKeyMatchesOrPersistFingerprint();
    return store;
  }

  async close(): Promise<void> {
    if (this.ownsConnection) await this.sql.end();
    this.km.destroy();
  }

  private async assertKeyMatchesOrPersistFingerprint(): Promise<void> {
    const expected = toBase64(this.km.storageKeyFingerprint());
    const rows = await this.sql<Array<{ value: string }>>`
      SELECT value FROM shade_meta_enc WHERE key = 'storage_key_fingerprint'
    `;
    if (rows.length === 0) {
      await this.sql`
        INSERT INTO shade_meta_enc (key, value) VALUES ('storage_key_fingerprint', ${expected})
      `;
      return;
    }
    if (rows[0]!.value !== expected) {
      throw new Error(
        'storage key mismatch — the supplied passphrase / keychain entry does not unlock this database',
      );
    }
  }

  // ─── Identity ──────────────────────────────────────────────

  async getIdentityKeyPair(): Promise<IdentityKeyPair | null> {
    const rows = await this.sql<Array<{ ciphertext: Uint8Array }>>`
      SELECT ciphertext FROM shade_identity_enc WHERE id = 1
    `;
    if (rows.length === 0) return null;
    return openIdentity(this.km, rows[0]!.ciphertext);
  }

  async saveIdentityKeyPair(kp: IdentityKeyPair): Promise<void> {
    const blob = await sealIdentity(this.km, kp);
    await this.sql`
      INSERT INTO shade_identity_enc (id, ciphertext) VALUES (1, ${blob})
      ON CONFLICT (id) DO UPDATE SET ciphertext = EXCLUDED.ciphertext
    `;
  }

  async getLocalRegistrationId(): Promise<number> {
    const rows = await this.sql<Array<{ ciphertext: Uint8Array }>>`
      SELECT ciphertext FROM shade_config_enc WHERE key = 'registrationId'
    `;
    if (rows.length === 0) return 0;
    return parseInt(await openConfig(this.km, 'registrationId', rows[0]!.ciphertext), 10);
  }

  async saveLocalRegistrationId(id: number): Promise<void> {
    const blob = await sealConfig(this.km, 'registrationId', String(id));
    await this.sql`
      INSERT INTO shade_config_enc (key, ciphertext) VALUES ('registrationId', ${blob})
      ON CONFLICT (key) DO UPDATE SET ciphertext = EXCLUDED.ciphertext
    `;
  }

  // ─── Signed PreKeys ────────────────────────────────────────

  async getSignedPreKey(keyId: number): Promise<SignedPreKey | null> {
    const rows = await this.sql<Array<{ ciphertext: Uint8Array }>>`
      SELECT ciphertext FROM shade_signed_prekeys_enc WHERE key_id = ${keyId}
    `;
    if (rows.length === 0) return null;
    return openSignedPreKey(this.km, keyId, rows[0]!.ciphertext);
  }

  async saveSignedPreKey(key: SignedPreKey): Promise<void> {
    const blob = await sealSignedPreKey(this.km, key);
    await this.sql`
      INSERT INTO shade_signed_prekeys_enc (key_id, ciphertext) VALUES (${key.keyId}, ${blob})
      ON CONFLICT (key_id) DO UPDATE SET ciphertext = EXCLUDED.ciphertext
    `;
  }

  async removeSignedPreKey(keyId: number): Promise<void> {
    await this.sql`DELETE FROM shade_signed_prekeys_enc WHERE key_id = ${keyId}`;
  }

  // ─── One-Time PreKeys ──────────────────────────────────────

  async getOneTimePreKey(keyId: number): Promise<OneTimePreKey | null> {
    const rows = await this.sql<Array<{ ciphertext: Uint8Array }>>`
      SELECT ciphertext FROM shade_one_time_prekeys_enc WHERE key_id = ${keyId}
    `;
    if (rows.length === 0) return null;
    return openOneTimePreKey(this.km, keyId, rows[0]!.ciphertext);
  }

  async saveOneTimePreKey(key: OneTimePreKey): Promise<void> {
    const blob = await sealOneTimePreKey(this.km, key);
    await this.sql`
      INSERT INTO shade_one_time_prekeys_enc (key_id, ciphertext) VALUES (${key.keyId}, ${blob})
      ON CONFLICT (key_id) DO UPDATE SET ciphertext = EXCLUDED.ciphertext
    `;
  }

  async removeOneTimePreKey(keyId: number): Promise<void> {
    await this.sql`DELETE FROM shade_one_time_prekeys_enc WHERE key_id = ${keyId}`;
  }

  async getOneTimePreKeyCount(): Promise<number> {
    const rows = await this.sql<Array<{ count: string }>>`
      SELECT COUNT(*)::text as count FROM shade_one_time_prekeys_enc
    `;
    return parseInt(rows[0]!.count, 10);
  }

  // ─── Sessions ──────────────────────────────────────────────

  async getSession(address: string): Promise<SessionState | null> {
    const rows = await this.sql<Array<{ ciphertext: Uint8Array }>>`
      SELECT ciphertext FROM shade_sessions_enc WHERE address = ${address}
    `;
    if (rows.length === 0) return null;
    return openSession(this.km, address, rows[0]!.ciphertext);
  }

  async saveSession(address: string, state: SessionState): Promise<void> {
    const blob = await sealSession(this.km, address, state);
    await this.sql`
      INSERT INTO shade_sessions_enc (address, ciphertext) VALUES (${address}, ${blob})
      ON CONFLICT (address) DO UPDATE SET ciphertext = EXCLUDED.ciphertext
    `;
  }

  async removeSession(address: string): Promise<void> {
    await this.sql`DELETE FROM shade_sessions_enc WHERE address = ${address}`;
  }

  // ─── Trust ─────────────────────────────────────────────────

  async isTrustedIdentity(address: string, identityKey: Uint8Array): Promise<boolean> {
    const rows = await this.sql<Array<{ ciphertext: Uint8Array }>>`
      SELECT ciphertext FROM shade_trusted_identities_enc WHERE address = ${address}
    `;
    if (rows.length === 0) return true; // TOFU
    const stored = await openTrust(this.km, address, rows[0]!.ciphertext);
    return constantTimeEqual(stored, identityKey);
  }

  async saveTrustedIdentity(address: string, identityKey: Uint8Array): Promise<void> {
    const blob = await sealTrust(this.km, address, identityKey);
    await this.sql`
      INSERT INTO shade_trusted_identities_enc (address, ciphertext) VALUES (${address}, ${blob})
      ON CONFLICT (address) DO UPDATE SET ciphertext = EXCLUDED.ciphertext
    `;
  }

  // ─── Identity History ──────────────────────────────────────

  async addRetiredIdentity(identity: RetiredIdentity): Promise<void> {
    const blob = await sealRetired(this.km, identity);
    await this.sql`
      INSERT INTO shade_retired_identities_enc (retired_at, ciphertext)
      VALUES (${identity.retiredAt}, ${blob})
      ON CONFLICT (retired_at) DO UPDATE SET ciphertext = EXCLUDED.ciphertext
    `;
  }

  async getRetiredIdentities(): Promise<RetiredIdentity[]> {
    const rows = await this.sql<Array<{ retired_at: string; ciphertext: Uint8Array }>>`
      SELECT retired_at, ciphertext FROM shade_retired_identities_enc ORDER BY retired_at DESC
    `;
    return Promise.all(rows.map((r) => openRetired(this.km, Number(r.retired_at), r.ciphertext)));
  }

  async pruneRetiredIdentities(olderThan: number): Promise<void> {
    await this.sql`DELETE FROM shade_retired_identities_enc WHERE retired_at < ${olderThan}`;
  }

  // ─── Stream-transfer resume state ──────────────────────────

  async saveStreamState(state: PersistedStreamState): Promise<void> {
    const blob = await sealStreamSensitive(this.km, state);
    await this.sql`
      INSERT INTO shade_stream_state_enc (
        stream_id, direction, peer_address, status, ciphertext, created_at, updated_at
      ) VALUES (
        ${state.streamId}, ${state.direction}, ${state.peerAddress}, ${state.status},
        ${blob}, ${state.createdAt}, ${state.updatedAt}
      )
      ON CONFLICT (stream_id) DO UPDATE SET
        direction    = EXCLUDED.direction,
        peer_address = EXCLUDED.peer_address,
        status       = EXCLUDED.status,
        ciphertext   = EXCLUDED.ciphertext,
        updated_at   = EXCLUDED.updated_at
    `;
  }

  async getStreamState(streamId: string): Promise<PersistedStreamState | null> {
    const rows = await this.sql<Array<StreamRow>>`
      SELECT * FROM shade_stream_state_enc WHERE stream_id = ${streamId}
    `;
    if (rows.length === 0) return null;
    return this.rowToStreamState(rows[0]!);
  }

  async removeStreamState(streamId: string): Promise<void> {
    await this.sql`DELETE FROM shade_stream_state_enc WHERE stream_id = ${streamId}`;
  }

  async listActiveStreamStates(direction?: 'send' | 'receive'): Promise<PersistedStreamState[]> {
    const rows =
      direction === undefined
        ? await this.sql<Array<StreamRow>>`
            SELECT * FROM shade_stream_state_enc
            WHERE status IN ('active','paused')
            ORDER BY updated_at DESC
          `
        : await this.sql<Array<StreamRow>>`
            SELECT * FROM shade_stream_state_enc
            WHERE status IN ('active','paused') AND direction = ${direction}
            ORDER BY updated_at DESC
          `;
    return Promise.all(rows.map((r) => this.rowToStreamState(r)));
  }

  async pruneStreamStates(olderThan: number): Promise<void> {
    await this.sql`
      DELETE FROM shade_stream_state_enc
      WHERE status IN ('finished','aborted') AND updated_at < ${olderThan}
    `;
  }

  // ─── Peer verifications (V3.3) ────────────────────────────
  // Fingerprints are public-by-design (intended for OOB display), so we
  // keep them plaintext alongside the encrypted tables for symmetry.

  async savePeerVerification(v: PeerVerification): Promise<void> {
    await this.sql`
      INSERT INTO shade_peer_verifications_enc
        (peer_address, fingerprint, verified_at, verified_by, identity_version)
      VALUES (${v.peerAddress}, ${v.fingerprint}, ${v.verifiedAt}, ${v.verifiedBy}, ${v.identityVersion})
      ON CONFLICT (peer_address) DO UPDATE SET
        fingerprint = EXCLUDED.fingerprint,
        verified_at = EXCLUDED.verified_at,
        verified_by = EXCLUDED.verified_by,
        identity_version = EXCLUDED.identity_version
    `;
  }

  async getPeerVerification(address: string): Promise<PeerVerification | null> {
    const rows = await this.sql<Array<{ peer_address: string; fingerprint: string; verified_at: string; verified_by: string; identity_version: string }>>`
      SELECT peer_address, fingerprint, verified_at, verified_by, identity_version
      FROM shade_peer_verifications_enc WHERE peer_address = ${address}
    `;
    if (rows.length === 0) return null;
    const r = rows[0]!;
    return {
      peerAddress: r.peer_address,
      fingerprint: r.fingerprint,
      verifiedAt: Number(r.verified_at),
      verifiedBy: r.verified_by as PeerVerificationSource,
      identityVersion: Number(r.identity_version),
    };
  }

  async removePeerVerification(address: string): Promise<void> {
    await this.sql`DELETE FROM shade_peer_verifications_enc WHERE peer_address = ${address}`;
  }

  async getPeerIdentityVersion(address: string): Promise<number> {
    const rows = await this.sql<Array<{ version: string }>>`
      SELECT version FROM shade_peer_identity_versions_enc WHERE peer_address = ${address}
    `;
    return rows.length ? Number(rows[0]!.version) : 1;
  }

  async bumpPeerIdentityVersion(address: string): Promise<number> {
    const current = await this.getPeerIdentityVersion(address);
    const next = current + 1;
    await this.sql`
      INSERT INTO shade_peer_identity_versions_enc (peer_address, version)
      VALUES (${address}, ${next})
      ON CONFLICT (peer_address) DO UPDATE SET version = EXCLUDED.version
    `;
    return next;
  }

  private async rowToStreamState(row: StreamRow): Promise<PersistedStreamState> {
    const sensitive = await openStreamSensitive(this.km, String(row.stream_id), row.ciphertext);
    const out: PersistedStreamState = {
      streamId: String(row.stream_id),
      direction: row.direction,
      peerAddress: String(row.peer_address),
      status: row.status,
      metadataJson: sensitive.metadataJson,
      partitionJson: sensitive.partitionJson,
      laneStateJson: sensitive.laneStateJson,
      ioDescriptorJson: sensitive.ioDescriptorJson,
      secretEnc: sensitive.secretEnc,
      secretNonce: sensitive.secretNonce,
      createdAt: Number(row.created_at),
      updatedAt: Number(row.updated_at),
    };
    if (sensitive.overallHashState !== undefined) out.overallHashState = sensitive.overallHashState;
    return out;
  }
}

interface StreamRow {
  stream_id: string;
  direction: 'send' | 'receive';
  peer_address: string;
  status: 'active' | 'paused' | 'finished' | 'aborted';
  ciphertext: Uint8Array;
  created_at: string | number;
  updated_at: string | number;
}

export async function ensureEncryptedClientTables(sql: Sql): Promise<void> {
  await sql`
    CREATE TABLE IF NOT EXISTS shade_meta_enc (
      key   TEXT PRIMARY KEY,
      value TEXT NOT NULL
    )
  `;
  await sql`
    CREATE TABLE IF NOT EXISTS shade_identity_enc (
      id INTEGER PRIMARY KEY CHECK (id = 1),
      ciphertext BYTEA NOT NULL
    )
  `;
  await sql`
    CREATE TABLE IF NOT EXISTS shade_config_enc (
      key TEXT PRIMARY KEY,
      ciphertext BYTEA NOT NULL
    )
  `;
  await sql`
    CREATE TABLE IF NOT EXISTS shade_signed_prekeys_enc (
      key_id INTEGER PRIMARY KEY,
      ciphertext BYTEA NOT NULL
    )
  `;
  await sql`
    CREATE TABLE IF NOT EXISTS shade_one_time_prekeys_enc (
      key_id INTEGER PRIMARY KEY,
      ciphertext BYTEA NOT NULL
    )
  `;
  await sql`
    CREATE TABLE IF NOT EXISTS shade_sessions_enc (
      address TEXT PRIMARY KEY,
      ciphertext BYTEA NOT NULL
    )
  `;
  await sql`
    CREATE TABLE IF NOT EXISTS shade_trusted_identities_enc (
      address TEXT PRIMARY KEY,
      ciphertext BYTEA NOT NULL
    )
  `;
  await sql`
    CREATE TABLE IF NOT EXISTS shade_retired_identities_enc (
      retired_at BIGINT PRIMARY KEY,
      ciphertext BYTEA NOT NULL
    )
  `;
  await sql`
    CREATE INDEX IF NOT EXISTS shade_retired_at_enc_idx
    ON shade_retired_identities_enc(retired_at)
  `;
  await sql`
    CREATE TABLE IF NOT EXISTS shade_stream_state_enc (
      stream_id    TEXT PRIMARY KEY,
      direction    TEXT NOT NULL CHECK (direction IN ('send','receive')),
      peer_address TEXT NOT NULL,
      status       TEXT NOT NULL CHECK (status IN ('active','paused','finished','aborted')),
      ciphertext   BYTEA NOT NULL,
      created_at   BIGINT NOT NULL,
      updated_at   BIGINT NOT NULL
    )
  `;
  await sql`
    CREATE INDEX IF NOT EXISTS shade_stream_enc_peer_idx
    ON shade_stream_state_enc(peer_address)
  `;
  await sql`
    CREATE INDEX IF NOT EXISTS shade_stream_enc_updated_idx
    ON shade_stream_state_enc(updated_at)
  `;
  await sql`
    CREATE INDEX IF NOT EXISTS shade_stream_enc_status_idx
    ON shade_stream_state_enc(status, direction)
  `;
  await sql`
    CREATE TABLE IF NOT EXISTS shade_peer_verifications_enc (
      peer_address     TEXT PRIMARY KEY,
      fingerprint      TEXT NOT NULL,
      verified_at      BIGINT NOT NULL,
      verified_by      TEXT NOT NULL,
      identity_version BIGINT NOT NULL
    )
  `;
  await sql`
    CREATE TABLE IF NOT EXISTS shade_peer_identity_versions_enc (
      peer_address TEXT PRIMARY KEY,
      version      BIGINT NOT NULL
    )
  `;
}