Shade/packages/shade-storage-indexeddb/tests/indexeddb-storage.test.ts

import 'fake-indexeddb/auto';
import { describe, test, expect, beforeEach, afterEach } from 'bun:test';
import { IndexedDBStorage } from '../src/indexeddb-storage.js';
import { SubtleCryptoProvider, MemoryStorage as _MemoryStorage } from '@shade/crypto-web';
import { ShadeSessionManager } from '@shade/core';
import type { IdentityKeyPair, SignedPreKey, OneTimePreKey } from '@shade/core';

void _MemoryStorage; // keep import side-effect-free reference shape

const crypto = new SubtleCryptoProvider();

function randBytes(n: number): Uint8Array {
  const buf = new Uint8Array(n);
  globalThis.crypto.getRandomValues(buf);
  return buf;
}

function uniqueDbName(): string {
  return `shade-test-${Date.now()}-${Math.random().toString(36).slice(2)}`;
}

async function deleteDb(dbName: string): Promise<void> {
  await new Promise<void>((resolve, reject) => {
    const req = indexedDB.deleteDatabase(dbName);
    req.onsuccess = () => resolve();
    req.onerror = () => reject(req.error);
    req.onblocked = () => resolve();
  });
}

describe('IndexedDBStorage', () => {
  let dbName: string;
  let storage: IndexedDBStorage;

  beforeEach(async () => {
    dbName = uniqueDbName();
    storage = await IndexedDBStorage.create({ dbName });
  });

  afterEach(async () => {
    await storage.close();
    await deleteDb(dbName);
  });

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

  describe('identity', () => {
    test('returns null when no identity stored', async () => {
      expect(await storage.getIdentityKeyPair()).toBeNull();
    });

    test('save and retrieve identity keypair', async () => {
      const ikp: IdentityKeyPair = {
        signingPublicKey: randBytes(32),
        signingPrivateKey: randBytes(32),
        dhPublicKey: randBytes(32),
        dhPrivateKey: randBytes(32),
      };
      await storage.saveIdentityKeyPair(ikp);
      const restored = await storage.getIdentityKeyPair();
      expect(restored).not.toBeNull();
      expect(restored!.signingPublicKey).toEqual(ikp.signingPublicKey);
      expect(restored!.dhPrivateKey).toEqual(ikp.dhPrivateKey);
    });

    test('registration ID roundtrip', async () => {
      expect(await storage.getLocalRegistrationId()).toBe(0);
      await storage.saveLocalRegistrationId(42);
      expect(await storage.getLocalRegistrationId()).toBe(42);
    });
  });

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

  describe('signed prekeys', () => {
    test('save, get, remove', async () => {
      const spk: SignedPreKey = {
        keyId: 1,
        keyPair: { publicKey: randBytes(32), privateKey: randBytes(32) },
        signature: randBytes(64),
        timestamp: Date.now(),
      };
      await storage.saveSignedPreKey(spk);
      const restored = await storage.getSignedPreKey(1);
      expect(restored).not.toBeNull();
      expect(restored!.keyId).toBe(1);
      expect(restored!.keyPair.publicKey).toEqual(spk.keyPair.publicKey);

      await storage.removeSignedPreKey(1);
      expect(await storage.getSignedPreKey(1)).toBeNull();
    });
  });

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

  describe('one-time prekeys', () => {
    test('save, get, remove, count', async () => {
      const otpk: OneTimePreKey = {
        keyId: 100,
        keyPair: { publicKey: randBytes(32), privateKey: randBytes(32) },
      };
      await storage.saveOneTimePreKey(otpk);
      expect(await storage.getOneTimePreKeyCount()).toBe(1);

      const restored = await storage.getOneTimePreKey(100);
      expect(restored!.keyId).toBe(100);

      await storage.removeOneTimePreKey(100);
      expect(await storage.getOneTimePreKeyCount()).toBe(0);
      expect(await storage.getOneTimePreKey(100)).toBeNull();
    });
  });

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

  describe('sessions', () => {
    test('save and restore session with skipped keys', async () => {
      const state = {
        remoteIdentityKey: randBytes(32),
        rootKey: randBytes(32),
        sendChain: { chainKey: randBytes(32), counter: 5 },
        receiveChain: { chainKey: randBytes(32), counter: 3 },
        dhSend: { publicKey: randBytes(32), privateKey: randBytes(32) },
        dhReceive: randBytes(32),
        previousSendCounter: 4,
        skippedKeys: new Map([['key:1', randBytes(32)]]),
      };

      await storage.saveSession('bob', state);
      const restored = await storage.getSession('bob');
      expect(restored).not.toBeNull();
      expect(restored!.sendChain.counter).toBe(5);
      expect(restored!.skippedKeys.size).toBe(1);
      expect(restored!.skippedKeys.get('key:1')).toEqual(state.skippedKeys.get('key:1')!);
    });

    test('remove session', async () => {
      await storage.saveSession('bob', {
        remoteIdentityKey: randBytes(32),
        rootKey: randBytes(32),
        sendChain: { chainKey: randBytes(32), counter: 0 },
        receiveChain: null,
        dhSend: { publicKey: randBytes(32), privateKey: randBytes(32) },
        dhReceive: null,
        previousSendCounter: 0,
        skippedKeys: new Map(),
      });
      await storage.removeSession('bob');
      expect(await storage.getSession('bob')).toBeNull();
    });
  });

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

  describe('trust', () => {
    test('TOFU: first use is trusted', async () => {
      expect(await storage.isTrustedIdentity('bob', randBytes(32))).toBe(true);
    });

    test('saved identity matches', async () => {
      const key = randBytes(32);
      await storage.saveTrustedIdentity('bob', key);
      expect(await storage.isTrustedIdentity('bob', key)).toBe(true);
      expect(await storage.isTrustedIdentity('bob', randBytes(32))).toBe(false);
    });
  });

  // ─── Retired identities ───────────────────────────────────

  describe('retired identities', () => {
    test('add, list (newest first), prune', async () => {
      const mk = (retiredAt: number): { keyPair: IdentityKeyPair; retiredAt: number } => ({
        keyPair: {
          signingPublicKey: randBytes(32),
          signingPrivateKey: randBytes(32),
          dhPublicKey: randBytes(32),
          dhPrivateKey: randBytes(32),
        },
        retiredAt,
      });

      await storage.addRetiredIdentity(mk(100));
      await storage.addRetiredIdentity(mk(300));
      await storage.addRetiredIdentity(mk(200));

      const list = await storage.getRetiredIdentities();
      expect(list.map((r) => r.retiredAt)).toEqual([300, 200, 100]);

      await storage.pruneRetiredIdentities(250);
      const after = await storage.getRetiredIdentities();
      expect(after.map((r) => r.retiredAt)).toEqual([300]);
    });
  });

  // ─── Crash Recovery ───────────────────────────────────────

  describe('persistence across close/reopen', () => {
    test('data survives close and reopen', async () => {
      const ikp: IdentityKeyPair = {
        signingPublicKey: randBytes(32),
        signingPrivateKey: randBytes(32),
        dhPublicKey: randBytes(32),
        dhPrivateKey: randBytes(32),
      };
      await storage.saveIdentityKeyPair(ikp);
      await storage.saveLocalRegistrationId(99);
      await storage.saveSession('alice', {
        remoteIdentityKey: randBytes(32),
        rootKey: randBytes(32),
        sendChain: { chainKey: randBytes(32), counter: 7 },
        receiveChain: null,
        dhSend: { publicKey: randBytes(32), privateKey: randBytes(32) },
        dhReceive: null,
        previousSendCounter: 0,
        skippedKeys: new Map(),
      });

      // Close and reopen against the same dbName
      await storage.close();
      storage = await IndexedDBStorage.create({ dbName });

      const restored = await storage.getIdentityKeyPair();
      expect(restored!.signingPublicKey).toEqual(ikp.signingPublicKey);
      expect(await storage.getLocalRegistrationId()).toBe(99);
      const session = await storage.getSession('alice');
      expect(session!.sendChain.counter).toBe(7);
    });
  });

  // ─── Full E2EE with IndexedDBStorage ──────────────────────

  describe('full E2EE conversation with persistent storage', () => {
    test('encrypt, close, reopen, continue conversation', async () => {
      const bobDbName = uniqueDbName();
      let bobStorage = await IndexedDBStorage.create({ dbName: bobDbName });

      try {
        const alice = new ShadeSessionManager(crypto, storage);
        let bob = new ShadeSessionManager(crypto, bobStorage);
        await alice.initialize();
        await bob.initialize();

        const otpks = await bob.generateOneTimePreKeys(5);
        const bundle = await bob.createPreKeyBundle();
        bundle.oneTimePreKey = { keyId: otpks[0]!.keyId, publicKey: otpks[0]!.keyPair.publicKey };

        await alice.initSessionFromBundle('bob', bundle);

        const env1 = await alice.encrypt('bob', 'Hello persistent!');
        expect(await bob.decrypt('alice', env1)).toBe('Hello persistent!');

        const env2 = await bob.encrypt('alice', 'Got it!');
        expect(await alice.decrypt('bob', env2)).toBe('Got it!');

        // "Crash" Bob — close the IDB and reopen
        await bobStorage.close();
        bobStorage = await IndexedDBStorage.create({ dbName: bobDbName });
        bob = new ShadeSessionManager(crypto, bobStorage);
        await bob.initialize();

        const env3 = await alice.encrypt('bob', 'After your restart');
        expect(await bob.decrypt('alice', env3)).toBe('After your restart');

        const env4 = await bob.encrypt('alice', 'I survived!');
        expect(await alice.decrypt('bob', env4)).toBe('I survived!');
      } finally {
        await bobStorage.close();
        await deleteDb(bobDbName);
      }
    });
  });
});