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

import postgres, { type Sql } from 'postgres';
import type { StorageProvider, IdentityKeyPair, SignedPreKey, OneTimePreKey, SessionState, RetiredIdentity } 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}`;
  }
}