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

import { Database } from 'bun:sqlite';
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';

/**
 * SQLite-backed StorageProvider with at-rest encryption (V3.2).
 *
 * Schema: parallel `_enc` tables alongside the unencrypted ones, so a
 * migration can run side-by-side and atomic-rename at the end. Sensitive
 * payloads are sealed with AES-256-GCM bound to (table, column, pk) AAD;
 * routing/timestamp fields stay plaintext to keep queries efficient.
 *
 * Bring your own KeyManager — see `KeyManager.open({ kind: 'passphrase' | 'keychain' | 'injected' })`.
 */
export class EncryptedSQLiteStorage implements StorageProvider {
  private readonly db: Database;
  private readonly km: KeyManager;
  private readonly ownsDb: boolean;

  // Prepared statements
  private stmts!: {
    getIdentity:                ReturnType<Database['prepare']>;
    saveIdentity:               ReturnType<Database['prepare']>;
    getConfig:                  ReturnType<Database['prepare']>;
    saveConfig:                 ReturnType<Database['prepare']>;
    getSignedPreKey:            ReturnType<Database['prepare']>;
    saveSignedPreKey:           ReturnType<Database['prepare']>;
    removeSignedPreKey:         ReturnType<Database['prepare']>;
    getOneTimePreKey:           ReturnType<Database['prepare']>;
    saveOneTimePreKey:          ReturnType<Database['prepare']>;
    removeOneTimePreKey:        ReturnType<Database['prepare']>;
    countOneTimePreKeys:        ReturnType<Database['prepare']>;
    getSession:                 ReturnType<Database['prepare']>;
    saveSession:                ReturnType<Database['prepare']>;
    removeSession:              ReturnType<Database['prepare']>;
    getTrust:                   ReturnType<Database['prepare']>;
    saveTrust:                  ReturnType<Database['prepare']>;
    addRetired:                 ReturnType<Database['prepare']>;
    listRetired:                ReturnType<Database['prepare']>;
    pruneRetired:               ReturnType<Database['prepare']>;
    saveStreamState:            ReturnType<Database['prepare']>;
    getStreamState:             ReturnType<Database['prepare']>;
    removeStreamState:          ReturnType<Database['prepare']>;
    listActiveStreamStates:     ReturnType<Database['prepare']>;
    listActiveByDirection:      ReturnType<Database['prepare']>;
    pruneStreamStates:          ReturnType<Database['prepare']>;
    getMeta:                    ReturnType<Database['prepare']>;
    setMeta:                    ReturnType<Database['prepare']>;
    savePeerVerification:       ReturnType<Database['prepare']>;
    getPeerVerification:        ReturnType<Database['prepare']>;
    removePeerVerification:     ReturnType<Database['prepare']>;
    getPeerIdentityVersion:     ReturnType<Database['prepare']>;
    upsertPeerIdentityVersion:  ReturnType<Database['prepare']>;
  };

  private constructor(db: Database, km: KeyManager, ownsDb: boolean) {
    this.db = db;
    this.km = km;
    this.ownsDb = ownsDb;
    this.ensureTables();
    this.prepareStatements();
  }

  /**
   * Open an encrypted SQLite store. The caller supplies the KeyManager
   * (so they control the key source) and the DB path.
   */
  static async open(opts: { dbPath?: string; keyManager: KeyManager }): Promise<EncryptedSQLiteStorage> {
    const path = opts.dbPath ?? process.env.SHADE_DB_PATH ?? '/data/shade-client.db';
    const db = new Database(path, { create: true });
    db.exec('PRAGMA journal_mode=WAL');
    const store = new EncryptedSQLiteStorage(db, opts.keyManager, true);
    await store.assertKeyMatchesOrPersistFingerprint();
    return store;
  }

  /** Wrap an existing bun:sqlite Database (caller owns it). */
  static async wrap(db: Database, km: KeyManager): Promise<EncryptedSQLiteStorage> {
    const store = new EncryptedSQLiteStorage(db, km, false);
    await store.assertKeyMatchesOrPersistFingerprint();
    return store;
  }

  private ensureTables() {
    this.db.exec(`
      CREATE TABLE IF NOT EXISTS shade_meta_enc (
        key   TEXT PRIMARY KEY,
        value TEXT NOT NULL
      );
      CREATE TABLE IF NOT EXISTS identity_enc (
        id INTEGER PRIMARY KEY CHECK (id = 1),
        ciphertext BLOB NOT NULL
      );
      CREATE TABLE IF NOT EXISTS config_enc (
        key TEXT PRIMARY KEY,
        ciphertext BLOB NOT NULL
      );
      CREATE TABLE IF NOT EXISTS signed_prekeys_enc (
        key_id INTEGER PRIMARY KEY,
        ciphertext BLOB NOT NULL
      );
      CREATE TABLE IF NOT EXISTS one_time_prekeys_enc (
        key_id INTEGER PRIMARY KEY,
        ciphertext BLOB NOT NULL
      );
      CREATE TABLE IF NOT EXISTS sessions_enc (
        address TEXT PRIMARY KEY,
        ciphertext BLOB NOT NULL
      );
      CREATE TABLE IF NOT EXISTS trusted_identities_enc (
        address TEXT PRIMARY KEY,
        ciphertext BLOB NOT NULL
      );
      CREATE TABLE IF NOT EXISTS retired_identities_enc (
        retired_at INTEGER PRIMARY KEY,
        ciphertext BLOB NOT NULL
      );
      CREATE INDEX IF NOT EXISTS idx_retired_at_enc ON retired_identities_enc(retired_at);
      CREATE TABLE IF NOT EXISTS stream_state_enc (
        stream_id      TEXT PRIMARY KEY,
        direction      TEXT NOT NULL,
        peer_address   TEXT NOT NULL,
        status         TEXT NOT NULL,
        ciphertext     BLOB NOT NULL,
        created_at     INTEGER NOT NULL,
        updated_at     INTEGER NOT NULL
      );
      CREATE INDEX IF NOT EXISTS idx_stream_enc_peer    ON stream_state_enc(peer_address);
      CREATE INDEX IF NOT EXISTS idx_stream_enc_updated ON stream_state_enc(updated_at);
      CREATE INDEX IF NOT EXISTS idx_stream_enc_status  ON stream_state_enc(status, direction);
      CREATE TABLE IF NOT EXISTS peer_verifications_enc (
        peer_address     TEXT PRIMARY KEY,
        fingerprint      TEXT NOT NULL,
        verified_at      INTEGER NOT NULL,
        verified_by      TEXT NOT NULL,
        identity_version INTEGER NOT NULL
      );
      CREATE TABLE IF NOT EXISTS peer_identity_versions_enc (
        peer_address TEXT PRIMARY KEY,
        version      INTEGER NOT NULL
      );
    `);
  }

  private prepareStatements() {
    this.stmts = {
      getIdentity:         this.db.prepare('SELECT ciphertext FROM identity_enc WHERE id = 1'),
      saveIdentity:        this.db.prepare('INSERT OR REPLACE INTO identity_enc (id, ciphertext) VALUES (1, ?)'),
      getConfig:           this.db.prepare('SELECT ciphertext FROM config_enc WHERE key = ?'),
      saveConfig:          this.db.prepare('INSERT OR REPLACE INTO config_enc (key, ciphertext) VALUES (?, ?)'),
      getSignedPreKey:     this.db.prepare('SELECT ciphertext FROM signed_prekeys_enc WHERE key_id = ?'),
      saveSignedPreKey:    this.db.prepare('INSERT OR REPLACE INTO signed_prekeys_enc (key_id, ciphertext) VALUES (?, ?)'),
      removeSignedPreKey:  this.db.prepare('DELETE FROM signed_prekeys_enc WHERE key_id = ?'),
      getOneTimePreKey:    this.db.prepare('SELECT ciphertext FROM one_time_prekeys_enc WHERE key_id = ?'),
      saveOneTimePreKey:   this.db.prepare('INSERT OR REPLACE INTO one_time_prekeys_enc (key_id, ciphertext) VALUES (?, ?)'),
      removeOneTimePreKey: this.db.prepare('DELETE FROM one_time_prekeys_enc WHERE key_id = ?'),
      countOneTimePreKeys: this.db.prepare('SELECT COUNT(*) as count FROM one_time_prekeys_enc'),
      getSession:          this.db.prepare('SELECT ciphertext FROM sessions_enc WHERE address = ?'),
      saveSession:         this.db.prepare('INSERT OR REPLACE INTO sessions_enc (address, ciphertext) VALUES (?, ?)'),
      removeSession:       this.db.prepare('DELETE FROM sessions_enc WHERE address = ?'),
      getTrust:            this.db.prepare('SELECT ciphertext FROM trusted_identities_enc WHERE address = ?'),
      saveTrust:           this.db.prepare('INSERT OR REPLACE INTO trusted_identities_enc (address, ciphertext) VALUES (?, ?)'),
      addRetired:          this.db.prepare('INSERT OR REPLACE INTO retired_identities_enc (retired_at, ciphertext) VALUES (?, ?)'),
      listRetired:         this.db.prepare('SELECT retired_at, ciphertext FROM retired_identities_enc ORDER BY retired_at DESC'),
      pruneRetired:        this.db.prepare('DELETE FROM retired_identities_enc WHERE retired_at < ?'),
      saveStreamState:     this.db.prepare(
        `INSERT OR REPLACE INTO stream_state_enc (
           stream_id, direction, peer_address, status, ciphertext, created_at, updated_at
         ) VALUES (?, ?, ?, ?, ?, ?, ?)`
      ),
      getStreamState:      this.db.prepare('SELECT * FROM stream_state_enc WHERE stream_id = ?'),
      removeStreamState:   this.db.prepare('DELETE FROM stream_state_enc WHERE stream_id = ?'),
      listActiveStreamStates: this.db.prepare(
        "SELECT * FROM stream_state_enc WHERE status IN ('active','paused') ORDER BY updated_at DESC"
      ),
      listActiveByDirection: this.db.prepare(
        "SELECT * FROM stream_state_enc WHERE status IN ('active','paused') AND direction = ? ORDER BY updated_at DESC"
      ),
      pruneStreamStates:   this.db.prepare(
        "DELETE FROM stream_state_enc WHERE status IN ('finished','aborted') AND updated_at < ?"
      ),
      getMeta:             this.db.prepare('SELECT value FROM shade_meta_enc WHERE key = ?'),
      setMeta:             this.db.prepare('INSERT OR REPLACE INTO shade_meta_enc (key, value) VALUES (?, ?)'),
      savePeerVerification: this.db.prepare(
        `INSERT OR REPLACE INTO peer_verifications_enc
           (peer_address, fingerprint, verified_at, verified_by, identity_version)
         VALUES (?, ?, ?, ?, ?)`,
      ),
      getPeerVerification: this.db.prepare(
        'SELECT peer_address, fingerprint, verified_at, verified_by, identity_version FROM peer_verifications_enc WHERE peer_address = ?',
      ),
      removePeerVerification: this.db.prepare('DELETE FROM peer_verifications_enc WHERE peer_address = ?'),
      getPeerIdentityVersion: this.db.prepare('SELECT version FROM peer_identity_versions_enc WHERE peer_address = ?'),
      upsertPeerIdentityVersion: this.db.prepare(
        `INSERT INTO peer_identity_versions_enc (peer_address, version) VALUES (?, ?)
         ON CONFLICT(peer_address) DO UPDATE SET version = excluded.version`,
      ),
    };
  }

  /**
   * On first open, persist a fingerprint of the storageKey. On subsequent
   * opens, compare and reject mismatches with a clear error rather than
   * silently writing data under the wrong key.
   */
  private async assertKeyMatchesOrPersistFingerprint(): Promise<void> {
    const expected = toBase64(this.km.storageKeyFingerprint());
    const row = this.stmts.getMeta.get('storage_key_fingerprint') as { value: string } | undefined;
    if (!row) {
      this.stmts.setMeta.run('storage_key_fingerprint', expected);
      return;
    }
    if (row.value !== expected) {
      throw new Error(
        'storage key mismatch — the supplied passphrase / keychain entry does not unlock this database',
      );
    }
  }

  close(): void {
    if (this.ownsDb) this.db.close();
    this.km.destroy();
  }

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

  async getIdentityKeyPair(): Promise<IdentityKeyPair | null> {
    const row = this.stmts.getIdentity.get() as { ciphertext: Uint8Array | ArrayBuffer } | undefined;
    if (!row) return null;
    return openIdentity(this.km, toBytes(row.ciphertext));
  }

  async saveIdentityKeyPair(kp: IdentityKeyPair): Promise<void> {
    const blob = await sealIdentity(this.km, kp);
    this.stmts.saveIdentity.run(blob);
  }

  async getLocalRegistrationId(): Promise<number> {
    const row = this.stmts.getConfig.get('registrationId') as { ciphertext: Uint8Array | ArrayBuffer } | undefined;
    if (!row) return 0;
    const v = await openConfig(this.km, 'registrationId', toBytes(row.ciphertext));
    return parseInt(v, 10);
  }

  async saveLocalRegistrationId(id: number): Promise<void> {
    const blob = await sealConfig(this.km, 'registrationId', String(id));
    this.stmts.saveConfig.run('registrationId', blob);
  }

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

  async getSignedPreKey(keyId: number): Promise<SignedPreKey | null> {
    const row = this.stmts.getSignedPreKey.get(keyId) as { ciphertext: Uint8Array | ArrayBuffer } | undefined;
    if (!row) return null;
    return openSignedPreKey(this.km, keyId, toBytes(row.ciphertext));
  }

  async saveSignedPreKey(key: SignedPreKey): Promise<void> {
    const blob = await sealSignedPreKey(this.km, key);
    this.stmts.saveSignedPreKey.run(key.keyId, blob);
  }

  async removeSignedPreKey(keyId: number): Promise<void> {
    this.stmts.removeSignedPreKey.run(keyId);
  }

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

  async getOneTimePreKey(keyId: number): Promise<OneTimePreKey | null> {
    const row = this.stmts.getOneTimePreKey.get(keyId) as { ciphertext: Uint8Array | ArrayBuffer } | undefined;
    if (!row) return null;
    return openOneTimePreKey(this.km, keyId, toBytes(row.ciphertext));
  }

  async saveOneTimePreKey(key: OneTimePreKey): Promise<void> {
    const blob = await sealOneTimePreKey(this.km, key);
    this.stmts.saveOneTimePreKey.run(key.keyId, blob);
  }

  async removeOneTimePreKey(keyId: number): Promise<void> {
    this.stmts.removeOneTimePreKey.run(keyId);
  }

  async getOneTimePreKeyCount(): Promise<number> {
    const row = this.stmts.countOneTimePreKeys.get() as { count: number };
    return row.count;
  }

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

  async getSession(address: string): Promise<SessionState | null> {
    const row = this.stmts.getSession.get(address) as { ciphertext: Uint8Array | ArrayBuffer } | undefined;
    if (!row) return null;
    return openSession(this.km, address, toBytes(row.ciphertext));
  }

  async saveSession(address: string, state: SessionState): Promise<void> {
    const blob = await sealSession(this.km, address, state);
    this.stmts.saveSession.run(address, blob);
  }

  async removeSession(address: string): Promise<void> {
    this.stmts.removeSession.run(address);
  }

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

  async isTrustedIdentity(address: string, identityKey: Uint8Array): Promise<boolean> {
    const row = this.stmts.getTrust.get(address) as { ciphertext: Uint8Array | ArrayBuffer } | undefined;
    if (!row) return true; // TOFU
    const stored = await openTrust(this.km, address, toBytes(row.ciphertext));
    return constantTimeEqual(stored, identityKey);
  }

  async saveTrustedIdentity(address: string, identityKey: Uint8Array): Promise<void> {
    const blob = await sealTrust(this.km, address, identityKey);
    this.stmts.saveTrust.run(address, blob);
  }

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

  async addRetiredIdentity(identity: RetiredIdentity): Promise<void> {
    const blob = await sealRetired(this.km, identity);
    this.stmts.addRetired.run(identity.retiredAt, blob);
  }

  async getRetiredIdentities(): Promise<RetiredIdentity[]> {
    const rows = this.stmts.listRetired.all() as { retired_at: number; ciphertext: Uint8Array | ArrayBuffer }[];
    return Promise.all(rows.map((r) => openRetired(this.km, Number(r.retired_at), toBytes(r.ciphertext))));
  }

  async pruneRetiredIdentities(olderThan: number): Promise<void> {
    this.stmts.pruneRetired.run(olderThan);
  }

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

  async saveStreamState(state: PersistedStreamState): Promise<void> {
    const blob = await sealStreamSensitive(this.km, state);
    this.stmts.saveStreamState.run(
      state.streamId,
      state.direction,
      state.peerAddress,
      state.status,
      blob,
      state.createdAt,
      state.updatedAt,
    );
  }

  async getStreamState(streamId: string): Promise<PersistedStreamState | null> {
    const row = this.stmts.getStreamState.get(streamId) as StreamRow | undefined;
    if (!row) return null;
    return this.rowToStreamState(row);
  }

  async removeStreamState(streamId: string): Promise<void> {
    this.stmts.removeStreamState.run(streamId);
  }

  async listActiveStreamStates(direction?: 'send' | 'receive'): Promise<PersistedStreamState[]> {
    const rows = (
      direction === undefined
        ? (this.stmts.listActiveStreamStates.all() as StreamRow[])
        : (this.stmts.listActiveByDirection.all(direction) as StreamRow[])
    );
    return Promise.all(rows.map((r) => this.rowToStreamState(r)));
  }

  async pruneStreamStates(olderThan: number): Promise<void> {
    this.stmts.pruneStreamStates.run(olderThan);
  }

  // ─── Peer verifications (V3.3) ────────────────────────────
  // Fingerprints are public-by-design; stored in plaintext for symmetry
  // with the unencrypted backend.

  async savePeerVerification(v: PeerVerification): Promise<void> {
    this.stmts.savePeerVerification.run(
      v.peerAddress,
      v.fingerprint,
      v.verifiedAt,
      v.verifiedBy,
      v.identityVersion,
    );
  }

  async getPeerVerification(address: string): Promise<PeerVerification | null> {
    const row = this.stmts.getPeerVerification.get(address) as
      | { peer_address: string; fingerprint: string; verified_at: number | bigint; verified_by: string; identity_version: number | bigint }
      | undefined;
    if (!row) return null;
    return {
      peerAddress: row.peer_address,
      fingerprint: row.fingerprint,
      verifiedAt: Number(row.verified_at),
      verifiedBy: row.verified_by as PeerVerificationSource,
      identityVersion: Number(row.identity_version),
    };
  }

  async removePeerVerification(address: string): Promise<void> {
    this.stmts.removePeerVerification.run(address);
  }

  async getPeerIdentityVersion(address: string): Promise<number> {
    const row = this.stmts.getPeerIdentityVersion.get(address) as { version: number | bigint } | undefined;
    return row ? Number(row.version) : 1;
  }

  async bumpPeerIdentityVersion(address: string): Promise<number> {
    const current = await this.getPeerIdentityVersion(address);
    const next = current + 1;
    this.stmts.upsertPeerIdentityVersion.run(address, next);
    return next;
  }

  private async rowToStreamState(row: StreamRow): Promise<PersistedStreamState> {
    const sensitive = await openStreamSensitive(this.km, row.stream_id, toBytes(row.ciphertext));
    const out: PersistedStreamState = {
      streamId: row.stream_id,
      direction: row.direction,
      peerAddress: 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 | ArrayBuffer;
  created_at: number | bigint;
  updated_at: number | bigint;
}

function toBytes(value: Uint8Array | ArrayBuffer | unknown): Uint8Array {
  if (value instanceof Uint8Array) return value;
  if (value instanceof ArrayBuffer) return new Uint8Array(value);
  if (Array.isArray(value)) return new Uint8Array(value as number[]);
  throw new Error(`Unsupported BLOB representation: ${typeof value}`);
}