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

import { Database } from 'bun:sqlite';
import type { StorageProvider, IdentityKeyPair, SignedPreKey, OneTimePreKey, SessionState, RetiredIdentity, PersistedStreamState } from '@shade/core';
import {
  toBase64, fromBase64,
  constantTimeEqual,
  serializeSessionState, deserializeSessionState,
  serializeSignedPreKey, deserializeSignedPreKey,
  serializeOneTimePreKey, deserializeOneTimePreKey,
  serializeIdentityKeyPair, deserializeIdentityKeyPair,
} from '@shade/core';

/**
 * SQLite-backed StorageProvider for Shade client-side key/session storage.
 *
 * Uses bun:sqlite (built-in, zero deps). Stores private keys — for trusted environments only.
 * WAL mode enabled for crash safety. Auto-creates tables on first use.
 *
 * Docker usage:
 *   Volume mount /data, set SHADE_DB_PATH=/data/shade-client.db
 */
export class SQLiteStorage implements StorageProvider {
  private db: Database;

  // 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']>;
    listActiveStreamStatesByDirection: ReturnType<Database['prepare']>;
    pruneStreamStates: ReturnType<Database['prepare']>;
  };

  constructor(dbPath?: string) {
    const path = dbPath ?? process.env.SHADE_DB_PATH ?? '/data/shade-client.db';
    this.db = new Database(path, { create: true });
    this.db.exec('PRAGMA journal_mode=WAL');
    this.ensureTables();
    this.prepareStatements();
  }

  private ensureTables() {
    this.db.exec(`
      CREATE TABLE IF NOT EXISTS identity (
        id INTEGER PRIMARY KEY CHECK (id = 1),
        signing_public_key TEXT NOT NULL,
        signing_private_key TEXT NOT NULL,
        dh_public_key TEXT NOT NULL,
        dh_private_key TEXT NOT NULL
      );
      CREATE TABLE IF NOT EXISTS config (
        key TEXT PRIMARY KEY,
        value TEXT NOT NULL
      );
      CREATE TABLE IF NOT EXISTS signed_prekeys (
        key_id INTEGER PRIMARY KEY,
        data_json TEXT NOT NULL
      );
      CREATE TABLE IF NOT EXISTS one_time_prekeys (
        key_id INTEGER PRIMARY KEY,
        data_json TEXT NOT NULL
      );
      CREATE TABLE IF NOT EXISTS sessions (
        address TEXT PRIMARY KEY,
        state_json TEXT NOT NULL
      );
      CREATE TABLE IF NOT EXISTS trusted_identities (
        address TEXT PRIMARY KEY,
        identity_key TEXT NOT NULL
      );
      CREATE TABLE IF NOT EXISTS retired_identities (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        data_json TEXT NOT NULL,
        retired_at INTEGER NOT NULL
      );
      CREATE INDEX IF NOT EXISTS idx_retired_at ON retired_identities(retired_at);
      CREATE TABLE IF NOT EXISTS stream_state (
        stream_id          TEXT PRIMARY KEY,
        direction          TEXT NOT NULL,
        peer_address       TEXT NOT NULL,
        status             TEXT NOT NULL,
        metadata_json      TEXT NOT NULL,
        partition_json     TEXT NOT NULL,
        lane_state_json    TEXT NOT NULL,
        io_descriptor_json TEXT NOT NULL,
        secret_enc         BLOB NOT NULL,
        secret_nonce       BLOB NOT NULL,
        overall_hash_state TEXT,
        created_at         INTEGER NOT NULL,
        updated_at         INTEGER NOT NULL
      );
      CREATE INDEX IF NOT EXISTS idx_stream_state_peer    ON stream_state(peer_address);
      CREATE INDEX IF NOT EXISTS idx_stream_state_updated ON stream_state(updated_at);
      CREATE INDEX IF NOT EXISTS idx_stream_state_status  ON stream_state(status, direction);
    `);
  }

  private prepareStatements() {
    this.stmts = {
      getIdentity: this.db.prepare('SELECT * FROM identity WHERE id = 1'),
      saveIdentity: this.db.prepare('INSERT OR REPLACE INTO identity (id, signing_public_key, signing_private_key, dh_public_key, dh_private_key) VALUES (1, ?, ?, ?, ?)'),
      getConfig: this.db.prepare('SELECT value FROM config WHERE key = ?'),
      saveConfig: this.db.prepare('INSERT OR REPLACE INTO config (key, value) VALUES (?, ?)'),
      getSignedPreKey: this.db.prepare('SELECT data_json FROM signed_prekeys WHERE key_id = ?'),
      saveSignedPreKey: this.db.prepare('INSERT OR REPLACE INTO signed_prekeys (key_id, data_json) VALUES (?, ?)'),
      removeSignedPreKey: this.db.prepare('DELETE FROM signed_prekeys WHERE key_id = ?'),
      getOneTimePreKey: this.db.prepare('SELECT data_json FROM one_time_prekeys WHERE key_id = ?'),
      saveOneTimePreKey: this.db.prepare('INSERT OR REPLACE INTO one_time_prekeys (key_id, data_json) VALUES (?, ?)'),
      removeOneTimePreKey: this.db.prepare('DELETE FROM one_time_prekeys WHERE key_id = ?'),
      countOneTimePreKeys: this.db.prepare('SELECT COUNT(*) as count FROM one_time_prekeys'),
      getSession: this.db.prepare('SELECT state_json FROM sessions WHERE address = ?'),
      saveSession: this.db.prepare('INSERT OR REPLACE INTO sessions (address, state_json) VALUES (?, ?)'),
      removeSession: this.db.prepare('DELETE FROM sessions WHERE address = ?'),
      getTrust: this.db.prepare('SELECT identity_key FROM trusted_identities WHERE address = ?'),
      saveTrust: this.db.prepare('INSERT OR REPLACE INTO trusted_identities (address, identity_key) VALUES (?, ?)'),
      addRetired: this.db.prepare('INSERT INTO retired_identities (data_json, retired_at) VALUES (?, ?)'),
      listRetired: this.db.prepare('SELECT data_json, retired_at FROM retired_identities ORDER BY retired_at DESC'),
      pruneRetired: this.db.prepare('DELETE FROM retired_identities WHERE retired_at < ?'),
      saveStreamState: this.db.prepare(
        `INSERT OR REPLACE INTO 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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
      ),
      getStreamState: this.db.prepare('SELECT * FROM stream_state WHERE stream_id = ?'),
      removeStreamState: this.db.prepare('DELETE FROM stream_state WHERE stream_id = ?'),
      listActiveStreamStates: this.db.prepare(
        "SELECT * FROM stream_state WHERE status IN ('active', 'paused') ORDER BY updated_at DESC",
      ),
      listActiveStreamStatesByDirection: this.db.prepare(
        "SELECT * FROM stream_state WHERE status IN ('active', 'paused') AND direction = ? ORDER BY updated_at DESC",
      ),
      pruneStreamStates: this.db.prepare(
        "DELETE FROM stream_state WHERE status IN ('finished', 'aborted') AND updated_at < ?",
      ),
    };
  }

  close() {
    this.db.close();
  }

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

  async getIdentityKeyPair(): Promise<IdentityKeyPair | null> {
    const row = this.stmts.getIdentity.get() as any;
    if (!row) return null;
    return {
      signingPublicKey: fromBase64(row.signing_public_key),
      signingPrivateKey: fromBase64(row.signing_private_key),
      dhPublicKey: fromBase64(row.dh_public_key),
      dhPrivateKey: fromBase64(row.dh_private_key),
    };
  }

  async saveIdentityKeyPair(kp: IdentityKeyPair): Promise<void> {
    this.stmts.saveIdentity.run(
      toBase64(kp.signingPublicKey),
      toBase64(kp.signingPrivateKey),
      toBase64(kp.dhPublicKey),
      toBase64(kp.dhPrivateKey),
    );
  }

  async getLocalRegistrationId(): Promise<number> {
    const row = this.stmts.getConfig.get('registrationId') as any;
    return row ? parseInt(row.value, 10) : 0;
  }

  async saveLocalRegistrationId(id: number): Promise<void> {
    this.stmts.saveConfig.run('registrationId', String(id));
  }

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

  async getSignedPreKey(keyId: number): Promise<SignedPreKey | null> {
    const row = this.stmts.getSignedPreKey.get(keyId) as any;
    if (!row) return null;
    return deserializeSignedPreKey(row.data_json);
  }

  async saveSignedPreKey(key: SignedPreKey): Promise<void> {
    this.stmts.saveSignedPreKey.run(key.keyId, serializeSignedPreKey(key));
  }

  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 any;
    if (!row) return null;
    return deserializeOneTimePreKey(row.data_json);
  }

  async saveOneTimePreKey(key: OneTimePreKey): Promise<void> {
    this.stmts.saveOneTimePreKey.run(key.keyId, serializeOneTimePreKey(key));
  }

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

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

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

  async getSession(address: string): Promise<SessionState | null> {
    const row = this.stmts.getSession.get(address) as any;
    if (!row) return null;
    return deserializeSessionState(row.state_json);
  }

  async saveSession(address: string, state: SessionState): Promise<void> {
    this.stmts.saveSession.run(address, serializeSessionState(state));
  }

  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 any;
    if (!row) return true; // TOFU
    const storedBytes = fromBase64(row.identity_key);
    return constantTimeEqual(storedBytes, identityKey);
  }

  async saveTrustedIdentity(address: string, identityKey: Uint8Array): Promise<void> {
    this.stmts.saveTrust.run(address, toBase64(identityKey));
  }

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

  async addRetiredIdentity(identity: RetiredIdentity): Promise<void> {
    this.stmts.addRetired.run(
      serializeIdentityKeyPair(identity.keyPair),
      identity.retiredAt,
    );
  }

  async getRetiredIdentities(): Promise<RetiredIdentity[]> {
    const rows = this.stmts.listRetired.all() as any[];
    return rows.map((r) => ({
      keyPair: deserializeIdentityKeyPair(r.data_json),
      retiredAt: r.retired_at,
    }));
  }

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

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

  async saveStreamState(state: PersistedStreamState): Promise<void> {
    this.stmts.saveStreamState.run(
      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,
    );
  }

  async getStreamState(streamId: string): Promise<PersistedStreamState | null> {
    const row = this.stmts.getStreamState.get(streamId) as any;
    if (!row) return null;
    return 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 any[])
        : (this.stmts.listActiveStreamStatesByDirection.all(direction) as any[])
    );
    return rows.map(rowToStreamState);
  }

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

function rowToStreamState(row: any): PersistedStreamState {
  const out: PersistedStreamState = {
    streamId: row.stream_id,
    direction: row.direction,
    peerAddress: row.peer_address,
    status: row.status,
    metadataJson: row.metadata_json,
    partitionJson: row.partition_json,
    laneStateJson: row.lane_state_json,
    ioDescriptorJson: row.io_descriptor_json,
    secretEnc: toBytes(row.secret_enc),
    secretNonce: toBytes(row.secret_nonce),
    createdAt: Number(row.created_at),
    updatedAt: Number(row.updated_at),
  };
  if (row.overall_hash_state !== null && row.overall_hash_state !== undefined) {
    out.overallHashState = row.overall_hash_state;
  }
  return out;
}

function toBytes(value: unknown): Uint8Array {
  if (value instanceof Uint8Array) return value;
  if (value instanceof ArrayBuffer) return new Uint8Array(value);
  if (Array.isArray(value)) return new Uint8Array(value);
  if (typeof value === 'string') {
    return new Uint8Array(value.split('').map((c) => c.charCodeAt(0)));
  }
  throw new Error(`Unsupported BLOB representation: ${typeof value}`);
}