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

import { describe, test, expect, beforeEach } from 'bun:test';
import {
  createPrekeyServerWithKT,
  KeyTransparencyService,
  MemoryPrekeyStore,
  signPayload,
} from '../src/index.js';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import { generateIdentityKeyPair } from '@shade/core';
import {
  MemoryKTLogStore,
  computeBundleHash,
  ktProofFromWire,
  sthFromWire,
  verifyBundleAbsence,
  verifyBundleInclusion,
  verifyConsistencyProof,
  type KTProofWire,
  type STHWire,
} from '@shade/key-transparency';

const crypto = new SubtleCryptoProvider();

function b64(bytes: Uint8Array): string {
  return Buffer.from(bytes).toString('base64');
}

function fromB64(s: string): Uint8Array {
  return new Uint8Array(Buffer.from(s, 'base64'));
}

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

async function makeBundleData() {
  const identity = await generateIdentityKeyPair(crypto);
  const signedPreKeyPub = randBytes(32);
  const signedPreKeySig = await crypto.sign(identity.signingPrivateKey, signedPreKeyPub);
  return {
    identity,
    signedPreKey: {
      keyId: 1,
      publicKey: signedPreKeyPub,
      signature: signedPreKeySig,
    },
  };
}

describe('Prekey server with KT enabled', () => {
  let app: any;
  let kt: KeyTransparencyService;
  let logKp: { publicKey: Uint8Array; privateKey: Uint8Array };

  beforeEach(async () => {
    logKp = await crypto.generateEd25519KeyPair();
    const result = await createPrekeyServerWithKT({
      crypto,
      store: new MemoryPrekeyStore(),
      disableRateLimit: true,
      keyTransparency: {
        store: new MemoryKTLogStore(),
        signingPrivateKey: logKp.privateKey,
        signingPublicKey: logKp.publicKey,
      },
    });
    app = result.app;
    kt = result.kt;
  });

  async function registerAddress(address: string) {
    const data = await makeBundleData();
    const body: any = {
      address,
      identitySigningKey: b64(data.identity.signingPublicKey),
      identityDHKey: b64(data.identity.dhPublicKey),
      signedPreKey: {
        keyId: data.signedPreKey.keyId,
        publicKey: b64(data.signedPreKey.publicKey),
        signature: b64(data.signedPreKey.signature),
      },
      oneTimePreKeys: [{ keyId: 100, publicKey: b64(randBytes(32)) }],
    };
    const signed = await signPayload(crypto, data.identity.signingPrivateKey, body);
    const res = await app.request('/v1/keys/register', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify(signed),
    });
    expect(res.status).toBe(200);
    return data;
  }

  test('GET /v1/kt/log_id returns logId + publicKey', async () => {
    const res = await app.request('/v1/kt/log_id');
    expect(res.status).toBe(200);
    const body = await res.json();
    expect(body.logId).toBeDefined();
    expect(body.publicKey).toBeDefined();
    expect(b64(kt.getSigningPublicKey())).toBe(body.publicKey);
  });

  test('GET /v1/kt/sth returns latest STH', async () => {
    await registerAddress('alice');
    const res = await app.request('/v1/kt/sth');
    expect(res.status).toBe(200);
    const wire = (await res.json()) as STHWire;
    const sth = sthFromWire(wire, fromB64);
    expect(sth.treeSize).toBe(1);
  });

  test('bundle response carries verified inclusion proof', async () => {
    const data = await registerAddress('alice');
    const res = await app.request('/v1/keys/bundle/alice');
    expect(res.status).toBe(200);
    const body = (await res.json()) as { ktProof: KTProofWire } & Record<string, unknown>;
    expect(body.ktProof).toBeDefined();

    const proof = ktProofFromWire(body.ktProof);
    expect(proof.body.kind).toBe('inclusion');

    await verifyBundleInclusion(
      { crypto, logPublicKey: logKp.publicKey },
      'alice',
      {
        identitySigningKey: data.identity.signingPublicKey,
        identityDHKey: data.identity.dhPublicKey,
        signedPreKey: data.signedPreKey,
      },
      proof,
    );

    // Sanity: bundle hash matches the proof's index commitment
    const expected = computeBundleHash({
      identitySigningKey: data.identity.signingPublicKey,
      identityDHKey: data.identity.dhPublicKey,
      signedPreKey: data.signedPreKey,
    });
    if (proof.body.kind === 'inclusion') {
      expect(b64(proof.body.indexProof.entry.bundleHash)).toBe(b64(expected));
    }
  });

  test('bundle for unknown address returns 404 + absence proof', async () => {
    await registerAddress('alice');
    const res = await app.request('/v1/keys/bundle/zeta');
    expect(res.status).toBe(404);
    const body = (await res.json()) as { ktProof?: KTProofWire };
    expect(body.ktProof).toBeDefined();
    const proof = ktProofFromWire(body.ktProof!);
    expect(proof.body.kind).toBe('absence');
    await verifyBundleAbsence({ crypto, logPublicKey: logKp.publicKey }, 'zeta', proof);
  });

  test('DELETE /v1/keys/:address commits a tombstone', async () => {
    const data = await registerAddress('alice');
    const sthBefore = await kt.getLatestSTH();

    const signed = await signPayload(crypto, data.identity.signingPrivateKey, { address: 'alice' });
    const res = await app.request('/v1/keys/alice', {
      method: 'DELETE',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify(signed),
    });
    expect(res.status).toBe(200);

    const sthAfter = await kt.getLatestSTH();
    expect(sthAfter.treeSize).toBeGreaterThan(sthBefore.treeSize);

    const fetched = await app.request('/v1/keys/bundle/alice');
    // Identity row was removed by the prekey-store; absence-proof returned.
    expect(fetched.status).toBe(404);
    const body = (await fetched.json()) as { ktProof?: KTProofWire };
    expect(body.ktProof).toBeDefined();
  });

  test('GET /v1/kt/consistency returns valid proof', async () => {
    await registerAddress('alice');
    const sth1 = await kt.getLatestSTH();
    await registerAddress('bob');
    const sth2 = await kt.getLatestSTH();

    const res = await app.request(`/v1/kt/consistency?from=${sth1.treeSize}&to=${sth2.treeSize}`);
    expect(res.status).toBe(200);
    const body = (await res.json()) as { proof: string[] };
    const proofBytes = body.proof.map(fromB64);
    expect(verifyConsistencyProof(sth1.treeSize, sth2.treeSize, sth1.rootHash, sth2.rootHash, proofBytes)).toBe(true);
  });

  test('GET /v1/kt/sth/:treeSize returns historical STH', async () => {
    await registerAddress('alice');
    const sth1 = await kt.getLatestSTH();
    await registerAddress('bob');

    const res = await app.request(`/v1/kt/sth/${sth1.treeSize}`);
    expect(res.status).toBe(200);
    const wire = (await res.json()) as STHWire;
    const back = sthFromWire(wire, fromB64);
    expect(b64(back.rootHash)).toBe(b64(sth1.rootHash));
  });

  test('rotation: latest STH proof verifies with new bundle, not old', async () => {
    const v1Data = await registerAddress('alice');

    // Register again with a new identity (rotation)
    const v2Data = await makeBundleData();
    const body: any = {
      address: 'alice',
      identitySigningKey: b64(v2Data.identity.signingPublicKey),
      identityDHKey: b64(v2Data.identity.dhPublicKey),
      signedPreKey: {
        keyId: v2Data.signedPreKey.keyId,
        publicKey: b64(v2Data.signedPreKey.publicKey),
        signature: b64(v2Data.signedPreKey.signature),
      },
    };
    const signed = await signPayload(crypto, v2Data.identity.signingPrivateKey, body);
    const reRes = await app.request('/v1/keys/register', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify(signed),
    });
    expect(reRes.status).toBe(200);

    // Now the bundle response should reflect the new identity
    const res = await app.request('/v1/keys/bundle/alice');
    const body2 = (await res.json()) as {
      identitySigningKey: string;
      ktProof: KTProofWire;
    };
    expect(body2.identitySigningKey).toBe(b64(v2Data.identity.signingPublicKey));

    const proof = ktProofFromWire(body2.ktProof);
    await verifyBundleInclusion(
      { crypto, logPublicKey: logKp.publicKey },
      'alice',
      {
        identitySigningKey: v2Data.identity.signingPublicKey,
        identityDHKey: v2Data.identity.dhPublicKey,
        signedPreKey: v2Data.signedPreKey,
      },
      proof,
    );
    // Verifying with v1 data should reject
    await expect(
      verifyBundleInclusion(
        { crypto, logPublicKey: logKp.publicKey },
        'alice',
        {
          identitySigningKey: v1Data.identity.signingPublicKey,
          identityDHKey: v1Data.identity.dhPublicKey,
          signedPreKey: v1Data.signedPreKey,
        },
        proof,
      ),
    ).rejects.toThrow();
  });
});