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

import postgres, { type Sql } from 'postgres';
import type { StorageProvider, IdentityKeyPair, SignedPreKey, OneTimePreKey, SessionState, RetiredIdentity, PersistedStreamState, PeerVerification, PeerVerificationSource } from '@shade/core';
import {
  toBase64, fromBase64,
  constantTimeEqual,
  serializeSessionState, deserializeSessionState,
  serializeSignedPreKey, deserializeSignedPreKey,
  serializeOneTimePreKey, deserializeOneTimePreKey,
  serializeIdentityKeyPair, deserializeIdentityKeyPair,
} from '@shade/core';
import { ensureClientTables } from './ensure-tables.js';

/**
 * PostgreSQL-backed StorageProvider for Shade client-side storage.
 *
 * Schema: tables prefixed `shade_` so they don't collide with existing
 * tables in projects that share the same PG instance (e.g. Orchestrator).
 *
 * Usage:
 * ```ts
 * const storage = await PostgresStorage.create('postgres://user:pass@host/db');
 * const manager = new ShadeSessionManager(crypto, storage);
 * ```
 */
export class PostgresStorage implements StorageProvider {
  private constructor(
    private readonly sql: Sql,
    private readonly ownsConnection: boolean,
  ) {}

  /** Create from a connection string (owns the connection) */
  static async create(connectionString: string): Promise<PostgresStorage> {
    const sql = postgres(connectionString);
    const storage = new PostgresStorage(sql, true);
    await ensureClientTables(sql);
    return storage;
  }

  /** Create from an existing postgres-js client (caller owns the connection) */
  static async fromClient(sql: Sql): Promise<PostgresStorage> {
    const storage = new PostgresStorage(sql, false);
    await ensureClientTables(sql);
    return storage;
  }

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

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

  async getIdentityKeyPair(): Promise<IdentityKeyPair | null> {
    const rows = await this.sql<Array<{ signing_public_key: string; signing_private_key: string; dh_public_key: string; dh_private_key: string }>>`
      SELECT signing_public_key, signing_private_key, dh_public_key, dh_private_key
      FROM shade_identity WHERE id = 1
    `;
    if (rows.length === 0) return null;
    const r = rows[0]!;
    return {
      signingPublicKey: fromBase64(r.signing_public_key),
      signingPrivateKey: fromBase64(r.signing_private_key),
      dhPublicKey: fromBase64(r.dh_public_key),
      dhPrivateKey: fromBase64(r.dh_private_key),
    };
  }

  async saveIdentityKeyPair(kp: IdentityKeyPair): Promise<void> {
    await this.sql`
      INSERT INTO shade_identity (id, signing_public_key, signing_private_key, dh_public_key, dh_private_key)
      VALUES (1, ${toBase64(kp.signingPublicKey)}, ${toBase64(kp.signingPrivateKey)}, ${toBase64(kp.dhPublicKey)}, ${toBase64(kp.dhPrivateKey)})
      ON CONFLICT (id) DO UPDATE SET
        signing_public_key = EXCLUDED.signing_public_key,
        signing_private_key = EXCLUDED.signing_private_key,
        dh_public_key = EXCLUDED.dh_public_key,
        dh_private_key = EXCLUDED.dh_private_key
    `;
  }

  async getLocalRegistrationId(): Promise<number> {
    const rows = await this.sql<Array<{ value: string }>>`
      SELECT value FROM shade_config WHERE key = 'registrationId'
    `;
    return rows.length ? parseInt(rows[0]!.value, 10) : 0;
  }

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

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

  async getSignedPreKey(keyId: number): Promise<SignedPreKey | null> {
    const rows = await this.sql<Array<{ data_json: string }>>`
      SELECT data_json FROM shade_signed_prekeys WHERE key_id = ${keyId}
    `;
    return rows.length ? deserializeSignedPreKey(rows[0]!.data_json) : null;
  }

  async saveSignedPreKey(key: SignedPreKey): Promise<void> {
    await this.sql`
      INSERT INTO shade_signed_prekeys (key_id, data_json)
      VALUES (${key.keyId}, ${serializeSignedPreKey(key)})
      ON CONFLICT (key_id) DO UPDATE SET data_json = EXCLUDED.data_json
    `;
  }

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

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

  async getOneTimePreKey(keyId: number): Promise<OneTimePreKey | null> {
    const rows = await this.sql<Array<{ data_json: string }>>`
      SELECT data_json FROM shade_one_time_prekeys WHERE key_id = ${keyId}
    `;
    return rows.length ? deserializeOneTimePreKey(rows[0]!.data_json) : null;
  }

  async saveOneTimePreKey(key: OneTimePreKey): Promise<void> {
    await this.sql`
      INSERT INTO shade_one_time_prekeys (key_id, data_json)
      VALUES (${key.keyId}, ${serializeOneTimePreKey(key)})
      ON CONFLICT (key_id) DO UPDATE SET data_json = EXCLUDED.data_json
    `;
  }

  async removeOneTimePreKey(keyId: number): Promise<void> {
    await this.sql`DELETE FROM shade_one_time_prekeys 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
    `;
    return parseInt(rows[0]!.count, 10);
  }

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

  async getSession(address: string): Promise<SessionState | null> {
    const rows = await this.sql<Array<{ state_json: string }>>`
      SELECT state_json FROM shade_sessions WHERE address = ${address}
    `;
    return rows.length ? deserializeSessionState(rows[0]!.state_json) : null;
  }

  async saveSession(address: string, state: SessionState): Promise<void> {
    await this.sql`
      INSERT INTO shade_sessions (address, state_json)
      VALUES (${address}, ${serializeSessionState(state)})
      ON CONFLICT (address) DO UPDATE SET state_json = EXCLUDED.state_json
    `;
  }

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

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

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

  async saveTrustedIdentity(address: string, identityKey: Uint8Array): Promise<void> {
    await this.sql`
      INSERT INTO shade_trusted_identities (address, identity_key)
      VALUES (${address}, ${toBase64(identityKey)})
      ON CONFLICT (address) DO UPDATE SET identity_key = EXCLUDED.identity_key
    `;
  }

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

  async addRetiredIdentity(identity: RetiredIdentity): Promise<void> {
    await this.sql`
      INSERT INTO shade_retired_identities (data_json, retired_at)
      VALUES (${serializeIdentityKeyPair(identity.keyPair)}, ${identity.retiredAt})
    `;
  }

  async getRetiredIdentities(): Promise<RetiredIdentity[]> {
    const rows = await this.sql<Array<{ data_json: string; retired_at: string }>>`
      SELECT data_json, retired_at FROM shade_retired_identities ORDER BY retired_at DESC
    `;
    return rows.map((r) => ({
      keyPair: deserializeIdentityKeyPair(r.data_json),
      retiredAt: Number(r.retired_at),
    }));
  }

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

  // ─── Stream-transfer resume state (v0.2.0) ────────────────

  async saveStreamState(state: PersistedStreamState): Promise<void> {
    await this.sql`
      INSERT INTO shade_stream_state (
        stream_id, direction, peer_address, status,
        metadata_json, partition_json, lane_state_json, io_descriptor_json,
        secret_enc, secret_nonce, overall_hash_state, created_at, updated_at
      ) VALUES (
        ${state.streamId}, ${state.direction}, ${state.peerAddress}, ${state.status},
        ${state.metadataJson}, ${state.partitionJson}, ${state.laneStateJson}, ${state.ioDescriptorJson},
        ${state.secretEnc}, ${state.secretNonce}, ${state.overallHashState ?? null},
        ${state.createdAt}, ${state.updatedAt}
      )
      ON CONFLICT (stream_id) DO UPDATE SET
        direction = EXCLUDED.direction,
        peer_address = EXCLUDED.peer_address,
        status = EXCLUDED.status,
        metadata_json = EXCLUDED.metadata_json,
        partition_json = EXCLUDED.partition_json,
        lane_state_json = EXCLUDED.lane_state_json,
        io_descriptor_json = EXCLUDED.io_descriptor_json,
        secret_enc = EXCLUDED.secret_enc,
        secret_nonce = EXCLUDED.secret_nonce,
        overall_hash_state = EXCLUDED.overall_hash_state,
        updated_at = EXCLUDED.updated_at
    `;
  }

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

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

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

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

  // ─── Peer verifications (V3.3) ────────────────────────────

  async savePeerVerification(v: PeerVerification): Promise<void> {
    await this.sql`
      INSERT INTO shade_peer_verifications
        (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 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 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 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 (peer_address, version)
      VALUES (${address}, ${next})
      ON CONFLICT (peer_address) DO UPDATE SET version = EXCLUDED.version
    `;
    return next;
  }
}

function rowToStreamState(row: Record<string, unknown>): PersistedStreamState {
  const out: PersistedStreamState = {
    streamId: String(row.stream_id),
    direction: row.direction as 'send' | 'receive',
    peerAddress: String(row.peer_address),
    status: row.status as 'active' | 'paused' | 'finished' | 'aborted',
    metadataJson: String(row.metadata_json),
    partitionJson: String(row.partition_json),
    laneStateJson: String(row.lane_state_json),
    ioDescriptorJson: String(row.io_descriptor_json),
    secretEnc: row.secret_enc as Uint8Array,
    secretNonce: row.secret_nonce as Uint8Array,
    createdAt: Number(row.created_at),
    updatedAt: Number(row.updated_at),
  };
  if (row.overall_hash_state !== null && row.overall_hash_state !== undefined) {
    out.overallHashState = String(row.overall_hash_state);
  }
  return out;
}