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

import { openDB, type IDBPDatabase, type DBSchema } from 'idb';
import type {
  StorageProvider,
  IdentityKeyPair,
  SignedPreKey,
  OneTimePreKey,
  SessionState,
  RetiredIdentity,
  PersistedStreamState,
  PeerVerification,
  PeerVerificationSource,
} from '@shade/core';
import {
  toBase64, fromBase64,
  constantTimeEqual,
  serializeSessionState, deserializeSessionState,
  serializeSignedPreKey, deserializeSignedPreKey,
  serializeOneTimePreKey, deserializeOneTimePreKey,
  serializeIdentityKeyPair, deserializeIdentityKeyPair,
} from '@shade/core';

/**
 * IndexedDB-backed StorageProvider for browser-side Shade clients.
 *
 * Persists identity, prekeys, sessions, retired identities, peer
 * verifications and stream-resume state across tab refresh and browser
 * restart. Same data shapes as `@shade/storage-sqlite` so cross-adapter
 * import/export remains feasible.
 *
 * Usage:
 * ```ts
 * const storage = await IndexedDBStorage.create({ dbName: 'my-app-shade' });
 * const manager = new ShadeSessionManager(crypto, storage);
 * ```
 */
export class IndexedDBStorage implements StorageProvider {
  private constructor(private db: IDBPDatabase<ShadeSchema>) {}

  /**
   * Open (or create) the IndexedDB database. Idempotent — repeated calls
   * with the same dbName resolve to a fresh connection sharing the same
   * underlying object stores.
   */
  static async create(opts: { dbName?: string } = {}): Promise<IndexedDBStorage> {
    const dbName = opts.dbName ?? 'shade';
    const db = await openDB<ShadeSchema>(dbName, SCHEMA_VERSION, {
      upgrade(db, oldVersion) {
        if (oldVersion < 1) {
          db.createObjectStore('identity', { keyPath: 'id' });
          db.createObjectStore('config', { keyPath: 'key' });
          db.createObjectStore('signedPreKeys', { keyPath: 'keyId' });
          db.createObjectStore('oneTimePreKeys', { keyPath: 'keyId' });
          db.createObjectStore('sessions', { keyPath: 'address' });
          db.createObjectStore('trustedIdentities', { keyPath: 'address' });

          const retired = db.createObjectStore('retiredIdentities', {
            keyPath: 'id',
            autoIncrement: true,
          });
          retired.createIndex('byRetiredAt', 'retiredAt');

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

          db.createObjectStore('peerVerifications', { keyPath: 'peerAddress' });
          db.createObjectStore('peerIdentityVersions', { keyPath: 'peerAddress' });
        }
      },
    });
    return new IndexedDBStorage(db);
  }

  /** Cleanly close the underlying connection. */
  async close(): Promise<void> {
    this.db.close();
  }

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

  async getIdentityKeyPair(): Promise<IdentityKeyPair | null> {
    const row = await this.db.get('identity', 1);
    if (!row) return null;
    return {
      signingPublicKey: fromBase64(row.signingPublicKey),
      signingPrivateKey: fromBase64(row.signingPrivateKey),
      dhPublicKey: fromBase64(row.dhPublicKey),
      dhPrivateKey: fromBase64(row.dhPrivateKey),
    };
  }

  async saveIdentityKeyPair(kp: IdentityKeyPair): Promise<void> {
    await this.db.put('identity', {
      id: 1,
      signingPublicKey: toBase64(kp.signingPublicKey),
      signingPrivateKey: toBase64(kp.signingPrivateKey),
      dhPublicKey: toBase64(kp.dhPublicKey),
      dhPrivateKey: toBase64(kp.dhPrivateKey),
    });
  }

  async getLocalRegistrationId(): Promise<number> {
    const row = await this.db.get('config', 'registrationId');
    return row ? parseInt(row.value, 10) : 0;
  }

  async saveLocalRegistrationId(id: number): Promise<void> {
    await this.db.put('config', { key: 'registrationId', value: String(id) });
  }

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

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

  async saveSignedPreKey(key: SignedPreKey): Promise<void> {
    await this.db.put('signedPreKeys', {
      keyId: key.keyId,
      dataJson: serializeSignedPreKey(key),
    });
  }

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

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

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

  async saveOneTimePreKey(key: OneTimePreKey): Promise<void> {
    await this.db.put('oneTimePreKeys', {
      keyId: key.keyId,
      dataJson: serializeOneTimePreKey(key),
    });
  }

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

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

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

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

  async saveSession(address: string, state: SessionState): Promise<void> {
    await this.db.put('sessions', { address, stateJson: serializeSessionState(state) });
  }

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

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

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

  async saveTrustedIdentity(address: string, identityKey: Uint8Array): Promise<void> {
    await this.db.put('trustedIdentities', { address, identityKey: toBase64(identityKey) });
  }

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

  async addRetiredIdentity(identity: RetiredIdentity): Promise<void> {
    // autoIncrement: omit `id` so IDB assigns one
    await this.db.add('retiredIdentities', {
      dataJson: serializeIdentityKeyPair(identity.keyPair),
      retiredAt: identity.retiredAt,
    } as RetiredIdentityRow);
  }

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

  async pruneRetiredIdentities(olderThan: number): Promise<void> {
    const tx = this.db.transaction('retiredIdentities', '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> {
    await this.db.put('streamStates', persistedToRow(state));
  }

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

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

  async listActiveStreamStates(direction?: 'send' | 'receive'): Promise<PersistedStreamState[]> {
    const idx = this.db.transaction('streamStates').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 filtered.map(rowToPersisted);
  }

  async pruneStreamStates(olderThan: number): Promise<void> {
    const tx = this.db.transaction('streamStates', '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) ────────────────────────────

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

  async getPeerVerification(address: string): Promise<PeerVerification | null> {
    const row = await this.db.get('peerVerifications', 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('peerVerifications', address);
  }

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

  /**
   * Atomic read-modify-write under one IDB transaction. SQLite's version
   * is a non-atomic read-then-upsert; the IDB version closes that race
   * because IDB transactions auto-commit only when control returns to
   * the event loop without pending requests.
   */
  async bumpPeerIdentityVersion(address: string): Promise<number> {
    const tx = this.db.transaction('peerIdentityVersions', '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;
  }
}

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

const SCHEMA_VERSION = 1;

interface IdentityRow {
  id: 1;
  signingPublicKey: string;
  signingPrivateKey: string;
  dhPublicKey: string;
  dhPrivateKey: string;
}

interface ConfigRow {
  key: string;
  value: string;
}

interface SignedPreKeyRow {
  keyId: number;
  dataJson: string;
}

interface OneTimePreKeyRow {
  keyId: number;
  dataJson: string;
}

interface SessionRow {
  address: string;
  stateJson: string;
}

interface TrustedIdentityRow {
  address: string;
  identityKey: string;
}

interface RetiredIdentityRow {
  id?: number;
  dataJson: string;
  retiredAt: number;
}

interface StreamStateRow {
  streamId: string;
  direction: 'send' | 'receive';
  peerAddress: string;
  status: 'active' | 'paused' | 'finished' | 'aborted';
  metadataJson: string;
  partitionJson: string;
  laneStateJson: string;
  ioDescriptorJson: string;
  secretEnc: Uint8Array;
  secretNonce: Uint8Array;
  overallHashState: string | null;
  createdAt: number;
  updatedAt: number;
}

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

interface PeerIdentityVersionRow {
  peerAddress: string;
  version: number;
}

interface ShadeSchema extends DBSchema {
  identity: { key: number; value: IdentityRow };
  config: { key: string; value: ConfigRow };
  signedPreKeys: { key: number; value: SignedPreKeyRow };
  oneTimePreKeys: { key: number; value: OneTimePreKeyRow };
  sessions: { key: string; value: SessionRow };
  trustedIdentities: { key: string; value: TrustedIdentityRow };
  retiredIdentities: {
    key: number;
    value: RetiredIdentityRow;
    indexes: { byRetiredAt: number };
  };
  streamStates: {
    key: string;
    value: StreamStateRow;
    indexes: {
      byStatus: string;
      byPeerAddress: string;
      byUpdatedAt: number;
    };
  };
  peerVerifications: { key: string; value: PeerVerificationRow };
  peerIdentityVersions: { key: string; value: PeerIdentityVersionRow };
}

// ─── Helpers ──────────────────────────────────────────────

function persistedToRow(s: PersistedStreamState): StreamStateRow {
  return {
    streamId: s.streamId,
    direction: s.direction,
    peerAddress: s.peerAddress,
    status: s.status,
    metadataJson: s.metadataJson,
    partitionJson: s.partitionJson,
    laneStateJson: s.laneStateJson,
    ioDescriptorJson: s.ioDescriptorJson,
    secretEnc: s.secretEnc,
    secretNonce: s.secretNonce,
    overallHashState: s.overallHashState ?? null,
    createdAt: s.createdAt,
    updatedAt: s.updatedAt,
  };
}

function rowToPersisted(r: StreamStateRow): PersistedStreamState {
  const out: PersistedStreamState = {
    streamId: r.streamId,
    direction: r.direction,
    peerAddress: r.peerAddress,
    status: r.status,
    metadataJson: r.metadataJson,
    partitionJson: r.partitionJson,
    laneStateJson: r.laneStateJson,
    ioDescriptorJson: r.ioDescriptorJson,
    secretEnc: r.secretEnc,
    secretNonce: r.secretNonce,
    createdAt: r.createdAt,
    updatedAt: r.updatedAt,
  };
  if (r.overallHashState !== null) out.overallHashState = r.overallHashState;
  return out;
}