Shade/packages/shade-core/src/ratchet.ts

import type { CryptoProvider } from './crypto.js';
import type { KeyPair, SessionState, ChainState, RatchetMessage } from './types.js';
import { MAX_SKIP, MAX_CACHED_SKIPPED_KEYS } from './types.js';
import { kdfRootKey, kdfChainKey } from './keys.js';
import { DecryptionError, MaxSkipExceededError } from './errors.js';

/**
 * Double Ratchet — per-message forward secrecy and post-compromise recovery.
 *
 * Combines a symmetric-key ratchet (chain keys → message keys) with a
 * Diffie-Hellman ratchet (new DH keypair per conversation turn).
 *
 * Reference: https://signal.org/docs/specifications/doubleratchet/
 */

// ─── Utility ─────────────────────────────────────────────────

function toBase64(buf: Uint8Array): string {
  // Use a simple hex encoding for map keys (avoids btoa issues with Uint8Array)
  return Array.from(buf, (b) => b.toString(16).padStart(2, '0')).join('');
}

function skippedKeyId(dhPublicKey: Uint8Array, counter: number): string {
  return `${toBase64(dhPublicKey)}:${counter}`;
}

/** Encode a RatchetMessage header as bytes for use as AES-GCM AAD */
function encodeHeader(msg: Pick<RatchetMessage, 'dhPublicKey' | 'previousCounter' | 'counter'>): Uint8Array {
  // dhPublicKey (32) + previousCounter (4, big-endian) + counter (4, big-endian)
  const buf = new Uint8Array(40);
  buf.set(msg.dhPublicKey, 0);
  new DataView(buf.buffer).setUint32(32, msg.previousCounter, false);
  new DataView(buf.buffer).setUint32(36, msg.counter, false);
  return buf;
}

// ─── Session Initialization ──────────────────────────────────

/**
 * Initialize a session as the sender (Alice, after X3DH).
 *
 * Alice knows the root key and Bob's signed prekey (initial remote DH key).
 * She generates her first DH ratchet keypair and performs the first DH ratchet step.
 */
export async function initSenderSession(
  crypto: CryptoProvider,
  rootKey: Uint8Array,
  remoteIdentityKey: Uint8Array,
  remoteDHPublicKey: Uint8Array,
): Promise<SessionState> {
  // Generate first DH ratchet keypair
  const dhSend = await crypto.generateX25519KeyPair();

  // First DH ratchet step
  const dhOutput = await crypto.x25519(dhSend.privateKey, remoteDHPublicKey);
  const { newRootKey, chainKey } = await kdfRootKey(crypto, rootKey, dhOutput);

  return {
    remoteIdentityKey,
    rootKey: newRootKey,
    sendChain: { chainKey, counter: 0 },
    receiveChain: null,
    dhSend,
    dhReceive: remoteDHPublicKey,
    previousSendCounter: 0,
    skippedKeys: new Map(),
  };
}

/**
 * Initialize a session as the receiver (Bob, after X3DH).
 *
 * Bob knows the root key and his own signed prekey (which was used as
 * the initial DH ratchet keypair). The keypair is COPIED into the
 * session — the receiving side's DH ratchet will eventually rotate
 * `dhSend` and zeroize the previous private key, and that scratch
 * buffer must NOT be the same memory as the persisted signed prekey
 * (which is shared with future X3DH establishments from other senders).
 */
export function initReceiverSession(
  rootKey: Uint8Array,
  remoteIdentityKey: Uint8Array,
  localDHKeyPair: KeyPair,
): SessionState {
  return {
    remoteIdentityKey,
    rootKey,
    sendChain: { chainKey: new Uint8Array(32), counter: 0 },
    receiveChain: null,
    dhSend: {
      publicKey: new Uint8Array(localDHKeyPair.publicKey),
      privateKey: new Uint8Array(localDHKeyPair.privateKey),
    },
    dhReceive: null,
    previousSendCounter: 0,
    skippedKeys: new Map(),
  };
}

// ─── Encrypt ─────────────────────────────────────────────────

/**
 * Encrypt a plaintext message using the Double Ratchet.
 *
 * Advances the sending chain by one step, derives a message key,
 * encrypts the plaintext with AES-256-GCM, and returns a RatchetMessage.
 */
export async function ratchetEncrypt(
  crypto: CryptoProvider,
  session: SessionState,
  plaintext: Uint8Array,
): Promise<RatchetMessage> {
  // Advance sending chain — old chain key is replaced, zero the old one
  const oldChainKey = session.sendChain.chainKey;
  const { newChainKey, messageKey } = await kdfChainKey(crypto, oldChainKey);
  crypto.zeroize(oldChainKey);
  const counter = session.sendChain.counter;

  // Build header for AAD
  const header: Pick<RatchetMessage, 'dhPublicKey' | 'previousCounter' | 'counter'> = {
    dhPublicKey: session.dhSend.publicKey,
    previousCounter: session.previousSendCounter,
    counter,
  };
  const aad = encodeHeader(header);

  // Encrypt, then zero the message key (used once)
  const { ciphertext, nonce } = await crypto.aesGcmEncrypt(messageKey, plaintext, aad);
  crypto.zeroize(messageKey);

  // Update session state
  session.sendChain.chainKey = newChainKey;
  session.sendChain.counter = counter + 1;

  return {
    dhPublicKey: session.dhSend.publicKey,
    previousCounter: session.previousSendCounter,
    counter,
    ciphertext,
    nonce,
  };
}

// ─── Decrypt ─────────────────────────────────────────────────

/**
 * Decrypt a RatchetMessage using the Double Ratchet.
 *
 * Handles three cases:
 * 1. Message from a skipped key (out-of-order delivery)
 * 2. Message from the current receiving chain
 * 3. Message with a new DH key (triggers a DH ratchet step)
 */
export async function ratchetDecrypt(
  crypto: CryptoProvider,
  session: SessionState,
  message: RatchetMessage,
): Promise<Uint8Array> {
  // Case 1: Try skipped keys first
  const skipId = skippedKeyId(message.dhPublicKey, message.counter);
  const skippedKey = session.skippedKeys.get(skipId);
  if (skippedKey) {
    session.skippedKeys.delete(skipId);
    try {
      return await decryptWithKey(crypto, skippedKey, message);
    } finally {
      crypto.zeroize(skippedKey);
    }
  }

  // Case 2 or 3: Check if this is a new DH ratchet
  const isNewRatchet = !session.dhReceive || !arraysEqual(message.dhPublicKey, session.dhReceive);

  if (isNewRatchet) {
    // Skip any remaining messages in the current receiving chain
    if (session.receiveChain && session.dhReceive) {
      await skipMessageKeys(crypto, session, session.dhReceive, session.receiveChain, message.previousCounter);
    }

    // Perform DH ratchet step
    await performDHRatchetStep(crypto, session, message.dhPublicKey);
  }

  // Skip to the message's counter in the current receiving chain
  if (!session.receiveChain) {
    throw new DecryptionError('No receiving chain available');
  }
  await skipMessageKeys(crypto, session, message.dhPublicKey, session.receiveChain, message.counter);

  // Advance the receiving chain one more step to get this message's key
  const oldChainKey = session.receiveChain.chainKey;
  const { newChainKey, messageKey } = await kdfChainKey(crypto, oldChainKey);
  crypto.zeroize(oldChainKey);
  session.receiveChain.chainKey = newChainKey;
  session.receiveChain.counter = message.counter + 1;

  try {
    return await decryptWithKey(crypto, messageKey, message);
  } finally {
    crypto.zeroize(messageKey);
  }
}

// ─── DH Ratchet Step ─────────────────────────────────────────

/**
 * Perform a DH ratchet step when receiving a message with a new DH public key.
 *
 * 1. DH(current send key, new remote key) → advance root key, get new receiving chain
 * 2. Generate new DH keypair
 * 3. DH(new keypair, remote key) → advance root key, get new sending chain
 */
async function performDHRatchetStep(
  crypto: CryptoProvider,
  session: SessionState,
  remoteDHKey: Uint8Array,
): Promise<void> {
  // Save previous send counter
  session.previousSendCounter = session.sendChain.counter;

  // Update remote DH key
  session.dhReceive = remoteDHKey;

  // DH with current send key → new receiving chain
  const dh1 = await crypto.x25519(session.dhSend.privateKey, remoteDHKey);
  const oldRootKey1 = session.rootKey;
  const recv = await kdfRootKey(crypto, oldRootKey1, dh1);
  crypto.zeroize(oldRootKey1);
  crypto.zeroize(dh1);
  session.rootKey = recv.newRootKey;
  session.receiveChain = { chainKey: recv.chainKey, counter: 0 };

  // Generate new DH keypair, zero the old one's private key
  const oldDhPrivate = session.dhSend.privateKey;
  session.dhSend = await crypto.generateX25519KeyPair();
  crypto.zeroize(oldDhPrivate);

  // DH with new send key → new sending chain
  const dh2 = await crypto.x25519(session.dhSend.privateKey, remoteDHKey);
  const oldRootKey2 = session.rootKey;
  const send = await kdfRootKey(crypto, oldRootKey2, dh2);
  crypto.zeroize(oldRootKey2);
  crypto.zeroize(dh2);
  session.rootKey = send.newRootKey;
  // Zero the old send chain key if present
  if (session.sendChain.chainKey.length > 0) {
    crypto.zeroize(session.sendChain.chainKey);
  }
  session.sendChain = { chainKey: send.chainKey, counter: 0 };
}

// ─── Skip Message Keys ──────────────────────────────────────

/**
 * Advance a chain, caching skipped message keys for out-of-order decryption.
 */
async function skipMessageKeys(
  crypto: CryptoProvider,
  session: SessionState,
  dhPublicKey: Uint8Array,
  chain: ChainState,
  untilCounter: number,
): Promise<void> {
  const toSkip = untilCounter - chain.counter;
  if (toSkip < 0) return; // already past this point
  if (toSkip > MAX_SKIP) {
    throw new MaxSkipExceededError(toSkip, MAX_SKIP);
  }

  for (let i = chain.counter; i < untilCounter; i++) {
    const { newChainKey, messageKey } = await kdfChainKey(crypto, chain.chainKey);
    const id = skippedKeyId(dhPublicKey, i);
    session.skippedKeys.set(id, messageKey);
    chain.chainKey = newChainKey;
    chain.counter = i + 1;

    // Evict oldest if we have too many cached keys
    if (session.skippedKeys.size > MAX_CACHED_SKIPPED_KEYS) {
      const firstKey = session.skippedKeys.keys().next().value;
      if (firstKey) session.skippedKeys.delete(firstKey);
    }
  }
}

// ─── Helpers ─────────────────────────────────────────────────

async function decryptWithKey(
  crypto: CryptoProvider,
  messageKey: Uint8Array,
  message: RatchetMessage,
): Promise<Uint8Array> {
  const aad = encodeHeader(message);
  try {
    return await crypto.aesGcmDecrypt(messageKey, message.ciphertext, message.nonce, aad);
  } catch {
    throw new DecryptionError('Failed to decrypt message — wrong key or tampered data');
  }
}

function arraysEqual(a: Uint8Array, b: Uint8Array): boolean {
  if (a.length !== b.length) return false;
  for (let i = 0; i < a.length; i++) {
    if (a[i] !== b[i]) return false;
  }
  return true;
}