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

import { describe, test, expect } from 'bun:test';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import {
  initSenderSession,
  initReceiverSession,
  ratchetEncrypt,
  ratchetDecrypt,
  MaxSkipExceededError,
  DecryptionError,
} from '../src/index.js';
import type { SessionState, RatchetMessage } from '../src/index.js';

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

/** Helper: set up Alice (sender) and Bob (receiver) sessions from a shared root key */
async function setupPair(): Promise<{ alice: SessionState; bob: SessionState }> {
  const rootKey = crypto.randomBytes(32);
  const remoteIdentityKey = crypto.randomBytes(32);

  // Bob's initial DH keypair (would be his signed prekey in real X3DH)
  const bobDH = await crypto.generateX25519KeyPair();

  const alice = await initSenderSession(crypto, rootKey, remoteIdentityKey, bobDH.publicKey);
  const bob = initReceiverSession(rootKey, remoteIdentityKey, bobDH);

  return { alice, bob };
}

describe('Double Ratchet', () => {
  // ─── Basic Send/Receive ──────────────────────────────────

  describe('basic send/receive', () => {
    test('Alice encrypts, Bob decrypts', async () => {
      const { alice, bob } = await setupPair();

      const msg = await ratchetEncrypt(crypto, alice, enc.encode('hello bob'));
      const plaintext = await ratchetDecrypt(crypto, bob, msg);

      expect(dec.decode(plaintext)).toBe('hello bob');
    });

    test('multiple messages in same direction', async () => {
      const { alice, bob } = await setupPair();

      const messages = ['first', 'second', 'third'];
      const encrypted: RatchetMessage[] = [];

      for (const text of messages) {
        encrypted.push(await ratchetEncrypt(crypto, alice, enc.encode(text)));
      }

      for (let i = 0; i < messages.length; i++) {
        const plaintext = await ratchetDecrypt(crypto, bob, encrypted[i]);
        expect(dec.decode(plaintext)).toBe(messages[i]);
      }
    });

    test('counter increments with each message', async () => {
      const { alice } = await setupPair();

      const m0 = await ratchetEncrypt(crypto, alice, enc.encode('a'));
      const m1 = await ratchetEncrypt(crypto, alice, enc.encode('b'));
      const m2 = await ratchetEncrypt(crypto, alice, enc.encode('c'));

      expect(m0.counter).toBe(0);
      expect(m1.counter).toBe(1);
      expect(m2.counter).toBe(2);

      // All use the same DH key (no ratchet step yet)
      expect(m0.dhPublicKey).toEqual(m1.dhPublicKey);
      expect(m1.dhPublicKey).toEqual(m2.dhPublicKey);
    });

    test('each message has a unique nonce', async () => {
      const { alice } = await setupPair();

      const m0 = await ratchetEncrypt(crypto, alice, enc.encode('a'));
      const m1 = await ratchetEncrypt(crypto, alice, enc.encode('a'));

      expect(m0.nonce).not.toEqual(m1.nonce);
      expect(m0.ciphertext).not.toEqual(m1.ciphertext);
    });
  });

  // ─── Ping-Pong (DH Ratchet) ──────────────────────────────

  describe('ping-pong conversation', () => {
    test('alternating messages trigger DH ratchets', async () => {
      const { alice, bob } = await setupPair();

      // Alice → Bob
      const m1 = await ratchetEncrypt(crypto, alice, enc.encode('hi bob'));
      expect(dec.decode(await ratchetDecrypt(crypto, bob, m1))).toBe('hi bob');

      // Bob → Alice (new DH key)
      const m2 = await ratchetEncrypt(crypto, bob, enc.encode('hi alice'));
      expect(m2.dhPublicKey).not.toEqual(m1.dhPublicKey); // DH ratchet happened
      expect(dec.decode(await ratchetDecrypt(crypto, alice, m2))).toBe('hi alice');

      // Alice → Bob (another new DH key)
      const m3 = await ratchetEncrypt(crypto, alice, enc.encode('how are you'));
      expect(m3.dhPublicKey).not.toEqual(m1.dhPublicKey);
      expect(m3.dhPublicKey).not.toEqual(m2.dhPublicKey);
      expect(dec.decode(await ratchetDecrypt(crypto, bob, m3))).toBe('how are you');

      // Bob → Alice
      const m4 = await ratchetEncrypt(crypto, bob, enc.encode('great!'));
      expect(dec.decode(await ratchetDecrypt(crypto, alice, m4))).toBe('great!');
    });

    test('extended conversation with many turns', async () => {
      const { alice, bob } = await setupPair();

      for (let i = 0; i < 20; i++) {
        const sender = i % 2 === 0 ? alice : bob;
        const receiver = i % 2 === 0 ? bob : alice;
        const text = `message ${i}`;

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

    test('burst messages then reply', async () => {
      const { alice, bob } = await setupPair();

      // Alice sends 5 messages
      const burst: RatchetMessage[] = [];
      for (let i = 0; i < 5; i++) {
        burst.push(await ratchetEncrypt(crypto, alice, enc.encode(`alice-${i}`)));
      }

      // Bob receives all 5
      for (let i = 0; i < 5; i++) {
        expect(dec.decode(await ratchetDecrypt(crypto, bob, burst[i]))).toBe(`alice-${i}`);
      }

      // Bob replies (triggers DH ratchet)
      const reply = await ratchetEncrypt(crypto, bob, enc.encode('got them all'));
      expect(dec.decode(await ratchetDecrypt(crypto, alice, reply))).toBe('got them all');

      // Alice sends more
      const m = await ratchetEncrypt(crypto, alice, enc.encode('great!'));
      expect(dec.decode(await ratchetDecrypt(crypto, bob, m))).toBe('great!');
    });
  });

  // ─── Out-of-Order Messages ────────────────────────────────

  describe('out-of-order delivery', () => {
    test('messages received in reverse order', async () => {
      const { alice, bob } = await setupPair();

      const m0 = await ratchetEncrypt(crypto, alice, enc.encode('first'));
      const m1 = await ratchetEncrypt(crypto, alice, enc.encode('second'));
      const m2 = await ratchetEncrypt(crypto, alice, enc.encode('third'));

      // Deliver in reverse
      expect(dec.decode(await ratchetDecrypt(crypto, bob, m2))).toBe('third');
      expect(dec.decode(await ratchetDecrypt(crypto, bob, m0))).toBe('first');
      expect(dec.decode(await ratchetDecrypt(crypto, bob, m1))).toBe('second');
    });

    test('skip some messages, then receive them later', async () => {
      const { alice, bob } = await setupPair();

      const messages: RatchetMessage[] = [];
      for (let i = 0; i < 10; i++) {
        messages.push(await ratchetEncrypt(crypto, alice, enc.encode(`msg-${i}`)));
      }

      // Receive only even-numbered messages first
      for (let i = 0; i < 10; i += 2) {
        expect(dec.decode(await ratchetDecrypt(crypto, bob, messages[i]))).toBe(`msg-${i}`);
      }

      // Then receive odd-numbered (skipped) messages
      for (let i = 1; i < 10; i += 2) {
        expect(dec.decode(await ratchetDecrypt(crypto, bob, messages[i]))).toBe(`msg-${i}`);
      }
    });

    test('out-of-order across DH ratchet boundaries', async () => {
      const { alice, bob } = await setupPair();

      // Alice sends 3 messages
      const a0 = await ratchetEncrypt(crypto, alice, enc.encode('a0'));
      const a1 = await ratchetEncrypt(crypto, alice, enc.encode('a1'));
      const a2 = await ratchetEncrypt(crypto, alice, enc.encode('a2'));

      // Bob receives only a2 (skips a0, a1)
      expect(dec.decode(await ratchetDecrypt(crypto, bob, a2))).toBe('a2');

      // Bob replies (DH ratchet)
      const b0 = await ratchetEncrypt(crypto, bob, enc.encode('b0'));
      expect(dec.decode(await ratchetDecrypt(crypto, alice, b0))).toBe('b0');

      // Now Bob receives the skipped a0 and a1 (from the old chain)
      expect(dec.decode(await ratchetDecrypt(crypto, bob, a0))).toBe('a0');
      expect(dec.decode(await ratchetDecrypt(crypto, bob, a1))).toBe('a1');
    });
  });

  // ─── Error Cases ──────────────────────────────────────────

  describe('error cases', () => {
    test('max skip exceeded throws', async () => {
      const { alice, bob } = await setupPair();

      // Encrypt 1002 messages but only try to decrypt the last one
      let lastMsg: RatchetMessage | undefined;
      for (let i = 0; i < 1002; i++) {
        lastMsg = await ratchetEncrypt(crypto, alice, enc.encode(`msg-${i}`));
      }

      expect(ratchetDecrypt(crypto, bob, lastMsg!)).rejects.toThrow(MaxSkipExceededError);
    });

    test('tampered ciphertext fails', async () => {
      const { alice, bob } = await setupPair();

      const msg = await ratchetEncrypt(crypto, alice, enc.encode('secret'));
      msg.ciphertext[0] ^= 0xff;

      expect(ratchetDecrypt(crypto, bob, msg)).rejects.toThrow(DecryptionError);
    });

    test('tampered header (counter) fails due to AAD', async () => {
      const { alice, bob } = await setupPair();

      const msg = await ratchetEncrypt(crypto, alice, enc.encode('secret'));
      msg.counter = 999; // tamper with counter

      expect(ratchetDecrypt(crypto, bob, msg)).rejects.toThrow();
    });
  });

  // ─── Long Conversation ────────────────────────────────────

  describe('stress test', () => {
    test('100+ message conversation with alternating turns', async () => {
      const { alice, bob } = await setupPair();

      for (let i = 0; i < 50; i++) {
        // Alice sends 2 messages
        for (let j = 0; j < 2; j++) {
          const text = `alice-${i}-${j}`;
          const msg = await ratchetEncrypt(crypto, alice, enc.encode(text));
          expect(dec.decode(await ratchetDecrypt(crypto, bob, msg))).toBe(text);
        }

        // Bob sends 1 message
        const text = `bob-${i}`;
        const msg = await ratchetEncrypt(crypto, bob, enc.encode(text));
        expect(dec.decode(await ratchetDecrypt(crypto, alice, msg))).toBe(text);
      }
    });
  });
});