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

import { describe, test, expect, beforeAll, afterAll } from 'bun:test';
import postgres, { type Sql } from 'postgres';
import { PostgresStorage } from '../src/index.js';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import { ShadeSessionManager } from '@shade/core';
import type { IdentityKeyPair } from '@shade/core';

const crypto = new SubtleCryptoProvider();

/**
 * These tests require a running PostgreSQL instance.
 * Skip if SHADE_TEST_PG_URL is not set.
 *
 * To run locally:
 *   docker run -d --name shade-test-pg -e POSTGRES_PASSWORD=test -p 5999:5432 postgres:16
 *   SHADE_TEST_PG_URL=postgres://postgres:test@localhost:5999/postgres bun test postgres-storage
 */
const PG_URL = process.env.SHADE_TEST_PG_URL;

if (!PG_URL) {
  describe.skip('PostgresStorage (skipped — set SHADE_TEST_PG_URL to enable)', () => {
    test('placeholder', () => {});
  });
} else {
  describe('PostgresStorage', () => {
    let sql: Sql;
    let storage: PostgresStorage;

    beforeAll(async () => {
      sql = postgres(PG_URL!);
      // Clean slate
      await sql`DROP TABLE IF EXISTS shade_identity CASCADE`;
      await sql`DROP TABLE IF EXISTS shade_config CASCADE`;
      await sql`DROP TABLE IF EXISTS shade_signed_prekeys CASCADE`;
      await sql`DROP TABLE IF EXISTS shade_one_time_prekeys CASCADE`;
      await sql`DROP TABLE IF EXISTS shade_sessions CASCADE`;
      await sql`DROP TABLE IF EXISTS shade_trusted_identities CASCADE`;
      await sql`DROP TABLE IF EXISTS shade_retired_identities CASCADE`;
      storage = await PostgresStorage.fromClient(sql);
    });

    afterAll(async () => {
      await sql.end();
    });

    test('identity save/load roundtrip', async () => {
      const ikp: IdentityKeyPair = {
        signingPublicKey: crypto.randomBytes(32),
        signingPrivateKey: crypto.randomBytes(32),
        dhPublicKey: crypto.randomBytes(32),
        dhPrivateKey: crypto.randomBytes(32),
      };
      await storage.saveIdentityKeyPair(ikp);
      const restored = await storage.getIdentityKeyPair();
      expect(restored!.signingPublicKey).toEqual(ikp.signingPublicKey);
      expect(restored!.dhPrivateKey).toEqual(ikp.dhPrivateKey);
    });

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

    test('signed prekey crud', async () => {
      await storage.saveSignedPreKey({
        keyId: 1,
        keyPair: { publicKey: crypto.randomBytes(32), privateKey: crypto.randomBytes(32) },
        signature: crypto.randomBytes(64),
        timestamp: Date.now(),
      });
      const spk = await storage.getSignedPreKey(1);
      expect(spk!.keyId).toBe(1);
      expect(spk!.signature.length).toBe(64);

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

    test('one-time prekey count and crud', async () => {
      await storage.saveOneTimePreKey({
        keyId: 100,
        keyPair: { publicKey: crypto.randomBytes(32), privateKey: crypto.randomBytes(32) },
      });
      const count = await storage.getOneTimePreKeyCount();
      expect(count).toBeGreaterThanOrEqual(1);
      await storage.removeOneTimePreKey(100);
    });

    test('session save/load with skipped keys', async () => {
      const state = {
        remoteIdentityKey: crypto.randomBytes(32),
        rootKey: crypto.randomBytes(32),
        sendChain: { chainKey: crypto.randomBytes(32), counter: 7 },
        receiveChain: null,
        dhSend: { publicKey: crypto.randomBytes(32), privateKey: crypto.randomBytes(32) },
        dhReceive: null,
        previousSendCounter: 0,
        skippedKeys: new Map([['key:1', crypto.randomBytes(32)]]),
      };
      await storage.saveSession('bob', state);
      const restored = await storage.getSession('bob');
      expect(restored!.sendChain.counter).toBe(7);
      expect(restored!.skippedKeys.size).toBe(1);
      await storage.removeSession('bob');
    });

    test('constant-time trust comparison', async () => {
      const key = crypto.randomBytes(32);
      // TOFU
      expect(await storage.isTrustedIdentity('peer1', key)).toBe(true);
      await storage.saveTrustedIdentity('peer1', key);
      expect(await storage.isTrustedIdentity('peer1', key)).toBe(true);
      expect(await storage.isTrustedIdentity('peer1', crypto.randomBytes(32))).toBe(false);
    });

    test('retired identity history', async () => {
      const ikp: IdentityKeyPair = {
        signingPublicKey: crypto.randomBytes(32),
        signingPrivateKey: crypto.randomBytes(32),
        dhPublicKey: crypto.randomBytes(32),
        dhPrivateKey: crypto.randomBytes(32),
      };
      const now = Date.now();
      await storage.addRetiredIdentity({ keyPair: ikp, retiredAt: now });

      const retired = await storage.getRetiredIdentities();
      expect(retired.length).toBeGreaterThanOrEqual(1);
      const found = retired.find((r) => r.retiredAt === now);
      expect(found).toBeDefined();
      expect(found!.keyPair.signingPublicKey).toEqual(ikp.signingPublicKey);

      await storage.pruneRetiredIdentities(now + 1);
      const after = await storage.getRetiredIdentities();
      expect(after.find((r) => r.retiredAt === now)).toBeUndefined();
    });

    test('full E2EE conversation using PostgresStorage', async () => {
      // Clean storage for this test
      await sql`TRUNCATE shade_identity, shade_config, shade_signed_prekeys, shade_one_time_prekeys, shade_sessions, shade_trusted_identities`;

      const aliceStorage = storage;
      const bobSql = postgres(PG_URL!);
      const bobStorage = await PostgresStorage.fromClient(bobSql);

      try {
        const alice = new ShadeSessionManager(crypto, aliceStorage);
        const bob = new ShadeSessionManager(crypto, bobStorage);

        // Bob uses a separate fake storage since we only have one PG schema
        // For this smoke test, we just verify the storage layer itself works
        await alice.initialize();
        expect(await alice.getIdentityFingerprint()).toBeTruthy();
      } finally {
        await bobSql.end();
      }
    });
  });
}