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

import { describe, test, expect, beforeEach, afterEach } from 'bun:test';
import { tmpdir } from 'os';
import { join } from 'path';
import { unlinkSync } from 'fs';
import { Database } from 'bun:sqlite';
import { EncryptedSQLiteStorage } from '../src/storage/encrypted-sqlite.js';
import { KeyManager } from '../src/crypto/key-manager.js';
import type { IdentityKeyPair, SignedPreKey, OneTimePreKey, SessionState, PersistedStreamState } from '@shade/core';

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

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

const KEY_BYTES = randBytes(32);
async function freshKM(): Promise<KeyManager> {
  return KeyManager.open({ kind: 'injected', key: KEY_BYTES });
}

function dummyIdentity(): IdentityKeyPair {
  return {
    signingPublicKey: randBytes(32),
    signingPrivateKey: randBytes(32),
    dhPublicKey: randBytes(32),
    dhPrivateKey: randBytes(32),
  };
}

function dummySignedPreKey(id: number): SignedPreKey {
  return {
    keyId: id,
    keyPair: { publicKey: randBytes(32), privateKey: randBytes(32) },
    signature: randBytes(64),
    timestamp: Date.now(),
  };
}

function dummyOTP(id: number): OneTimePreKey {
  return {
    keyId: id,
    keyPair: { publicKey: randBytes(32), privateKey: randBytes(32) },
  };
}

function dummySession(): SessionState {
  return {
    remoteIdentityKey: randBytes(32),
    rootKey: randBytes(32),
    sendChain: { chainKey: randBytes(32), counter: 0 },
    receiveChain: { chainKey: randBytes(32), counter: 0 },
    dhSend: { publicKey: randBytes(32), privateKey: randBytes(32) },
    dhReceive: randBytes(32),
    previousSendCounter: 0,
    skippedKeys: new Map(),
  };
}

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

  beforeEach(async () => {
    dbPath = tempDb();
    store = await EncryptedSQLiteStorage.open({ dbPath, keyManager: await freshKM() });
  });

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

  test('identity round-trip', async () => {
    expect(await store.getIdentityKeyPair()).toBeNull();
    const kp = dummyIdentity();
    await store.saveIdentityKeyPair(kp);
    const got = await store.getIdentityKeyPair();
    expect(got).toEqual(kp);
  });

  test('registrationId round-trip', async () => {
    expect(await store.getLocalRegistrationId()).toBe(0);
    await store.saveLocalRegistrationId(12345);
    expect(await store.getLocalRegistrationId()).toBe(12345);
  });

  test('signed prekey round-trip + remove', async () => {
    expect(await store.getSignedPreKey(7)).toBeNull();
    const k = dummySignedPreKey(7);
    await store.saveSignedPreKey(k);
    const got = await store.getSignedPreKey(7);
    expect(got?.keyId).toBe(7);
    expect(got?.keyPair.privateKey).toEqual(k.keyPair.privateKey);
    await store.removeSignedPreKey(7);
    expect(await store.getSignedPreKey(7)).toBeNull();
  });

  test('one-time prekey round-trip + count + remove', async () => {
    expect(await store.getOneTimePreKeyCount()).toBe(0);
    await store.saveOneTimePreKey(dummyOTP(1));
    await store.saveOneTimePreKey(dummyOTP(2));
    expect(await store.getOneTimePreKeyCount()).toBe(2);
    expect(await store.getOneTimePreKey(1)).not.toBeNull();
    await store.removeOneTimePreKey(1);
    expect(await store.getOneTimePreKey(1)).toBeNull();
    expect(await store.getOneTimePreKeyCount()).toBe(1);
  });

  test('session round-trip + remove', async () => {
    const s = dummySession();
    await store.saveSession('device:abc', s);
    const got = await store.getSession('device:abc');
    expect(got?.rootKey).toEqual(s.rootKey);
    await store.removeSession('device:abc');
    expect(await store.getSession('device:abc')).toBeNull();
  });

  test('TOFU + trust check', async () => {
    const ik = randBytes(32);
    expect(await store.isTrustedIdentity('peer-1', ik)).toBe(true); // TOFU
    await store.saveTrustedIdentity('peer-1', ik);
    expect(await store.isTrustedIdentity('peer-1', ik)).toBe(true);
    expect(await store.isTrustedIdentity('peer-1', randBytes(32))).toBe(false);
  });

  test('retired identities are sorted DESC', async () => {
    await store.addRetiredIdentity({ keyPair: dummyIdentity(), retiredAt: 100 });
    await store.addRetiredIdentity({ keyPair: dummyIdentity(), retiredAt: 200 });
    const list = await store.getRetiredIdentities();
    expect(list.length).toBe(2);
    expect(list[0]!.retiredAt).toBe(200);
    await store.pruneRetiredIdentities(150);
    expect((await store.getRetiredIdentities()).length).toBe(1);
  });

  test('stream-state round-trip + listActive + prune', async () => {
    const s: PersistedStreamState = {
      streamId: 'stream-1',
      direction: 'send',
      peerAddress: 'device:bob',
      status: 'active',
      metadataJson: '{"name":"file.bin"}',
      partitionJson: '[]',
      laneStateJson: '[]',
      ioDescriptorJson: '{"path":"/tmp/x"}',
      secretEnc: randBytes(32),
      secretNonce: randBytes(12),
      createdAt: 1,
      updatedAt: 2,
    };
    await store.saveStreamState(s);
    const got = await store.getStreamState('stream-1');
    expect(got).toEqual(s);
    const active = await store.listActiveStreamStates();
    expect(active.length).toBe(1);
    expect((await store.listActiveStreamStates('receive')).length).toBe(0);

    await store.saveStreamState({ ...s, streamId: 'stream-2', status: 'finished', updatedAt: 50 });
    expect((await store.listActiveStreamStates()).length).toBe(1); // only stream-1 still active
    await store.pruneStreamStates(100);
    expect(await store.getStreamState('stream-2')).toBeNull();
    expect(await store.getStreamState('stream-1')).not.toBeNull(); // active rows untouched
  });

  test('rejects open with wrong key (fingerprint mismatch)', async () => {
    await store.saveIdentityKeyPair(dummyIdentity());
    store.close();
    const otherKey = randBytes(32);
    await expect(EncryptedSQLiteStorage.open({
      dbPath,
      keyManager: await KeyManager.open({ kind: 'injected', key: otherKey }),
    })).rejects.toThrow(/storage key mismatch/);
    // Reopen with original key for afterEach
    store = await EncryptedSQLiteStorage.open({ dbPath, keyManager: await freshKM() });
  });
});

describe('EncryptedSQLiteStorage — tamper detection', () => {
  test('flipped ciphertext byte → decrypt fails', async () => {
    const dbPath = tempDb();
    const store = await EncryptedSQLiteStorage.open({ dbPath, keyManager: await freshKM() });
    await store.saveIdentityKeyPair(dummyIdentity());
    store.close();

    // Tamper with the ciphertext directly via raw SQLite.
    const db = new Database(dbPath);
    const row = db.prepare('SELECT ciphertext FROM identity_enc WHERE id = 1').get() as { ciphertext: Uint8Array };
    const ct = new Uint8Array(row.ciphertext);
    ct[ct.length - 1]! ^= 0x01;
    db.prepare('UPDATE identity_enc SET ciphertext = ? WHERE id = 1').run(ct);
    db.close();

    const reopened = await EncryptedSQLiteStorage.open({ dbPath, keyManager: await freshKM() });
    await expect(reopened.getIdentityKeyPair()).rejects.toThrow();
    reopened.close();
    try { unlinkSync(dbPath); } catch {}
    try { unlinkSync(dbPath + '-wal'); } catch {}
    try { unlinkSync(dbPath + '-shm'); } catch {}
  });

  test('row swap (sessions) → decrypt fails due to AAD mismatch', async () => {
    const dbPath = tempDb();
    const store = await EncryptedSQLiteStorage.open({ dbPath, keyManager: await freshKM() });
    await store.saveSession('alice', dummySession());
    await store.saveSession('bob',   dummySession());
    store.close();

    // Swap the ciphertexts.
    const db = new Database(dbPath);
    const aliceRow = db.prepare('SELECT ciphertext FROM sessions_enc WHERE address = ?').get('alice') as { ciphertext: Uint8Array };
    const bobRow   = db.prepare('SELECT ciphertext FROM sessions_enc WHERE address = ?').get('bob')   as { ciphertext: Uint8Array };
    db.prepare('UPDATE sessions_enc SET ciphertext = ? WHERE address = ?').run(bobRow.ciphertext, 'alice');
    db.prepare('UPDATE sessions_enc SET ciphertext = ? WHERE address = ?').run(aliceRow.ciphertext, 'bob');
    db.close();

    const reopened = await EncryptedSQLiteStorage.open({ dbPath, keyManager: await freshKM() });
    await expect(reopened.getSession('alice')).rejects.toThrow();
    await expect(reopened.getSession('bob')).rejects.toThrow();
    reopened.close();
    try { unlinkSync(dbPath); } catch {}
    try { unlinkSync(dbPath + '-wal'); } catch {}
    try { unlinkSync(dbPath + '-shm'); } catch {}
  });
});