Shade/packages/shade-core/tests/x3dh.test.ts

import { describe, test, expect, beforeEach } from 'bun:test';
import { SubtleCryptoProvider, MemoryStorage } from '@shade/crypto-web';
import {
  generateIdentityKeyPair,
  generateSignedPreKey,
  generateOneTimePreKeys,
  createPreKeyBundle,
  processPreKeyBundle,
  processPreKeyMessage,
  InvalidSignatureError,
  PreKeyNotFoundError,
} from '../src/index.js';
import type { RatchetMessage } from '../src/index.js';

const crypto = new SubtleCryptoProvider();

/** Create a dummy RatchetMessage for testing (X3DH doesn't care about the content) */
function dummyRatchetMessage(): RatchetMessage {
  return {
    dhPublicKey: crypto.randomBytes(32),
    previousCounter: 0,
    counter: 0,
    ciphertext: crypto.randomBytes(48),
    nonce: crypto.randomBytes(12),
  };
}

describe('X3DH', () => {
  let aliceStorage: MemoryStorage;
  let bobStorage: MemoryStorage;

  beforeEach(() => {
    aliceStorage = new MemoryStorage();
    bobStorage = new MemoryStorage();
  });

  // ─── Key Generation ────────────────────────────────────────

  describe('key generation', () => {
    test('generates identity keypair with correct lengths', async () => {
      const id = await generateIdentityKeyPair(crypto);
      expect(id.signingPublicKey.length).toBe(32);
      expect(id.signingPrivateKey.length).toBe(32);
      expect(id.dhPublicKey.length).toBe(32);
      expect(id.dhPrivateKey.length).toBe(32);
    });

    test('signing and DH keys are different', async () => {
      const id = await generateIdentityKeyPair(crypto);
      expect(id.signingPublicKey).not.toEqual(id.dhPublicKey);
      expect(id.signingPrivateKey).not.toEqual(id.dhPrivateKey);
    });

    test('generates signed prekey with valid signature', async () => {
      const id = await generateIdentityKeyPair(crypto);
      const spk = await generateSignedPreKey(crypto, id, 1);

      expect(spk.keyId).toBe(1);
      expect(spk.keyPair.publicKey.length).toBe(32);
      expect(spk.signature.length).toBe(64);

      // Verify the signature
      const valid = await crypto.verify(id.signingPublicKey, spk.keyPair.publicKey, spk.signature);
      expect(valid).toBe(true);
    });

    test('generates batch of one-time prekeys', async () => {
      const otpks = await generateOneTimePreKeys(crypto, 100, 5);
      expect(otpks.length).toBe(5);

      for (let i = 0; i < 5; i++) {
        expect(otpks[i].keyId).toBe(100 + i);
        expect(otpks[i].keyPair.publicKey.length).toBe(32);
      }

      // All keys are unique
      const pubKeys = otpks.map((k) => Array.from(k.keyPair.publicKey).join(','));
      expect(new Set(pubKeys).size).toBe(5);
    });
  });

  // ─── Full Handshake ────────────────────────────────────────

  describe('full handshake', () => {
    test('Alice and Bob derive the same root key (with one-time prekey)', async () => {
      // Bob generates keys
      const bobIdentity = await generateIdentityKeyPair(crypto);
      const bobSignedPreKey = await generateSignedPreKey(crypto, bobIdentity, 1);
      const bobOneTimePreKeys = await generateOneTimePreKeys(crypto, 100, 3);

      // Bob stores his keys
      await bobStorage.saveIdentityKeyPair(bobIdentity);
      await bobStorage.saveSignedPreKey(bobSignedPreKey);
      for (const otpk of bobOneTimePreKeys) {
        await bobStorage.saveOneTimePreKey(otpk);
      }

      // Bob publishes a bundle (server would store this)
      const bundle = createPreKeyBundle(42, bobIdentity, bobSignedPreKey, bobOneTimePreKeys[0]);

      // Alice generates her identity
      const aliceIdentity = await generateIdentityKeyPair(crypto);
      await aliceStorage.saveIdentityKeyPair(aliceIdentity);

      // Alice processes the bundle
      const aliceResult = await processPreKeyBundle(crypto, aliceStorage, bundle);

      expect(aliceResult.rootKey.length).toBe(32);
      expect(aliceResult.signedPreKeyId).toBe(1);
      expect(aliceResult.preKeyId).toBe(100);
      expect(aliceResult.ephemeralPublicKey.length).toBe(32);

      // Alice creates a PreKeyMessage
      const preKeyMessage = {
        registrationId: 1,
        preKeyId: aliceResult.preKeyId,
        signedPreKeyId: aliceResult.signedPreKeyId,
        ephemeralKey: aliceResult.ephemeralPublicKey,
        identityDHKey: aliceIdentity.dhPublicKey,
        message: dummyRatchetMessage(),
      };

      // Bob processes the PreKeyMessage
      const bobResult = await processPreKeyMessage(crypto, bobStorage, preKeyMessage);

      // Both derive the same root key
      expect(bobResult.rootKey).toEqual(aliceResult.rootKey);
      expect(bobResult.remoteIdentityKey).toEqual(aliceIdentity.dhPublicKey);
      expect(bobResult.remoteEphemeralKey).toEqual(aliceResult.ephemeralPublicKey);
    });

    test('Alice and Bob derive the same root key (without one-time prekey)', async () => {
      const bobIdentity = await generateIdentityKeyPair(crypto);
      const bobSignedPreKey = await generateSignedPreKey(crypto, bobIdentity, 1);
      await bobStorage.saveIdentityKeyPair(bobIdentity);
      await bobStorage.saveSignedPreKey(bobSignedPreKey);

      // Bundle without one-time prekey
      const bundle = createPreKeyBundle(42, bobIdentity, bobSignedPreKey);

      const aliceIdentity = await generateIdentityKeyPair(crypto);
      await aliceStorage.saveIdentityKeyPair(aliceIdentity);

      const aliceResult = await processPreKeyBundle(crypto, aliceStorage, bundle);
      expect(aliceResult.preKeyId).toBeUndefined();

      const preKeyMessage = {
        registrationId: 1,
        signedPreKeyId: aliceResult.signedPreKeyId,
        ephemeralKey: aliceResult.ephemeralPublicKey,
        identityDHKey: aliceIdentity.dhPublicKey,
        message: dummyRatchetMessage(),
      };

      const bobResult = await processPreKeyMessage(crypto, bobStorage, preKeyMessage);
      expect(bobResult.rootKey).toEqual(aliceResult.rootKey);
    });

    test('different handshakes produce different root keys', async () => {
      const bobIdentity = await generateIdentityKeyPair(crypto);
      const bobSignedPreKey = await generateSignedPreKey(crypto, bobIdentity, 1);
      await bobStorage.saveIdentityKeyPair(bobIdentity);
      await bobStorage.saveSignedPreKey(bobSignedPreKey);

      const bundle = createPreKeyBundle(42, bobIdentity, bobSignedPreKey);

      // Alice 1
      const alice1Id = await generateIdentityKeyPair(crypto);
      const alice1Storage = new MemoryStorage();
      await alice1Storage.saveIdentityKeyPair(alice1Id);
      const result1 = await processPreKeyBundle(crypto, alice1Storage, bundle);

      // Alice 2 (different identity)
      const alice2Id = await generateIdentityKeyPair(crypto);
      const alice2Storage = new MemoryStorage();
      await alice2Storage.saveIdentityKeyPair(alice2Id);
      const result2 = await processPreKeyBundle(crypto, alice2Storage, bundle);

      expect(result1.rootKey).not.toEqual(result2.rootKey);
    });
  });

  // ─── Signature Verification ────────────────────────────────

  describe('signature verification', () => {
    test('rejects bundle with invalid signed prekey signature', async () => {
      const bobIdentity = await generateIdentityKeyPair(crypto);
      const bobSignedPreKey = await generateSignedPreKey(crypto, bobIdentity, 1);

      // Tamper with the signature
      const bundle = createPreKeyBundle(42, bobIdentity, bobSignedPreKey);
      bundle.signedPreKey.signature[0] ^= 0xff;

      const aliceIdentity = await generateIdentityKeyPair(crypto);
      await aliceStorage.saveIdentityKeyPair(aliceIdentity);

      expect(processPreKeyBundle(crypto, aliceStorage, bundle)).rejects.toThrow(InvalidSignatureError);
    });

    test('rejects bundle with wrong identity key signing', async () => {
      const bobIdentity = await generateIdentityKeyPair(crypto);
      const eveIdentity = await generateIdentityKeyPair(crypto);

      // Eve signs the prekey, but claims to be Bob
      const eveSignedPreKey = await generateSignedPreKey(crypto, eveIdentity, 1);
      const bundle = createPreKeyBundle(42, bobIdentity, eveSignedPreKey);

      const aliceIdentity = await generateIdentityKeyPair(crypto);
      await aliceStorage.saveIdentityKeyPair(aliceIdentity);

      expect(processPreKeyBundle(crypto, aliceStorage, bundle)).rejects.toThrow(InvalidSignatureError);
    });
  });

  // ─── One-Time Prekey Consumption ───────────────────────────

  describe('one-time prekey consumption', () => {
    test('one-time prekey is deleted after use', async () => {
      const bobIdentity = await generateIdentityKeyPair(crypto);
      const bobSignedPreKey = await generateSignedPreKey(crypto, bobIdentity, 1);
      const bobOTPKs = await generateOneTimePreKeys(crypto, 100, 3);
      await bobStorage.saveIdentityKeyPair(bobIdentity);
      await bobStorage.saveSignedPreKey(bobSignedPreKey);
      for (const otpk of bobOTPKs) await bobStorage.saveOneTimePreKey(otpk);

      expect(await bobStorage.getOneTimePreKeyCount()).toBe(3);

      const bundle = createPreKeyBundle(42, bobIdentity, bobSignedPreKey, bobOTPKs[0]);

      const aliceIdentity = await generateIdentityKeyPair(crypto);
      await aliceStorage.saveIdentityKeyPair(aliceIdentity);

      const aliceResult = await processPreKeyBundle(crypto, aliceStorage, bundle);
      const preKeyMessage = {
        registrationId: 1,
        preKeyId: aliceResult.preKeyId,
        signedPreKeyId: aliceResult.signedPreKeyId,
        ephemeralKey: aliceResult.ephemeralPublicKey,
        identityDHKey: aliceIdentity.dhPublicKey,
        message: dummyRatchetMessage(),
      };

      await processPreKeyMessage(crypto, bobStorage, preKeyMessage);

      // One-time prekey 100 should be consumed
      expect(await bobStorage.getOneTimePreKeyCount()).toBe(2);
      expect(await bobStorage.getOneTimePreKey(100)).toBeNull();
      // Others remain
      expect(await bobStorage.getOneTimePreKey(101)).not.toBeNull();
      expect(await bobStorage.getOneTimePreKey(102)).not.toBeNull();
    });

    test('fails when referenced one-time prekey does not exist', async () => {
      const bobIdentity = await generateIdentityKeyPair(crypto);
      const bobSignedPreKey = await generateSignedPreKey(crypto, bobIdentity, 1);
      await bobStorage.saveIdentityKeyPair(bobIdentity);
      await bobStorage.saveSignedPreKey(bobSignedPreKey);
      // No one-time prekeys stored

      const aliceIdentity = await generateIdentityKeyPair(crypto);
      await aliceStorage.saveIdentityKeyPair(aliceIdentity);

      const preKeyMessage = {
        registrationId: 1,
        preKeyId: 999, // doesn't exist
        signedPreKeyId: 1,
        ephemeralKey: crypto.randomBytes(32),
        identityDHKey: aliceIdentity.dhPublicKey,
        message: dummyRatchetMessage(),
      };

      expect(processPreKeyMessage(crypto, bobStorage, preKeyMessage)).rejects.toThrow(PreKeyNotFoundError);
    });

    test('fails when referenced signed prekey does not exist', async () => {
      const bobIdentity = await generateIdentityKeyPair(crypto);
      await bobStorage.saveIdentityKeyPair(bobIdentity);
      // No signed prekey stored

      const preKeyMessage = {
        registrationId: 1,
        signedPreKeyId: 999,
        ephemeralKey: crypto.randomBytes(32),
        identityDHKey: crypto.randomBytes(32),
        message: dummyRatchetMessage(),
      };

      expect(processPreKeyMessage(crypto, bobStorage, preKeyMessage)).rejects.toThrow(PreKeyNotFoundError);
    });
  });

  // ─── PreKey Bundle Assembly ────────────────────────────────

  describe('createPreKeyBundle', () => {
    test('assembles bundle with one-time prekey', async () => {
      const id = await generateIdentityKeyPair(crypto);
      const spk = await generateSignedPreKey(crypto, id, 5);
      const otpk = (await generateOneTimePreKeys(crypto, 200, 1))[0];

      const bundle = createPreKeyBundle(42, id, spk, otpk);

      expect(bundle.registrationId).toBe(42);
      expect(bundle.identitySigningKey).toEqual(id.signingPublicKey);
      expect(bundle.identityDHKey).toEqual(id.dhPublicKey);
      expect(bundle.signedPreKey.keyId).toBe(5);
      expect(bundle.signedPreKey.publicKey).toEqual(spk.keyPair.publicKey);
      expect(bundle.signedPreKey.signature).toEqual(spk.signature);
      expect(bundle.oneTimePreKey?.keyId).toBe(200);
      expect(bundle.oneTimePreKey?.publicKey).toEqual(otpk.keyPair.publicKey);
    });

    test('assembles bundle without one-time prekey', async () => {
      const id = await generateIdentityKeyPair(crypto);
      const spk = await generateSignedPreKey(crypto, id, 1);

      const bundle = createPreKeyBundle(42, id, spk);

      expect(bundle.oneTimePreKey).toBeUndefined();
    });
  });
});