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

import { openDB, type IDBPDatabase, type DBSchema } from 'idb';
import type {
  BroadcastChannelRecord,
  BroadcastMemberRecord,
  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 {
  openBroadcastChannelSensitive, openConfig, openIdentity, openOneTimePreKey,
  openRetired, openSession, openSignedPreKey, openStreamSensitive, openTrust,
  sealBroadcastChannelSensitive, sealConfig, sealIdentity, sealOneTimePreKey,
  sealRetired, sealSession, sealSignedPreKey, sealStreamSensitive, sealTrust,
} from '../crypto/row-codec.js';

/**
 * IndexedDB-backed StorageProvider with at-rest encryption.
 *
 * Schema is the IndexedDB equivalent of `EncryptedSQLiteStorage`: one object
 * store per `_enc` table, sealed payloads stored as `Uint8Array` in a
 * `ciphertext` field, routing/timestamp fields kept plaintext where SQLite
 * does so for query efficiency. Crypto stack — `KeyManager`, `aeadSeal`,
 * `aeadOpen`, row-codec sealers, AAD scheme — is shared verbatim with the
 * SQLite/Postgres backends, so a row sealed under one backend decrypts
 * under another given the same `KeyManager`.
 *
 * Browser-safe: imports `idb` (~12 kB, pure JS) and SubtleCrypto only.
 *
 * Usage:
 * ```ts
 * import { KeyManager } from '@shade/storage-encrypted/crypto';
 * import { EncryptedIndexedDBStorage } from '@shade/storage-encrypted/idb';
 *
 * const km = await KeyManager.open({
 *   kind: 'composite',
 *   sources: [
 *     { kind: 'passphrase', passphrase, salt: pwSalt },
 *     { kind: 'argon2id', secret: pin, salt: pinSalt },
 *   ],
 * });
 * const storage = await EncryptedIndexedDBStorage.open({
 *   dbName: 'my-app-shade',
 *   keyManager: km,
 * });
 * ```
 */
export class EncryptedIndexedDBStorage implements StorageProvider {
  private constructor(
    private db: IDBPDatabase<EncryptedShadeSchema>,
    private km: KeyManager,
  ) {}

  /**
   * Open (or create) the encrypted IndexedDB database. On first open the
   * storageKey fingerprint is persisted; subsequent opens with a different
   * KeyManager (wrong passphrase / PIN) reject with a clear error rather
   * than silently writing data under the wrong key.
   */
  static async open(opts: {
    dbName?: string;
    keyManager: KeyManager;
  }): Promise<EncryptedIndexedDBStorage> {
    const dbName = opts.dbName ?? 'shade-encrypted';
    const db = await openDB<EncryptedShadeSchema>(dbName, SCHEMA_VERSION, {
      upgrade(db, oldVersion) {
        if (oldVersion < 1) {
          db.createObjectStore('meta_enc', { keyPath: 'key' });
          db.createObjectStore('identity_enc', { keyPath: 'id' });
          db.createObjectStore('config_enc', { keyPath: 'key' });
          db.createObjectStore('signed_prekeys_enc', { keyPath: 'keyId' });
          db.createObjectStore('one_time_prekeys_enc', { keyPath: 'keyId' });
          db.createObjectStore('sessions_enc', { keyPath: 'address' });
          db.createObjectStore('trusted_identities_enc', { keyPath: 'address' });

          const retired = db.createObjectStore('retired_identities_enc', {
            keyPath: 'retiredAt',
          });
          retired.createIndex('byRetiredAt', 'retiredAt');

          const stream = db.createObjectStore('stream_state_enc', { keyPath: 'streamId' });
          stream.createIndex('byStatus', 'status');
          stream.createIndex('byPeerAddress', 'peerAddress');
          stream.createIndex('byUpdatedAt', 'updatedAt');

          db.createObjectStore('peer_verifications_enc', { keyPath: 'peerAddress' });
          db.createObjectStore('peer_identity_versions_enc', { keyPath: 'peerAddress' });
        }
        if (oldVersion < 2) {
          db.createObjectStore('broadcast_channels_enc', { keyPath: 'channelId' });
          const members = db.createObjectStore('broadcast_members_enc', {
            keyPath: ['channelId', 'peerAddress'],
          });
          members.createIndex('byChannelId', 'channelId');
        }
      },
    });
    const store = new EncryptedIndexedDBStorage(db, opts.keyManager);
    await store.assertKeyMatchesOrPersistFingerprint();
    return store;
  }

  /** Cleanly close the underlying connection. KeyManager is destroyed. */
  close(): void {
    this.db.close();
    this.km.destroy();
  }

  /**
   * 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 = await this.db.get('meta_enc', 'storage_key_fingerprint');
    if (!row) {
      await this.db.put('meta_enc', { key: 'storage_key_fingerprint', value: expected });
      return;
    }
    if (row.value !== expected) {
      throw new Error(
        'storage key mismatch — the supplied passphrase / keychain entry does not unlock this database',
      );
    }
  }

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

  async getIdentityKeyPair(): Promise<IdentityKeyPair | null> {
    const row = await this.db.get('identity_enc', 1);
    if (!row) return null;
    return openIdentity(this.km, row.ciphertext);
  }

  async saveIdentityKeyPair(kp: IdentityKeyPair): Promise<void> {
    const blob = await sealIdentity(this.km, kp);
    await this.db.put('identity_enc', { id: 1, ciphertext: blob });
  }

  async getLocalRegistrationId(): Promise<number> {
    const row = await this.db.get('config_enc', 'registrationId');
    if (!row) return 0;
    const v = await openConfig(this.km, 'registrationId', row.ciphertext);
    return parseInt(v, 10);
  }

  async saveLocalRegistrationId(id: number): Promise<void> {
    const blob = await sealConfig(this.km, 'registrationId', String(id));
    await this.db.put('config_enc', { key: 'registrationId', ciphertext: blob });
  }

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

  async getSignedPreKey(keyId: number): Promise<SignedPreKey | null> {
    const row = await this.db.get('signed_prekeys_enc', keyId);
    if (!row) return null;
    return openSignedPreKey(this.km, keyId, row.ciphertext);
  }

  async saveSignedPreKey(key: SignedPreKey): Promise<void> {
    const blob = await sealSignedPreKey(this.km, key);
    await this.db.put('signed_prekeys_enc', { keyId: key.keyId, ciphertext: blob });
  }

  async removeSignedPreKey(keyId: number): Promise<void> {
    await this.db.delete('signed_prekeys_enc', keyId);
  }

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

  async getOneTimePreKey(keyId: number): Promise<OneTimePreKey | null> {
    const row = await this.db.get('one_time_prekeys_enc', keyId);
    if (!row) return null;
    return openOneTimePreKey(this.km, keyId, row.ciphertext);
  }

  async saveOneTimePreKey(key: OneTimePreKey): Promise<void> {
    const blob = await sealOneTimePreKey(this.km, key);
    await this.db.put('one_time_prekeys_enc', { keyId: key.keyId, ciphertext: blob });
  }

  async removeOneTimePreKey(keyId: number): Promise<void> {
    await this.db.delete('one_time_prekeys_enc', keyId);
  }

  async getOneTimePreKeyCount(): Promise<number> {
    return this.db.count('one_time_prekeys_enc');
  }

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

  async getSession(address: string): Promise<SessionState | null> {
    const row = await this.db.get('sessions_enc', address);
    if (!row) return null;
    return openSession(this.km, address, row.ciphertext);
  }

  async saveSession(address: string, state: SessionState): Promise<void> {
    const blob = await sealSession(this.km, address, state);
    await this.db.put('sessions_enc', { address, ciphertext: blob });
  }

  async removeSession(address: string): Promise<void> {
    await this.db.delete('sessions_enc', address);
  }

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

  async isTrustedIdentity(address: string, identityKey: Uint8Array): Promise<boolean> {
    const row = await this.db.get('trusted_identities_enc', address);
    if (!row) return true; // TOFU
    const stored = await openTrust(this.km, address, row.ciphertext);
    return constantTimeEqual(stored, identityKey);
  }

  async saveTrustedIdentity(address: string, identityKey: Uint8Array): Promise<void> {
    const blob = await sealTrust(this.km, address, identityKey);
    await this.db.put('trusted_identities_enc', { address, ciphertext: blob });
  }

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

  async addRetiredIdentity(identity: RetiredIdentity): Promise<void> {
    const blob = await sealRetired(this.km, identity);
    await this.db.put('retired_identities_enc', {
      retiredAt: identity.retiredAt,
      ciphertext: blob,
    });
  }

  async getRetiredIdentities(): Promise<RetiredIdentity[]> {
    // Mirror SQLite's `ORDER BY retired_at DESC`
    const rows = await this.db.getAllFromIndex('retired_identities_enc', 'byRetiredAt');
    rows.reverse();
    return Promise.all(
      rows.map((r) => openRetired(this.km, r.retiredAt, r.ciphertext)),
    );
  }

  async pruneRetiredIdentities(olderThan: number): Promise<void> {
    const tx = this.db.transaction('retired_identities_enc', 'readwrite');
    const idx = tx.store.index('byRetiredAt');
    const range = IDBKeyRange.upperBound(olderThan, true);
    let cursor = await idx.openCursor(range);
    while (cursor) {
      await cursor.delete();
      cursor = await cursor.continue();
    }
    await tx.done;
  }

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

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

  async getStreamState(streamId: string): Promise<PersistedStreamState | null> {
    const row = await this.db.get('stream_state_enc', streamId);
    if (!row) return null;
    return this.rowToStreamState(row);
  }

  async removeStreamState(streamId: string): Promise<void> {
    await this.db.delete('stream_state_enc', streamId);
  }

  async listActiveStreamStates(direction?: 'send' | 'receive'): Promise<PersistedStreamState[]> {
    const tx = this.db.transaction('stream_state_enc');
    const idx = tx.store.index('byStatus');
    const active = await idx.getAll(IDBKeyRange.only('active'));
    const paused = await idx.getAll(IDBKeyRange.only('paused'));
    const merged = [...active, ...paused];
    const filtered = direction === undefined
      ? merged
      : merged.filter((r) => r.direction === direction);
    filtered.sort((a, b) => b.updatedAt - a.updatedAt);
    return Promise.all(filtered.map((r) => this.rowToStreamState(r)));
  }

  async pruneStreamStates(olderThan: number): Promise<void> {
    const tx = this.db.transaction('stream_state_enc', 'readwrite');
    const idx = tx.store.index('byUpdatedAt');
    const range = IDBKeyRange.upperBound(olderThan, true);
    let cursor = await idx.openCursor(range);
    while (cursor) {
      const row = cursor.value;
      if (row.status === 'finished' || row.status === 'aborted') {
        await cursor.delete();
      }
      cursor = await cursor.continue();
    }
    await tx.done;
  }

  // ─── Peer verifications (V3.3) ────────────────────────────
  // Fingerprints are public-by-design; stored in plaintext for symmetry
  // with the SQLite/Postgres encrypted backends.

  async savePeerVerification(v: PeerVerification): Promise<void> {
    await this.db.put('peer_verifications_enc', { ...v });
  }

  async getPeerVerification(address: string): Promise<PeerVerification | null> {
    const row = await this.db.get('peer_verifications_enc', address);
    if (!row) return null;
    return {
      peerAddress: row.peerAddress,
      fingerprint: row.fingerprint,
      verifiedAt: row.verifiedAt,
      verifiedBy: row.verifiedBy as PeerVerificationSource,
      identityVersion: row.identityVersion,
    };
  }

  async removePeerVerification(address: string): Promise<void> {
    await this.db.delete('peer_verifications_enc', address);
  }

  async getPeerIdentityVersion(address: string): Promise<number> {
    const row = await this.db.get('peer_identity_versions_enc', address);
    return row ? row.version : 1;
  }

  async bumpPeerIdentityVersion(address: string): Promise<number> {
    // Atomic read-modify-write under one IDB transaction. Closes the race
    // that exists in the SQLite version's non-atomic read-then-upsert.
    const tx = this.db.transaction('peer_identity_versions_enc', 'readwrite');
    const existing = await tx.store.get(address);
    const next = (existing ? existing.version : 1) + 1;
    await tx.store.put({ peerAddress: address, version: next });
    await tx.done;
    return next;
  }

  // ─── Broadcast channels (V4.6) ────────────────────────────

  async saveBroadcastChannel(channel: BroadcastChannelRecord): Promise<void> {
    const sealed = await sealBroadcastChannelSensitive(this.km, channel.channelId, {
      chainKey: channel.chainKey,
      iteration: channel.iteration,
      signingPublicKey: channel.signingPublicKey,
      ...(channel.signingPrivateKey !== undefined ? { signingPrivateKey: channel.signingPrivateKey } : {}),
    });
    await this.db.put('broadcast_channels_enc', {
      channelId: channel.channelId,
      ownerRole: channel.ownerRole,
      ownerAddress: channel.ownerAddress,
      label: channel.label ?? null,
      generation: channel.generation,
      ciphertext: sealed,
      createdAt: channel.createdAt,
      updatedAt: channel.updatedAt,
    });
  }

  async getBroadcastChannel(channelId: string): Promise<BroadcastChannelRecord | null> {
    const row = await this.db.get('broadcast_channels_enc', channelId);
    if (!row) return null;
    return this.encRowToChannel(row);
  }

  async listBroadcastChannels(): Promise<BroadcastChannelRecord[]> {
    const rows = await this.db.getAll('broadcast_channels_enc');
    rows.sort((a, b) => a.createdAt - b.createdAt);
    return Promise.all(rows.map((r) => this.encRowToChannel(r)));
  }

  async removeBroadcastChannel(channelId: string): Promise<void> {
    const tx = this.db.transaction(
      ['broadcast_channels_enc', 'broadcast_members_enc'],
      'readwrite',
    );
    const memIdx = tx.objectStore('broadcast_members_enc').index('byChannelId');
    let cursor = await memIdx.openCursor(IDBKeyRange.only(channelId));
    while (cursor) {
      await cursor.delete();
      cursor = await cursor.continue();
    }
    await tx.objectStore('broadcast_channels_enc').delete(channelId);
    await tx.done;
  }

  async saveBroadcastMember(member: BroadcastMemberRecord): Promise<void> {
    await this.db.put('broadcast_members_enc', { ...member });
  }

  async getBroadcastMembers(channelId: string): Promise<BroadcastMemberRecord[]> {
    const rows = await this.db.getAllFromIndex(
      'broadcast_members_enc',
      'byChannelId',
      IDBKeyRange.only(channelId),
    );
    rows.sort((a, b) => a.joinedAt - b.joinedAt);
    return rows.map((r) => ({
      channelId: r.channelId,
      peerAddress: r.peerAddress,
      joinedAt: r.joinedAt,
      removedAt: r.removedAt,
    }));
  }

  async removeBroadcastMember(channelId: string, peerAddress: string): Promise<void> {
    await this.db.delete('broadcast_members_enc', [channelId, peerAddress]);
  }

  private async encRowToChannel(row: BroadcastChannelEncRow): Promise<BroadcastChannelRecord> {
    const sensitive = await openBroadcastChannelSensitive(
      this.km,
      row.channelId,
      row.ciphertext,
    );
    const out: BroadcastChannelRecord = {
      channelId: row.channelId,
      ownerRole: row.ownerRole,
      ownerAddress: row.ownerAddress,
      generation: row.generation,
      chainKey: sensitive.chainKey,
      iteration: sensitive.iteration,
      signingPublicKey: sensitive.signingPublicKey,
      createdAt: row.createdAt,
      updatedAt: row.updatedAt,
    };
    if (row.label !== null) out.label = row.label;
    if (sensitive.signingPrivateKey !== undefined) out.signingPrivateKey = sensitive.signingPrivateKey;
    return out;
  }

  private async rowToStreamState(row: StreamStateEncRow): Promise<PersistedStreamState> {
    const sensitive = await openStreamSensitive(this.km, row.streamId, row.ciphertext);
    const out: PersistedStreamState = {
      streamId: row.streamId,
      direction: row.direction,
      peerAddress: row.peerAddress,
      status: row.status,
      metadataJson: sensitive.metadataJson,
      partitionJson: sensitive.partitionJson,
      laneStateJson: sensitive.laneStateJson,
      ioDescriptorJson: sensitive.ioDescriptorJson,
      secretEnc: sensitive.secretEnc,
      secretNonce: sensitive.secretNonce,
      createdAt: row.createdAt,
      updatedAt: row.updatedAt,
    };
    if (sensitive.overallHashState !== undefined) out.overallHashState = sensitive.overallHashState;
    return out;
  }
}

// ─── Schema ────────────────────────────────────────────────

const SCHEMA_VERSION = 2;

interface MetaRow { key: string; value: string }
interface IdentityRow { id: 1; ciphertext: Uint8Array }
interface ConfigRow { key: string; ciphertext: Uint8Array }
interface SignedPreKeyRow { keyId: number; ciphertext: Uint8Array }
interface OneTimePreKeyRow { keyId: number; ciphertext: Uint8Array }
interface SessionRow { address: string; ciphertext: Uint8Array }
interface TrustedIdentityRow { address: string; ciphertext: Uint8Array }
interface RetiredIdentityRow { retiredAt: number; ciphertext: Uint8Array }

interface StreamStateEncRow {
  streamId: string;
  direction: 'send' | 'receive';
  peerAddress: string;
  status: 'active' | 'paused' | 'finished' | 'aborted';
  ciphertext: Uint8Array;
  createdAt: number;
  updatedAt: number;
}

interface PeerVerificationRow {
  peerAddress: string;
  fingerprint: string;
  verifiedAt: number;
  verifiedBy: string;
  identityVersion: number;
}

interface PeerIdentityVersionRow { peerAddress: string; version: number }

interface BroadcastChannelEncRow {
  channelId: string;
  ownerRole: 'sender' | 'receiver';
  ownerAddress: string;
  label: string | null;
  generation: number;
  ciphertext: Uint8Array;
  createdAt: number;
  updatedAt: number;
}

interface BroadcastMemberEncRow {
  channelId: string;
  peerAddress: string;
  joinedAt: number;
  removedAt: number | null;
}

interface EncryptedShadeSchema extends DBSchema {
  meta_enc: { key: string; value: MetaRow };
  identity_enc: { key: number; value: IdentityRow };
  config_enc: { key: string; value: ConfigRow };
  signed_prekeys_enc: { key: number; value: SignedPreKeyRow };
  one_time_prekeys_enc: { key: number; value: OneTimePreKeyRow };
  sessions_enc: { key: string; value: SessionRow };
  trusted_identities_enc: { key: string; value: TrustedIdentityRow };
  retired_identities_enc: {
    key: number;
    value: RetiredIdentityRow;
    indexes: { byRetiredAt: number };
  };
  stream_state_enc: {
    key: string;
    value: StreamStateEncRow;
    indexes: {
      byStatus: string;
      byPeerAddress: string;
      byUpdatedAt: number;
    };
  };
  peer_verifications_enc: { key: string; value: PeerVerificationRow };
  peer_identity_versions_enc: { key: string; value: PeerIdentityVersionRow };
  broadcast_channels_enc: { key: string; value: BroadcastChannelEncRow };
  broadcast_members_enc: {
    key: [string, string];
    value: BroadcastMemberEncRow;
    indexes: { byChannelId: string };
  };
}