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

import { Database } from 'bun:sqlite';
import type { PrekeyStore } from '@shade/server/src/store.js';
import { toBase64, fromBase64 } from '@shade/core';

/**
 * SQLite-backed PrekeyStore for the Shade Prekey Server.
 *
 * Stores PUBLIC keys only (never private keys). Used by the prekey server
 * Docker container to persist registered identities and prekey bundles.
 *
 * Docker usage:
 *   Volume mount /data, set SHADE_PREKEY_DB_PATH=/data/shade-prekeys.db
 */
export class SqlitePrekeyStore implements PrekeyStore {
  private db: Database;

  private stmts!: {
    saveIdentity: ReturnType<Database['prepare']>;
    getIdentity: ReturnType<Database['prepare']>;
    saveSignedPreKey: ReturnType<Database['prepare']>;
    getSignedPreKey: ReturnType<Database['prepare']>;
    insertOTPK: ReturnType<Database['prepare']>;
    consumeOTPK: ReturnType<Database['prepare']>;
    countOTPK: ReturnType<Database['prepare']>;
    deleteIdentity: ReturnType<Database['prepare']>;
    deleteSignedPreKey: ReturnType<Database['prepare']>;
    deleteOTPKs: ReturnType<Database['prepare']>;
  };

  constructor(dbPath?: string) {
    const path = dbPath ?? process.env.SHADE_PREKEY_DB_PATH ?? '/data/shade-prekeys.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 identities (
        address TEXT PRIMARY KEY,
        identity_signing_key TEXT NOT NULL,
        identity_dh_key TEXT NOT NULL
      );
      CREATE TABLE IF NOT EXISTS signed_prekeys (
        address TEXT PRIMARY KEY,
        key_id INTEGER NOT NULL,
        public_key TEXT NOT NULL,
        signature TEXT NOT NULL
      );
      CREATE TABLE IF NOT EXISTS one_time_prekeys (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        address TEXT NOT NULL,
        key_id INTEGER NOT NULL,
        public_key TEXT NOT NULL
      );
      CREATE INDEX IF NOT EXISTS idx_otp_address ON one_time_prekeys(address);
    `);
  }

  private prepareStatements() {
    this.stmts = {
      saveIdentity: this.db.prepare('INSERT OR REPLACE INTO identities (address, identity_signing_key, identity_dh_key) VALUES (?, ?, ?)'),
      getIdentity: this.db.prepare('SELECT identity_signing_key, identity_dh_key FROM identities WHERE address = ?'),
      saveSignedPreKey: this.db.prepare('INSERT OR REPLACE INTO signed_prekeys (address, key_id, public_key, signature) VALUES (?, ?, ?, ?)'),
      getSignedPreKey: this.db.prepare('SELECT key_id, public_key, signature FROM signed_prekeys WHERE address = ?'),
      insertOTPK: this.db.prepare('INSERT INTO one_time_prekeys (address, key_id, public_key) VALUES (?, ?, ?)'),
      consumeOTPK: this.db.prepare('DELETE FROM one_time_prekeys WHERE id = (SELECT id FROM one_time_prekeys WHERE address = ? ORDER BY id LIMIT 1) RETURNING key_id, public_key'),
      countOTPK: this.db.prepare('SELECT COUNT(*) as count FROM one_time_prekeys WHERE address = ?'),
      deleteIdentity: this.db.prepare('DELETE FROM identities WHERE address = ?'),
      deleteSignedPreKey: this.db.prepare('DELETE FROM signed_prekeys WHERE address = ?'),
      deleteOTPKs: this.db.prepare('DELETE FROM one_time_prekeys WHERE address = ?'),
    };
  }

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

  async saveIdentity(address: string, identitySigningKey: Uint8Array, identityDHKey: Uint8Array): Promise<void> {
    this.stmts.saveIdentity.run(address, toBase64(identitySigningKey), toBase64(identityDHKey));
  }

  async getIdentity(address: string): Promise<{ identitySigningKey: Uint8Array; identityDHKey: Uint8Array } | null> {
    const row = this.stmts.getIdentity.get(address) as any;
    if (!row) return null;
    return {
      identitySigningKey: fromBase64(row.identity_signing_key),
      identityDHKey: fromBase64(row.identity_dh_key),
    };
  }

  async saveSignedPreKey(address: string, keyId: number, publicKey: Uint8Array, signature: Uint8Array): Promise<void> {
    this.stmts.saveSignedPreKey.run(address, keyId, toBase64(publicKey), toBase64(signature));
  }

  async getSignedPreKey(address: string): Promise<{ keyId: number; publicKey: Uint8Array; signature: Uint8Array } | null> {
    const row = this.stmts.getSignedPreKey.get(address) as any;
    if (!row) return null;
    return {
      keyId: row.key_id,
      publicKey: fromBase64(row.public_key),
      signature: fromBase64(row.signature),
    };
  }

  async saveOneTimePreKeys(address: string, keys: Array<{ keyId: number; publicKey: Uint8Array }>): Promise<void> {
    const insertMany = this.db.transaction(() => {
      for (const k of keys) {
        this.stmts.insertOTPK.run(address, k.keyId, toBase64(k.publicKey));
      }
    });
    insertMany();
  }

  async consumeOneTimePreKey(address: string): Promise<{ keyId: number; publicKey: Uint8Array } | null> {
    const row = this.stmts.consumeOTPK.get(address) as any;
    if (!row) return null;
    return {
      keyId: row.key_id,
      publicKey: fromBase64(row.public_key),
    };
  }

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

  async deleteAll(address: string): Promise<void> {
    const deleteAllTx = this.db.transaction(() => {
      this.stmts.deleteIdentity.run(address);
      this.stmts.deleteSignedPreKey.run(address);
      this.stmts.deleteOTPKs.run(address);
    });
    deleteAllTx();
  }
}