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

import { describe, test, expect } from 'bun:test';
import { SubtleCryptoProvider, MemoryStorage } from '@shade/crypto-web';
import {
  generateIdentityKeyPair,
  generateSignedPreKey,
  generateOneTimePreKeys,
  createPreKeyBundle,
  processPreKeyBundle,
  processPreKeyMessage,
  initSenderSession,
  initReceiverSession,
  ratchetEncrypt,
  ratchetDecrypt,
} from '../src/index.js';

const crypto = new SubtleCryptoProvider();
const enc = new TextEncoder();
const dec = new TextDecoder();

describe('Full E2EE Integration: X3DH → Double Ratchet', () => {
  test('complete conversation between Alice and Bob', async () => {
    // ─── Setup Bob (publishes prekey bundle) ─────────────────
    const bobStorage = new MemoryStorage();
    const bobIdentity = await generateIdentityKeyPair(crypto);
    await bobStorage.saveIdentityKeyPair(bobIdentity);
    await bobStorage.saveLocalRegistrationId(42);

    const bobSignedPreKey = await generateSignedPreKey(crypto, bobIdentity, 1);
    await bobStorage.saveSignedPreKey(bobSignedPreKey);

    const bobOTPKs = await generateOneTimePreKeys(crypto, 100, 10);
    for (const otpk of bobOTPKs) await bobStorage.saveOneTimePreKey(otpk);

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

    // ─── Alice initiates (processes bundle, creates session) ──
    const aliceStorage = new MemoryStorage();
    const aliceIdentity = await generateIdentityKeyPair(crypto);
    await aliceStorage.saveIdentityKeyPair(aliceIdentity);

    const x3dhResult = await processPreKeyBundle(crypto, aliceStorage, bundle);

    // Alice initializes her ratchet session
    const aliceSession = await initSenderSession(
      crypto,
      x3dhResult.rootKey,
      x3dhResult.remoteIdentityKey,
      x3dhResult.remoteSignedPreKey, // Bob's signed prekey = initial DH ratchet key
    );

    // Alice encrypts her first message
    const firstMsg = await ratchetEncrypt(crypto, aliceSession, enc.encode('Hello Bob! This is E2EE.'));

    // Alice sends a PreKeyMessage to Bob
    const preKeyMessage = {
      registrationId: 1,
      preKeyId: x3dhResult.preKeyId,
      signedPreKeyId: x3dhResult.signedPreKeyId,
      ephemeralKey: x3dhResult.ephemeralPublicKey,
      identityDHKey: aliceIdentity.dhPublicKey,
      message: firstMsg,
    };

    // ─── Bob receives and processes ──────────────────────────
    const bobX3dh = await processPreKeyMessage(crypto, bobStorage, preKeyMessage);
    expect(bobX3dh.rootKey).toEqual(x3dhResult.rootKey);

    // Bob initializes his ratchet session
    const bobSession = initReceiverSession(
      bobX3dh.rootKey,
      bobX3dh.remoteIdentityKey,
      bobSignedPreKey.keyPair, // Bob's signed prekey as his initial DH keypair
    );

    // Bob decrypts Alice's first message
    const plaintext1 = await ratchetDecrypt(crypto, bobSession, bobX3dh.initialMessage);
    expect(dec.decode(plaintext1)).toBe('Hello Bob! This is E2EE.');

    // ─── Full conversation ───────────────────────────────────

    // Alice sends more
    const m2 = await ratchetEncrypt(crypto, aliceSession, enc.encode('Are you there?'));
    expect(dec.decode(await ratchetDecrypt(crypto, bobSession, m2))).toBe('Are you there?');

    // Bob replies (DH ratchet triggers)
    const m3 = await ratchetEncrypt(crypto, bobSession, enc.encode('Yes! Forward secrecy is active.'));
    expect(dec.decode(await ratchetDecrypt(crypto, aliceSession, m3))).toBe('Yes! Forward secrecy is active.');

    // Alice replies
    const m4 = await ratchetEncrypt(crypto, aliceSession, enc.encode('Every message has a unique key.'));
    expect(dec.decode(await ratchetDecrypt(crypto, bobSession, m4))).toBe('Every message has a unique key.');

    // Multiple back-and-forth
    for (let i = 0; i < 10; i++) {
      const sender = i % 2 === 0 ? aliceSession : bobSession;
      const receiver = i % 2 === 0 ? bobSession : aliceSession;
      const text = `Turn ${i}: ${i % 2 === 0 ? 'Alice' : 'Bob'} speaking`;

      const msg = await ratchetEncrypt(crypto, sender, enc.encode(text));
      expect(dec.decode(await ratchetDecrypt(crypto, receiver, msg))).toBe(text);
    }
  });

  test('works without one-time prekey', async () => {
    const bobStorage = new MemoryStorage();
    const bobIdentity = await generateIdentityKeyPair(crypto);
    await bobStorage.saveIdentityKeyPair(bobIdentity);
    const bobSignedPreKey = await generateSignedPreKey(crypto, bobIdentity, 1);
    await bobStorage.saveSignedPreKey(bobSignedPreKey);

    // No one-time prekeys
    const bundle = createPreKeyBundle(42, bobIdentity, bobSignedPreKey);

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

    const x3dhResult = await processPreKeyBundle(crypto, aliceStorage, bundle);
    const aliceSession = await initSenderSession(
      crypto, x3dhResult.rootKey, x3dhResult.remoteIdentityKey, x3dhResult.remoteSignedPreKey,
    );

    const firstMsg = await ratchetEncrypt(crypto, aliceSession, enc.encode('No OTPK needed'));

    const preKeyMessage = {
      registrationId: 1,
      signedPreKeyId: x3dhResult.signedPreKeyId,
      ephemeralKey: x3dhResult.ephemeralPublicKey,
      identityDHKey: aliceIdentity.dhPublicKey,
      message: firstMsg,
    };

    const bobX3dh = await processPreKeyMessage(crypto, bobStorage, preKeyMessage);
    const bobSession = initReceiverSession(
      bobX3dh.rootKey, bobX3dh.remoteIdentityKey, bobSignedPreKey.keyPair,
    );

    expect(dec.decode(await ratchetDecrypt(crypto, bobSession, bobX3dh.initialMessage)))
      .toBe('No OTPK needed');

    // Continue conversation
    const reply = await ratchetEncrypt(crypto, bobSession, enc.encode('Got it!'));
    expect(dec.decode(await ratchetDecrypt(crypto, aliceSession, reply))).toBe('Got it!');
  });

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

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

    // Alice uses OTPK 100
    const bundle = createPreKeyBundle(42, bobIdentity, bobSignedPreKey, bobOTPKs[0]);
    const aliceStorage = new MemoryStorage();
    const aliceIdentity = await generateIdentityKeyPair(crypto);
    await aliceStorage.saveIdentityKeyPair(aliceIdentity);

    const x3dhResult = await processPreKeyBundle(crypto, aliceStorage, bundle);
    const aliceSession = await initSenderSession(
      crypto, x3dhResult.rootKey, x3dhResult.remoteIdentityKey, x3dhResult.remoteSignedPreKey,
    );
    const firstMsg = await ratchetEncrypt(crypto, aliceSession, enc.encode('test'));

    await processPreKeyMessage(crypto, bobStorage, {
      registrationId: 1,
      preKeyId: 100,
      signedPreKeyId: 1,
      ephemeralKey: x3dhResult.ephemeralPublicKey,
      identityDHKey: aliceIdentity.dhPublicKey,
      message: firstMsg,
    });

    // OTPK 100 consumed, 101 and 102 remain
    expect(await bobStorage.getOneTimePreKeyCount()).toBe(2);
    expect(await bobStorage.getOneTimePreKey(100)).toBeNull();
    expect(await bobStorage.getOneTimePreKey(101)).not.toBeNull();
  });
});