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

import { describe, test, expect, beforeEach, afterEach } from 'bun:test';
import { tmpdir } from 'os';
import { join } from 'path';
import { unlinkSync } from 'fs';
import { SQLiteStorage } from '../src/sqlite-storage.js';
import { SubtleCryptoProvider, MemoryStorage } from '@shade/crypto-web';
import { ShadeSessionManager } from '@shade/core';
import type { IdentityKeyPair, SignedPreKey, OneTimePreKey } from '@shade/core';

const crypto = new SubtleCryptoProvider();

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

function tempDbPath(): string {
  return join(tmpdir(), `shade-test-${Date.now()}-${Math.random().toString(36).slice(2)}.db`);
}

describe('SQLiteStorage', () => {
  let dbPath: string;
  let storage: SQLiteStorage;

  beforeEach(() => {
    dbPath = tempDbPath();
    storage = new SQLiteStorage(dbPath);
  });

  afterEach(() => {
    storage.close();
    try { unlinkSync(dbPath); } catch {}
    try { unlinkSync(dbPath + '-wal'); } catch {}
    try { unlinkSync(dbPath + '-shm'); } catch {}
  });

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

  // ─── 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
      storage.close();
      storage = new SQLiteStorage(dbPath);

      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 SQLiteStorage ─────────────────────────

  describe('full E2EE conversation with persistent storage', () => {
    test('encrypt, close, reopen, continue conversation', async () => {
      const bobDbPath = tempDbPath();
      let bobStorage = new SQLiteStorage(bobDbPath);

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

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

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

        // "Crash" Bob — close DB and reopen
        bobStorage.close();
        bobStorage = new SQLiteStorage(bobDbPath);
        bob = new ShadeSessionManager(crypto, bobStorage);
        await bob.initialize();

        // Continue conversation after "crash"
        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 {
        bobStorage.close();
        try { unlinkSync(bobDbPath); } catch {}
        try { unlinkSync(bobDbPath + '-wal'); } catch {}
        try { unlinkSync(bobDbPath + '-shm'); } catch {}
      }
    });
  });
});