Shade/packages/shade-transport/tests/kt-transport.test.ts

import { describe, test, expect } from 'bun:test';
import { SubtleCryptoProvider, MemoryStorage } from '@shade/crypto-web';
import { ShadeSessionManager } from '@shade/core';
import {
  createPrekeyServerWithKT,
  MemoryPrekeyStore,
} from '@shade/server';
import {
  KTVerificationError,
  LightWitness,
  MemoryKTLogStore,
  computeLogId,
  signSth,
  type SignedTreeHead,
} from '@shade/key-transparency';
import { ShadeFetchTransport } from '../src/fetch-transport.js';

const crypto = new SubtleCryptoProvider();

describe('ShadeFetchTransport with KT verifier', () => {
  test('fetch verifies inclusion proof against pinned log key', async () => {
    const logKp = await crypto.generateEd25519KeyPair();
    const { app, kt } = await createPrekeyServerWithKT({
      crypto,
      store: new MemoryPrekeyStore(),
      disableRateLimit: true,
      keyTransparency: {
        store: new MemoryKTLogStore(),
        signingPrivateKey: logKp.privateKey,
        signingPublicKey: logKp.publicKey,
      },
    });

    const port = 20100 + Math.floor(Math.random() * 500);
    const handle = Bun.serve({ port, fetch: app.fetch });
    try {
      const baseUrl = `http://localhost:${port}`;

      // Bob registers
      const bobStorage = new MemoryStorage();
      const bobManager = new ShadeSessionManager(crypto, bobStorage);
      await bobManager.initialize();
      const bobIdentity = await bobStorage.getIdentityKeyPair();
      const bobTransport = new ShadeFetchTransport({
        baseUrl,
        crypto,
        signingPrivateKey: bobIdentity!.signingPrivateKey,
      });
      const bobOTPKs = await bobManager.generateOneTimePreKeys(3);
      const bobBundle = await bobManager.createPreKeyBundle();
      await bobTransport.register('bob', bobManager.getPublicIdentity(), bobBundle.signedPreKey, bobOTPKs);

      // Alice fetches with KT verifier — should succeed
      const observed: SignedTreeHead[] = [];
      const aliceTransport = new ShadeFetchTransport({
        baseUrl,
        crypto,
        keyTransparency: {
          mode: 'observe-strict',
          logPublicKey: logKp.publicKey,
          onObserveSth: async (sth) => {
            observed.push(sth);
          },
        },
      });
      const result = await aliceTransport.fetchBundleVerified('bob');
      expect(result.bundle.identityDHKey).toEqual(bobManager.getPublicIdentity().dhKey);
      expect(result.ktSth).toBeDefined();
      expect(result.ktSth!.treeSize).toBe(1);
      expect(observed.length).toBe(1);

      // Sanity: server-side latest STH matches what client observed
      const serverSth = await kt.getLatestSTH();
      expect(Buffer.from(serverSth.rootHash).toString('hex')).toBe(
        Buffer.from(result.ktSth!.rootHash).toString('hex'),
      );
    } finally {
      handle.stop();
    }
  });

  test('observe-strict throws KTVerificationError on missing proof', async () => {
    // Plain server (no KT)
    const { createPrekeyServer, MemoryPrekeyStore } = await import('@shade/server');
    const app = createPrekeyServer({
      crypto,
      store: new MemoryPrekeyStore(),
      disableRateLimit: true,
    });
    const port = 20800 + Math.floor(Math.random() * 200);
    const handle = Bun.serve({ port, fetch: app.fetch });
    try {
      const baseUrl = `http://localhost:${port}`;

      const bobStorage = new MemoryStorage();
      const bobManager = new ShadeSessionManager(crypto, bobStorage);
      await bobManager.initialize();
      const bobIdentity = await bobStorage.getIdentityKeyPair();
      const bobTransport = new ShadeFetchTransport({
        baseUrl,
        crypto,
        signingPrivateKey: bobIdentity!.signingPrivateKey,
      });
      const bobOTPKs = await bobManager.generateOneTimePreKeys(2);
      const bobBundle = await bobManager.createPreKeyBundle();
      await bobTransport.register('bob', bobManager.getPublicIdentity(), bobBundle.signedPreKey, bobOTPKs);

      const logKp = await crypto.generateEd25519KeyPair();
      const aliceTransport = new ShadeFetchTransport({
        baseUrl,
        crypto,
        keyTransparency: { mode: 'observe-strict', logPublicKey: logKp.publicKey },
      });
      await expect(aliceTransport.fetchBundle('bob')).rejects.toBeInstanceOf(KTVerificationError);
    } finally {
      handle.stop();
    }
  });

  test('forged STH (server signed with wrong key) is rejected', async () => {
    const realLogKp = await crypto.generateEd25519KeyPair();
    const evilLogKp = await crypto.generateEd25519KeyPair();
    const { app } = await createPrekeyServerWithKT({
      crypto,
      store: new MemoryPrekeyStore(),
      disableRateLimit: true,
      keyTransparency: {
        store: new MemoryKTLogStore(),
        // Operator signs with the EVIL key
        signingPrivateKey: evilLogKp.privateKey,
        signingPublicKey: evilLogKp.publicKey,
      },
    });
    const port = 21100 + Math.floor(Math.random() * 200);
    const handle = Bun.serve({ port, fetch: app.fetch });
    try {
      const baseUrl = `http://localhost:${port}`;

      const bobStorage = new MemoryStorage();
      const bobManager = new ShadeSessionManager(crypto, bobStorage);
      await bobManager.initialize();
      const bobIdentity = await bobStorage.getIdentityKeyPair();
      const bobTransport = new ShadeFetchTransport({
        baseUrl,
        crypto,
        signingPrivateKey: bobIdentity!.signingPrivateKey,
      });
      const bobOTPKs = await bobManager.generateOneTimePreKeys(2);
      const bobBundle = await bobManager.createPreKeyBundle();
      await bobTransport.register('bob', bobManager.getPublicIdentity(), bobBundle.signedPreKey, bobOTPKs);

      // Client pinned the REAL log key — verification must fail
      const aliceTransport = new ShadeFetchTransport({
        baseUrl,
        crypto,
        keyTransparency: { mode: 'observe-strict', logPublicKey: realLogKp.publicKey },
      });
      await expect(aliceTransport.fetchBundle('bob')).rejects.toThrow();
    } finally {
      handle.stop();
    }
  });

  test('observed STH feeds LightWitness; subsequent split-view detected', async () => {
    const logKp = await crypto.generateEd25519KeyPair();
    const { app, kt } = await createPrekeyServerWithKT({
      crypto,
      store: new MemoryPrekeyStore(),
      disableRateLimit: true,
      keyTransparency: {
        store: new MemoryKTLogStore(),
        signingPrivateKey: logKp.privateKey,
        signingPublicKey: logKp.publicKey,
      },
    });
    const port = 21500 + Math.floor(Math.random() * 200);
    const handle = Bun.serve({ port, fetch: app.fetch });
    try {
      const baseUrl = `http://localhost:${port}`;

      // Register Bob
      const bobStorage = new MemoryStorage();
      const bobManager = new ShadeSessionManager(crypto, bobStorage);
      await bobManager.initialize();
      const bobIdentity = await bobStorage.getIdentityKeyPair();
      const bobTransport = new ShadeFetchTransport({
        baseUrl,
        crypto,
        signingPrivateKey: bobIdentity!.signingPrivateKey,
      });
      const bobOTPKs = await bobManager.generateOneTimePreKeys(2);
      const bobBundle = await bobManager.createPreKeyBundle();
      await bobTransport.register('bob', bobManager.getPublicIdentity(), bobBundle.signedPreKey, bobOTPKs);

      // Witness backed by the real /v1/kt/* endpoints
      const witness = new LightWitness({
        crypto,
        logPublicKey: logKp.publicKey,
        fetcher: {
          async fetchLatestSTH() {
            const res = await fetch(`${baseUrl}/v1/kt/sth`);
            return res.json();
          },
          async fetchConsistencyProof(from, to) {
            const res = await fetch(`${baseUrl}/v1/kt/consistency?from=${from}&to=${to}`);
            return res.json();
          },
        },
      });
      const observedSth = await witness.pollOnce();
      expect(observedSth.treeSize).toBe(1);

      // Forge a divergent STH at the same tree_size and feed it to the
      // witness (this simulates a malicious second view)
      const realSth = await kt.getLatestSTH();
      const tampered = new Uint8Array(realSth.rootHash);
      tampered[0] ^= 0xff;
      const forged = await signSth(crypto, logKp.privateKey, {
        treeSize: realSth.treeSize,
        timestampMs: realSth.timestampMs,
        rootHash: tampered,
        indexRoot: realSth.indexRoot,
        logId: computeLogId(logKp.publicKey),
      });
      await expect(witness.observe(forged)).rejects.toThrow(/Split view/);
    } finally {
      handle.stop();
    }
  });
});