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

import type { CryptoProvider } from './crypto.js';
import type { StorageProvider } from './storage.js';
import type {
  IdentityKeyPair,
  KeyPair,
  SignedPreKey,
  OneTimePreKey,
  PreKeyBundle,
  PreKeyMessage,
  RatchetMessage,
  SessionState,
} from './types.js';
import { deriveInitialRootKey } from './keys.js';
import { InvalidSignatureError, PreKeyNotFoundError, UntrustedIdentityError } from './errors.js';

/**
 * X3DH — Extended Triple Diffie-Hellman key agreement.
 *
 * Establishes a shared secret between two parties (Alice and Bob) even when
 * Bob is offline, using prekey bundles published to a server.
 *
 * Reference: https://signal.org/docs/specifications/x3dh/
 */

// ─── Key Generation ──────────────────────────────────────────

/** Generate a new identity keypair: Ed25519 (signing) + X25519 (DH) */
export async function generateIdentityKeyPair(crypto: CryptoProvider): Promise<IdentityKeyPair> {
  const [signing, dh] = await Promise.all([
    crypto.generateEd25519KeyPair(),
    crypto.generateX25519KeyPair(),
  ]);
  return {
    signingPublicKey: signing.publicKey,
    signingPrivateKey: signing.privateKey,
    dhPublicKey: dh.publicKey,
    dhPrivateKey: dh.privateKey,
  };
}

/** Generate a signed prekey: X25519 keypair signed by the identity key */
export async function generateSignedPreKey(
  crypto: CryptoProvider,
  identityKey: IdentityKeyPair,
  keyId: number,
): Promise<SignedPreKey> {
  const keyPair = await crypto.generateX25519KeyPair();
  const signature = await crypto.sign(identityKey.signingPrivateKey, keyPair.publicKey);
  return {
    keyId,
    keyPair,
    signature,
    timestamp: Date.now(),
  };
}

/** Generate a batch of one-time prekeys */
export async function generateOneTimePreKeys(
  crypto: CryptoProvider,
  startId: number,
  count: number,
): Promise<OneTimePreKey[]> {
  const keys: OneTimePreKey[] = [];
  for (let i = 0; i < count; i++) {
    const keyPair = await crypto.generateX25519KeyPair();
    keys.push({ keyId: startId + i, keyPair });
  }
  return keys;
}

/** Assemble a prekey bundle for publishing to the prekey server */
export function createPreKeyBundle(
  registrationId: number,
  identityKey: IdentityKeyPair,
  signedPreKey: SignedPreKey,
  oneTimePreKey?: OneTimePreKey,
): PreKeyBundle {
  return {
    registrationId,
    identitySigningKey: identityKey.signingPublicKey,
    identityDHKey: identityKey.dhPublicKey,
    signedPreKey: {
      keyId: signedPreKey.keyId,
      publicKey: signedPreKey.keyPair.publicKey,
      signature: signedPreKey.signature,
    },
    oneTimePreKey: oneTimePreKey
      ? { keyId: oneTimePreKey.keyId, publicKey: oneTimePreKey.keyPair.publicKey }
      : undefined,
  };
}

// ─── Alice: Initiate Session ─────────────────────────────────

/**
 * Process a prekey bundle to establish a new session (Alice's side).
 *
 * Steps:
 * 1. Verify the signed prekey's signature using Bob's identity signing key
 * 2. Generate an ephemeral X25519 keypair
 * 3. Compute 3 or 4 DH shared secrets:
 *    DH1 = DH(Alice identity DH, Bob signed prekey)
 *    DH2 = DH(Alice ephemeral, Bob identity DH)
 *    DH3 = DH(Alice ephemeral, Bob signed prekey)
 *    DH4 = DH(Alice ephemeral, Bob one-time prekey) — if available
 * 4. Derive the initial root key from the concatenated DH outputs
 * 5. Return the X3DH result with all info needed to create a session + PreKeyMessage
 */
export interface X3DHInitResult {
  /** Initial root key (32 bytes) — seeds the Double Ratchet */
  rootKey: Uint8Array;
  /** Alice's ephemeral X25519 public key (sent to Bob in the PreKeyMessage) */
  ephemeralPublicKey: Uint8Array;
  /** Bob's signed prekey ID (so Bob can look up the key) */
  signedPreKeyId: number;
  /** Bob's one-time prekey ID if consumed */
  preKeyId?: number;
  /** Bob's identity DH public key */
  remoteIdentityKey: Uint8Array;
  /** Bob's signed prekey public key (used as initial DH ratchet remote key) */
  remoteSignedPreKey: Uint8Array;
}

export async function processPreKeyBundle(
  crypto: CryptoProvider,
  storage: StorageProvider,
  bundle: PreKeyBundle,
): Promise<X3DHInitResult> {
  // 1. Verify signed prekey signature
  const valid = await crypto.verify(
    bundle.identitySigningKey,
    bundle.signedPreKey.publicKey,
    bundle.signedPreKey.signature,
  );
  if (!valid) {
    throw new InvalidSignatureError('Signed prekey signature is invalid');
  }

  // 2. Check trust (TOFU or pinned)
  // We trust based on the DH key since that's what we use in the protocol
  const identityKey = await storage.getIdentityKeyPair();
  if (!identityKey) throw new Error('No local identity key — call initialize() first');

  // 3. Generate ephemeral keypair
  const ephemeral = await crypto.generateX25519KeyPair();

  // 4. Compute DH shared secrets
  const dh1 = await crypto.x25519(identityKey.dhPrivateKey, bundle.signedPreKey.publicKey);
  const dh2 = await crypto.x25519(ephemeral.privateKey, bundle.identityDHKey);
  const dh3 = await crypto.x25519(ephemeral.privateKey, bundle.signedPreKey.publicKey);

  const secrets = [dh1, dh2, dh3];

  let preKeyId: number | undefined;
  if (bundle.oneTimePreKey) {
    const dh4 = await crypto.x25519(ephemeral.privateKey, bundle.oneTimePreKey.publicKey);
    secrets.push(dh4);
    preKeyId = bundle.oneTimePreKey.keyId;
  }

  // 5. Derive initial root key
  const rootKey = await deriveInitialRootKey(crypto, secrets);

  // 6. Save trust for remote identity
  await storage.saveTrustedIdentity('pending', bundle.identityDHKey);

  return {
    rootKey,
    ephemeralPublicKey: ephemeral.publicKey,
    signedPreKeyId: bundle.signedPreKey.keyId,
    preKeyId,
    remoteIdentityKey: bundle.identityDHKey,
    remoteSignedPreKey: bundle.signedPreKey.publicKey,
  };
}

// ─── Bob: Respond to PreKeyMessage ───────────────────────────

/**
 * Process an incoming PreKeyMessage to establish a session (Bob's side).
 *
 * Steps:
 * 1. Look up the signed prekey and optionally the one-time prekey
 * 2. Compute the same 3 or 4 DH shared secrets (from Bob's perspective)
 * 3. Derive the same initial root key
 * 4. Delete the consumed one-time prekey
 * 5. Return the X3DH result to seed the Double Ratchet
 */
export interface X3DHResponseResult {
  /** Initial root key (32 bytes) — must match Alice's */
  rootKey: Uint8Array;
  /** Alice's identity DH public key */
  remoteIdentityKey: Uint8Array;
  /** Alice's ephemeral public key (used as initial DH ratchet remote key) */
  remoteEphemeralKey: Uint8Array;
  /** The embedded first ratchet message to decrypt */
  initialMessage: RatchetMessage;
}

export async function processPreKeyMessage(
  crypto: CryptoProvider,
  storage: StorageProvider,
  message: PreKeyMessage,
): Promise<X3DHResponseResult> {
  const identityKey = await storage.getIdentityKeyPair();
  if (!identityKey) throw new Error('No local identity key — call initialize() first');

  // 1. Look up signed prekey
  const signedPreKey = await storage.getSignedPreKey(message.signedPreKeyId);
  if (!signedPreKey) {
    throw new PreKeyNotFoundError(message.signedPreKeyId, 'signed');
  }

  // 2. Compute DH shared secrets (Bob's perspective — mirrored from Alice)
  const dh1 = await crypto.x25519(signedPreKey.keyPair.privateKey, message.identityDHKey);
  const dh2 = await crypto.x25519(identityKey.dhPrivateKey, message.ephemeralKey);
  const dh3 = await crypto.x25519(signedPreKey.keyPair.privateKey, message.ephemeralKey);

  const secrets = [dh1, dh2, dh3];

  // 3. If a one-time prekey was used, include DH4
  if (message.preKeyId != null) {
    const oneTimePreKey = await storage.getOneTimePreKey(message.preKeyId);
    if (!oneTimePreKey) {
      throw new PreKeyNotFoundError(message.preKeyId, 'one-time');
    }
    const dh4 = await crypto.x25519(oneTimePreKey.keyPair.privateKey, message.ephemeralKey);
    secrets.push(dh4);

    // 4. Consume (delete) the one-time prekey
    await storage.removeOneTimePreKey(message.preKeyId);
  }

  // 5. Derive the initial root key (should match Alice's)
  const rootKey = await deriveInitialRootKey(crypto, secrets);

  // 6. Save trust for remote identity
  await storage.saveTrustedIdentity('pending', message.identityDHKey);

  return {
    rootKey,
    remoteIdentityKey: message.identityDHKey,
    remoteEphemeralKey: message.ephemeralKey,
    initialMessage: message.message,
  };
}