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('initReceiverSession isolation', () => {
    /**
     * Regression — the V3.10 multi-sender recovery flow surfaced a
     * bug where `initReceiverSession` shared a reference to the
     * receiver's signed-prekey keypair with the new session. The
     * first DH ratchet step zeroed the session's stale send-key
     * private bytes — which were the SAME backing buffer as the
     * persisted signed prekey. A subsequent X3DH from a different
     * sender then derived a divergent root key and decryption
     * failed.
     *
     * Fix: `initReceiverSession` copies the keypair into the
     * session. Verify here.
     */
    test('does not mutate the caller-provided keypair after a DH ratchet step', async () => {
      const rootKey = crypto.randomBytes(32);
      const remoteIdentityKey = crypto.randomBytes(32);
      const bobDH = await crypto.generateX25519KeyPair();
      const originalPrivate = new Uint8Array(bobDH.privateKey);
      const originalPublic = new Uint8Array(bobDH.publicKey);

      const bob = initReceiverSession(rootKey, remoteIdentityKey, bobDH);

      // Drive a full receive that triggers `performDHRatchetStep`
      // via a message with a fresh dhPublicKey.
      const aliceFirst = await initSenderSession(crypto, rootKey, remoteIdentityKey, bobDH.publicKey);
      const msg = await ratchetEncrypt(crypto, aliceFirst, enc.encode('hi'));
      await ratchetDecrypt(crypto, bob, msg);

      // The ORIGINAL keypair must not have been touched, so a
      // second X3DH-style establishment using the same prekey
      // material still succeeds.
      expect(Array.from(bobDH.privateKey)).toEqual(Array.from(originalPrivate));
      expect(Array.from(bobDH.publicKey)).toEqual(Array.from(originalPublic));

      // Sanity-check: a second receiver session built from the
      // same keypair should still decrypt fresh sender traffic.
      const bob2 = initReceiverSession(rootKey, remoteIdentityKey, bobDH);
      const aliceSecond = await initSenderSession(crypto, rootKey, remoteIdentityKey, bobDH.publicKey);
      const msg2 = await ratchetEncrypt(crypto, aliceSecond, enc.encode('hi again'));
      const plain2 = await ratchetDecrypt(crypto, bob2, msg2);
      expect(dec.decode(plain2)).toBe('hi again');
    });
  });

  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();
    });

    /**
     * Regression — the v4.2.0 OutboundQueue waiter-since bug delivered
     * the same envelope twice to `manager.decrypt`. The first decrypt
     * succeeded via a cached skipped key; the second one fell into the
     * `message.counter < chain.counter` path with no skipped key
     * available, advanced the chainKey ONCE and rewound `chain.counter`
     * to `message.counter + 1`, leaving the ratchet permanently
     * desynced. ratchetDecrypt now rejects without mutating state when
     * a same-DH message is behind the chain and not in skippedKeys, so
     * a downstream replay (transport bug, retry, etc.) cannot poison
     * the session for everyone else.
     */
    test('same-DH stale message after consumed skipped key fails without corrupting state', async () => {
      const { alice, bob } = await setupPair();

      // Alice sends 3 messages on the same DH chain.
      const m0 = await ratchetEncrypt(crypto, alice, enc.encode('m0'));
      const m1 = await ratchetEncrypt(crypto, alice, enc.encode('m1'));
      const m2 = await ratchetEncrypt(crypto, alice, enc.encode('m2'));

      // Bob receives m1 first, caching m0's key. Then m0 (delivered
      // via the cache). After this, m0's skipped key is consumed.
      expect(dec.decode(await ratchetDecrypt(crypto, bob, m1))).toBe('m1');
      expect(dec.decode(await ratchetDecrypt(crypto, bob, m0))).toBe('m0');

      // Replay of m0: skippedKey is gone, chain.counter is past m0.
      // Pre-fix: this would corrupt Bob's chain state; post-fix it
      // throws cleanly.
      await expect(ratchetDecrypt(crypto, bob, m0)).rejects.toThrow(DecryptionError);

      // Bob can still decrypt the remaining valid message — chain
      // state was NOT mutated by the rejected replay.
      expect(dec.decode(await ratchetDecrypt(crypto, bob, m2))).toBe('m2');
    });
  });

  // ─── 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);
      }
    });
  });
});