Shade/packages/shade-core/src/sender-keys.ts

import type { CryptoProvider } from './crypto.js';
import { kdfChainKey } from './keys.js';

/**
 * Signal-style Sender Keys for group messaging.
 *
 * Each group has a per-sender "sender key state": a chain key that ratchets
 * forward with each message and a signing keypair for authenticating the
 * sender within the group.
 *
 * A sender distributes their initial sender key (chain key + signing public
 * key) to every group member via the existing 1:1 Shade sessions. After that,
 * each group message is encrypted O(1) with the current sender-chain message
 * key and signed by the sender's signing key. Each recipient advances their
 * copy of the chain key to decrypt.
 *
 * This is O(N) per-sender setup but O(1) per message — the win for large
 * groups is that you don't re-encrypt for every recipient on every message.
 *
 * Reference: https://signal.org/docs/specifications/sesame/
 */

export interface SenderKeyState {
  /** Current chain key for this sender's messages */
  chainKey: Uint8Array;
  /** Current iteration (counter) */
  iteration: number;
  /** Sender's signing public key (Ed25519) */
  signingPublicKey: Uint8Array;
  /** Sender's signing private key (only present for the sender themselves) */
  signingPrivateKey?: Uint8Array;
}

export interface GroupSession {
  groupId: string;
  /** Map from sender address → their sender key state */
  senderKeys: Map<string, SenderKeyState>;
}

/** A group message to distribute. */
export interface SenderKeyMessage {
  senderAddress: string;
  iteration: number;
  ciphertext: Uint8Array;
  nonce: Uint8Array;
  signature: Uint8Array;
}

/** Initial sender key distribution message (sent via 1:1 Shade session) */
export interface SenderKeyDistribution {
  groupId: string;
  senderAddress: string;
  chainKey: Uint8Array;
  iteration: number;
  signingPublicKey: Uint8Array;
}

/** Create a fresh sender key state for a new sender in a group */
export async function createSenderKey(
  crypto: CryptoProvider,
): Promise<SenderKeyState> {
  const chainKey = crypto.randomBytes(32);
  const { publicKey, privateKey } = await crypto.generateEd25519KeyPair();
  return {
    chainKey,
    iteration: 0,
    signingPublicKey: publicKey,
    signingPrivateKey: privateKey,
  };
}

/** Build a distribution message to send to group members (copies chainKey) */
export function buildDistribution(
  groupId: string,
  senderAddress: string,
  state: SenderKeyState,
): SenderKeyDistribution {
  return {
    groupId,
    senderAddress,
    chainKey: new Uint8Array(state.chainKey), // defensive copy
    iteration: state.iteration,
    signingPublicKey: new Uint8Array(state.signingPublicKey),
  };
}

/** Install a received distribution into a group session (on receiving side) */
export function installDistribution(
  session: GroupSession,
  dist: SenderKeyDistribution,
): void {
  session.senderKeys.set(dist.senderAddress, {
    chainKey: new Uint8Array(dist.chainKey), // defensive copy
    iteration: dist.iteration,
    signingPublicKey: new Uint8Array(dist.signingPublicKey),
    // No signing private key on the receiving side
  });
}

/**
 * Encrypt a plaintext for the group using our sender-key chain.
 * Advances the chain by one step.
 */
export async function senderKeyEncrypt(
  crypto: CryptoProvider,
  session: GroupSession,
  senderAddress: string,
  plaintext: Uint8Array,
): Promise<SenderKeyMessage> {
  const state = session.senderKeys.get(senderAddress);
  if (!state) throw new Error(`No sender key for ${senderAddress} in group ${session.groupId}`);
  if (!state.signingPrivateKey) throw new Error('Cannot send: no signing private key');

  // Advance chain
  const { newChainKey, messageKey } = await kdfChainKey(crypto, state.chainKey);
  crypto.zeroize(state.chainKey);
  const iteration = state.iteration;

  // Encrypt with AAD = groupId || senderAddress || iteration for binding
  const aad = encodeHeader(session.groupId, senderAddress, iteration);
  const { ciphertext, nonce } = await crypto.aesGcmEncrypt(messageKey, plaintext, aad);
  crypto.zeroize(messageKey);

  // Sign the ciphertext + header for authenticity
  const toSign = new Uint8Array(aad.length + ciphertext.length);
  toSign.set(aad, 0);
  toSign.set(ciphertext, aad.length);
  const signature = await crypto.sign(state.signingPrivateKey, toSign);

  // Update state
  state.chainKey = newChainKey;
  state.iteration = iteration + 1;

  return { senderAddress, iteration, ciphertext, nonce, signature };
}

/**
 * Decrypt a sender key message. Verifies signature then advances the
 * appropriate chain key.
 *
 * Handles out-of-order delivery via skipped-key cache (simpler than Double
 * Ratchet since there's no DH ratchet).
 */
export async function senderKeyDecrypt(
  crypto: CryptoProvider,
  session: GroupSession,
  groupId: string,
  message: SenderKeyMessage,
): Promise<Uint8Array> {
  if (groupId !== session.groupId) throw new Error('Group ID mismatch');

  const state = session.senderKeys.get(message.senderAddress);
  if (!state) {
    throw new Error(`Unknown sender ${message.senderAddress} in group ${session.groupId}`);
  }

  // Verify signature
  const aad = encodeHeader(session.groupId, message.senderAddress, message.iteration);
  const signedBytes = new Uint8Array(aad.length + message.ciphertext.length);
  signedBytes.set(aad, 0);
  signedBytes.set(message.ciphertext, aad.length);

  const valid = await crypto.verify(state.signingPublicKey, signedBytes, message.signature);
  if (!valid) throw new Error('Invalid signature on group message');

  // Advance chain to the message's iteration
  if (message.iteration < state.iteration) {
    throw new Error(`Stale iteration ${message.iteration} (current ${state.iteration})`);
  }
  const skip = message.iteration - state.iteration;
  if (skip > 1000) throw new Error('Too many skipped group messages');

  let chainKey = state.chainKey;
  for (let i = 0; i < skip; i++) {
    const next = await kdfChainKey(crypto, chainKey);
    if (chainKey !== state.chainKey) crypto.zeroize(chainKey);
    chainKey = next.newChainKey;
  }

  const { newChainKey, messageKey } = await kdfChainKey(crypto, chainKey);
  if (chainKey !== state.chainKey) crypto.zeroize(chainKey);

  let plaintext: Uint8Array;
  try {
    plaintext = await crypto.aesGcmDecrypt(messageKey, message.ciphertext, message.nonce, aad);
  } finally {
    crypto.zeroize(messageKey);
  }

  // Update state
  if (state.chainKey !== newChainKey) crypto.zeroize(state.chainKey);
  state.chainKey = newChainKey;
  state.iteration = message.iteration + 1;

  return plaintext;
}

function encodeHeader(groupId: string, senderAddress: string, iteration: number): Uint8Array {
  const encoder = new TextEncoder();
  const gBytes = encoder.encode(groupId);
  const sBytes = encoder.encode(senderAddress);
  const buf = new Uint8Array(2 + gBytes.length + 2 + sBytes.length + 4);
  let offset = 0;
  new DataView(buf.buffer).setUint16(offset, gBytes.length, false); offset += 2;
  buf.set(gBytes, offset); offset += gBytes.length;
  new DataView(buf.buffer).setUint16(offset, sBytes.length, false); offset += 2;
  buf.set(sBytes, offset); offset += sBytes.length;
  new DataView(buf.buffer).setUint32(offset, iteration, false);
  return buf;
}