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

import { describe, test, expect, beforeEach, afterEach } from 'bun:test';
import { tmpdir } from 'os';
import { join } from 'path';
import { unlinkSync, existsSync, readdirSync } from 'fs';
import { Database } from 'bun:sqlite';
import { SQLiteStorage } from '@shade/storage-sqlite';
import { EncryptedSQLiteStorage } from '../src/storage/encrypted-sqlite.js';
import { KeyManager } from '../src/crypto/key-manager.js';
import { migrateSqliteToEncrypted, rotateSqliteEncryptionKey } from '../src/migrate/migrate-sqlite.js';

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

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

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

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

function dummySession() {
  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<string, Uint8Array>(),
  };
}

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

  beforeEach(() => {
    dbPath = tempDb();
  });

  afterEach(() => {
    for (const f of [dbPath, dbPath + '-wal', dbPath + '-shm']) {
      try { unlinkSync(f); } catch {}
    }
    // Clean up any .bak files left in the temp dir.
    const dir = tmpdir();
    for (const name of readdirSync(dir)) {
      if (name.startsWith(`shade-migrate-`) && name.includes('.bak.')) {
        try { unlinkSync(join(dir, name)); } catch {}
      }
    }
  });

  test('migrates a populated unencrypted DB', async () => {
    const id = dummyIdentity();
    const sk = dummySignedPreKey(1);
    const sess = dummySession();

    const src = new SQLiteStorage(dbPath);
    await src.saveIdentityKeyPair(id);
    await src.saveLocalRegistrationId(99);
    await src.saveSignedPreKey(sk);
    await src.saveSession('alice', sess);
    await src.saveTrustedIdentity('alice', id.dhPublicKey);
    await src.addRetiredIdentity({ keyPair: dummyIdentity(), retiredAt: 1234 });
    src.close();

    const masterKey = randBytes(32);
    const km = await KeyManager.open({ kind: 'injected', key: masterKey });
    const report = await migrateSqliteToEncrypted({ dbPath, keyManager: km, backup: false });

    expect(report.identity).toBe(1);
    expect(report.config).toBe(1);
    expect(report.signedPrekeys).toBe(1);
    expect(report.sessions).toBe(1);
    expect(report.trustedIdentities).toBe(1);
    expect(report.retiredIdentities).toBe(1);

    // Verify we can read everything back with the same masterKey.
    const km2 = await KeyManager.open({ kind: 'injected', key: masterKey });
    const enc = await EncryptedSQLiteStorage.open({ dbPath, keyManager: km2 });
    expect(await enc.getIdentityKeyPair()).not.toBeNull();
    expect(await enc.getLocalRegistrationId()).toBe(99);
    expect(await enc.getSession('alice')).not.toBeNull();
    expect(await enc.getSignedPreKey(1)).not.toBeNull();
    expect((await enc.getRetiredIdentities()).length).toBe(1);
    enc.close();
  });

  test('--dry-run leaves DB unchanged', async () => {
    const src = new SQLiteStorage(dbPath);
    await src.saveIdentityKeyPair(dummyIdentity());
    await src.saveSession('alice', dummySession());
    src.close();

    const km = await KeyManager.open({ kind: 'injected', key: randBytes(32) });
    const report = await migrateSqliteToEncrypted({ dbPath, keyManager: km, dryRun: true, backup: false });
    expect(report.dryRun).toBe(true);

    // Original tables still present and populated.
    const db = new Database(dbPath);
    const idCount = (db.prepare('SELECT COUNT(*) as c FROM identity').get() as { c: number }).c;
    expect(idCount).toBe(1);
    const sessCount = (db.prepare('SELECT COUNT(*) as c FROM sessions').get() as { c: number }).c;
    expect(sessCount).toBe(1);
    db.close();
  });

  test('drops unencrypted tables after successful migration', async () => {
    const src = new SQLiteStorage(dbPath);
    await src.saveIdentityKeyPair(dummyIdentity());
    src.close();

    const km = await KeyManager.open({ kind: 'injected', key: randBytes(32) });
    await migrateSqliteToEncrypted({ dbPath, keyManager: km, backup: false });

    const db = new Database(dbPath);
    const r = db.prepare("SELECT name FROM sqlite_master WHERE type='table' AND name='identity'").get();
    expect(r).toBeNull();
    db.close();
  });

  test('produces .bak file when backup enabled', async () => {
    const src = new SQLiteStorage(dbPath);
    await src.saveIdentityKeyPair(dummyIdentity());
    src.close();

    const km = await KeyManager.open({ kind: 'injected', key: randBytes(32) });
    const report = await migrateSqliteToEncrypted({ dbPath, keyManager: km, backup: true });
    expect(report.backupPath).toBeDefined();
    expect(existsSync(report.backupPath!)).toBe(true);
    try { unlinkSync(report.backupPath!); } catch {}
  });
});

describe('rotateSqliteEncryptionKey', () => {
  test('re-keys all rows; old key no longer opens DB', async () => {
    const dbPath = tempDb();
    const oldKeyBytes = randBytes(32);
    const newKeyBytes = randBytes(32);

    const oldKm = await KeyManager.open({ kind: 'injected', key: oldKeyBytes });
    const store = await EncryptedSQLiteStorage.open({ dbPath, keyManager: oldKm });
    await store.saveIdentityKeyPair(dummyIdentity());
    await store.saveSession('alice', dummySession());
    await store.saveSignedPreKey(dummySignedPreKey(1));
    store.close();

    const oldKmAgain = await KeyManager.open({ kind: 'injected', key: oldKeyBytes });
    const newKm = await KeyManager.open({ kind: 'injected', key: newKeyBytes });
    const result = await rotateSqliteEncryptionKey({
      dbPath,
      oldKeyManager: oldKmAgain,
      newKeyManager: newKm,
    });
    expect(result.rowsRotated).toBeGreaterThan(0);

    // Old key is rejected.
    const oldKmOnceMore = await KeyManager.open({ kind: 'injected', key: oldKeyBytes });
    await expect(EncryptedSQLiteStorage.open({ dbPath, keyManager: oldKmOnceMore }))
      .rejects.toThrow(/storage key mismatch/);

    // New key works.
    const newKmAgain = await KeyManager.open({ kind: 'injected', key: newKeyBytes });
    const reopened = await EncryptedSQLiteStorage.open({ dbPath, keyManager: newKmAgain });
    expect(await reopened.getIdentityKeyPair()).not.toBeNull();
    expect(await reopened.getSession('alice')).not.toBeNull();
    reopened.close();

    for (const f of [dbPath, dbPath + '-wal', dbPath + '-shm']) {
      try { unlinkSync(f); } catch {}
    }
  });
});