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

import { describe, test, expect, beforeEach, afterEach } from 'bun:test';
import { tmpdir } from 'os';
import { join } from 'path';
import { unlinkSync } from 'fs';
import { SqlitePrekeyStore } from '../src/sqlite-prekey-store.js';

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

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

describe('SqlitePrekeyStore', () => {
  let dbPath: string;
  let store: SqlitePrekeyStore;

  beforeEach(() => {
    dbPath = tempDbPath();
    store = new SqlitePrekeyStore(dbPath);
  });

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

  describe('identity', () => {
    test('save and get identity', async () => {
      const sigKey = randBytes(32);
      const dhKey = randBytes(32);
      await store.saveIdentity('bob', sigKey, dhKey);
      const id = await store.getIdentity('bob');
      expect(id).not.toBeNull();
      expect(id!.identitySigningKey).toEqual(sigKey);
      expect(id!.identityDHKey).toEqual(dhKey);
    });

    test('returns null for unknown address', async () => {
      expect(await store.getIdentity('nobody')).toBeNull();
    });
  });

  describe('signed prekeys', () => {
    test('save and get signed prekey', async () => {
      const pubKey = randBytes(32);
      const sig = randBytes(64);
      await store.saveSignedPreKey('bob', 1, pubKey, sig);
      const spk = await store.getSignedPreKey('bob');
      expect(spk!.keyId).toBe(1);
      expect(spk!.publicKey).toEqual(pubKey);
      expect(spk!.signature).toEqual(sig);
    });
  });

  describe('one-time prekeys', () => {
    test('FIFO consumption', async () => {
      await store.saveOneTimePreKeys('bob', [
        { keyId: 100, publicKey: randBytes(32) },
        { keyId: 101, publicKey: randBytes(32) },
        { keyId: 102, publicKey: randBytes(32) },
      ]);

      expect(await store.getOneTimePreKeyCount('bob')).toBe(3);

      const k1 = await store.consumeOneTimePreKey('bob');
      expect(k1!.keyId).toBe(100);
      expect(await store.getOneTimePreKeyCount('bob')).toBe(2);

      const k2 = await store.consumeOneTimePreKey('bob');
      expect(k2!.keyId).toBe(101);

      const k3 = await store.consumeOneTimePreKey('bob');
      expect(k3!.keyId).toBe(102);

      const k4 = await store.consumeOneTimePreKey('bob');
      expect(k4).toBeNull();
      expect(await store.getOneTimePreKeyCount('bob')).toBe(0);
    });

    test('replenish adds more keys', async () => {
      await store.saveOneTimePreKeys('bob', [{ keyId: 100, publicKey: randBytes(32) }]);
      await store.saveOneTimePreKeys('bob', [{ keyId: 200, publicKey: randBytes(32) }, { keyId: 201, publicKey: randBytes(32) }]);
      expect(await store.getOneTimePreKeyCount('bob')).toBe(3);
    });

    test('different addresses are isolated', async () => {
      await store.saveOneTimePreKeys('alice', [{ keyId: 1, publicKey: randBytes(32) }]);
      await store.saveOneTimePreKeys('bob', [{ keyId: 1, publicKey: randBytes(32) }, { keyId: 2, publicKey: randBytes(32) }]);

      expect(await store.getOneTimePreKeyCount('alice')).toBe(1);
      expect(await store.getOneTimePreKeyCount('bob')).toBe(2);
    });
  });

  describe('deleteAll', () => {
    test('removes everything for an address', async () => {
      await store.saveIdentity('bob', randBytes(32), randBytes(32));
      await store.saveSignedPreKey('bob', 1, randBytes(32), randBytes(64));
      await store.saveOneTimePreKeys('bob', [{ keyId: 100, publicKey: randBytes(32) }]);

      await store.deleteAll('bob');

      expect(await store.getIdentity('bob')).toBeNull();
      expect(await store.getSignedPreKey('bob')).toBeNull();
      expect(await store.getOneTimePreKeyCount('bob')).toBe(0);
    });

    test('does not affect other addresses', async () => {
      await store.saveIdentity('alice', randBytes(32), randBytes(32));
      await store.saveIdentity('bob', randBytes(32), randBytes(32));

      await store.deleteAll('bob');

      expect(await store.getIdentity('alice')).not.toBeNull();
      expect(await store.getIdentity('bob')).toBeNull();
    });
  });

  describe('persistence', () => {
    test('data survives close and reopen', async () => {
      await store.saveIdentity('bob', randBytes(32), randBytes(32));
      await store.saveOneTimePreKeys('bob', [
        { keyId: 1, publicKey: randBytes(32) },
        { keyId: 2, publicKey: randBytes(32) },
      ]);

      store.close();
      store = new SqlitePrekeyStore(dbPath);

      expect(await store.getIdentity('bob')).not.toBeNull();
      expect(await store.getOneTimePreKeyCount('bob')).toBe(2);
    });
  });
});