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

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

/**
 * 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']>;
  };

  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);
    `);
  }

  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 < ?'),
    };
  }

  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);
  }
}