End-to-end encryption: Signal protocols and MLS

// A symmetric-key ratchet as section 2.2 and section 7.2 of the Double Ratchet specification describe it: each message
// gets a key from HMAC-SHA-256 of the chain key with the byte 0x01, the chain key moves on with the byte 0x02, and the
// old chain key is overwritten. Three designs are then tested: a thief copies each design's state after message 5 and
// tries to open all eight ciphertexts. The chain key and the messages are lab values.
import { createCipheriv, createDecipheriv, createHash, createHmac, hkdfSync } from 'node:crypto';

const hmac = (key, byte) => createHmac('sha256', key).update(Buffer.from([byte])).digest();
const step = (ck) => ({ mk: hmac(ck, 0x01), ck: hmac(ck, 0x02) });
const fingerprint = (bytes) => createHash('sha256').update(bytes).digest('hex').slice(0, 8);
// All a thief can do with a chain key is run it forward.
function forward(ck, count) {
  const out = [];
  for (let i = 0; i < count; i++) {
    const s = step(ck);
    out.push(s.mk);
    ck = s.ck;
  }
  return out;
}

// One message key encrypts one message, so the nonce can come from the same HKDF call as the AES key.
function seal(mk, plaintext, ad) {
  const okm = Buffer.from(hkdfSync('sha256', mk, Buffer.alloc(32), 'msc ratchet demo v1', 44));
  const cipher = createCipheriv('aes-256-gcm', okm.subarray(0, 32), okm.subarray(32, 44));
  cipher.setAAD(Buffer.from(ad));
  return Buffer.concat([cipher.update(plaintext, 'utf8'), cipher.final(), cipher.getAuthTag()]);
}
function open(mk, blob, ad) {
  const okm = Buffer.from(hkdfSync('sha256', mk, Buffer.alloc(32), 'msc ratchet demo v1', 44));
  const decipher = createDecipheriv('aes-256-gcm', okm.subarray(0, 32), okm.subarray(32, 44));
  decipher.setAAD(Buffer.from(ad));
  decipher.setAuthTag(blob.subarray(-16));
  try {
    return Buffer.concat([decipher.update(blob.subarray(0, -16)), decipher.final()]).toString('utf8');
  } catch {
    return null;
  }
}

const ROOT_CK = Buffer.alloc(32, 0x5a);
const MESSAGES = ['hello', 'are you there', 'meet at six', 'bring the files', 'ok', 'new number is 98xx', 'delete this chat', 'bye'];
const THEFT_AFTER = 5; // the state is copied once five messages are sent

// Alice and Bob walk the same chain: the first five message keys agree, and each side forgets what it has used.
let a = ROOT_CK;
let b = ROOT_CK;
const keys = [];
for (let n = 0; n < 5; n++) {
  const sa = step(a);
  const sb = step(b);
  a = sa.ck;
  b = sb.ck;
  keys.push(`#${n} ${fingerprint(sa.mk)}${fingerprint(sa.mk) === fingerprint(sb.mk) ? '' : ' MISMATCH'}`);
}
console.log(`first five message keys, Alice and Bob agree: ${keys.join('  ')}`);

// The three designs. send(n, text) returns a ciphertext, state() is what a thief copies from the device.
const designs = {
  'one key for every message': () => {
    const key = ROOT_CK;
    return { send: (n, text) => seal(key, text, `msg ${n}`), state: () => ({ key }), derive: (s) => [s.key] };
  },
  'chain, used keys kept': () => {
    let ck = ROOT_CK;
    const used = [];
    return {
      send: (n, text) => {
        const s = step(ck);
        ck = s.ck;
        used.push(s.mk); // kept "in case of a retry": the mistake
        return seal(s.mk, text, `msg ${n}`);
      },
      state: () => ({ ck, used: [...used] }),
      derive: (s) => [...s.used, ...forward(s.ck, 100)],
    };
  },
  'chain, used keys deleted': () => {
    let ck = ROOT_CK;
    return {
      send: (n, text) => {
        const s = step(ck);
        ck = s.ck; // the old chain key and the message key are not kept anywhere
        return seal(s.mk, text, `msg ${n}`);
      },
      state: () => ({ ck }),
      derive: (s) => forward(s.ck, 100),
    };
  },
};

console.log(`\nA thief copies each device after ${THEFT_AFTER} of ${MESSAGES.length} messages and tries every key it can get:`);
console.log('design                       before the theft   after the theft');
for (const [name, make] of Object.entries(designs)) {
  const d = make();
  const sealed = [];
  let stolen;
  MESSAGES.forEach((text, n) => {
    if (n === THEFT_AFTER) stolen = d.state();
    sealed.push(d.send(n, text));
  });
  const candidates = d.derive(stolen);
  const opened = sealed.map((blob, n) => candidates.some((mk) => open(mk, blob, `msg ${n}`) === MESSAGES[n]));
  const before = opened.slice(0, THEFT_AFTER).filter(Boolean).length;
  const after = opened.slice(THEFT_AFTER).filter(Boolean).length;
  console.log(`${name.padEnd(28)} ${`${before} of ${THEFT_AFTER}`.padEnd(18)} ${after} of ${MESSAGES.length - THEFT_AFTER}`);
}

// The AEAD still refuses a changed ciphertext or a ciphertext moved to another message number.
const mk0 = step(ROOT_CK).mk;
const blob = seal(mk0, 'hello', 'msg 0');
const flipped = Buffer.from(blob);
flipped[0] ^= 1;
const show = (text) => text ?? 'refused';
console.log(`\nunchanged: ${show(open(mk0, blob, 'msg 0'))}, one bit changed: ${show(open(mk0, flipped, 'msg 0'))}, same bytes as message 1: ${show(open(mk0, blob, 'msg 1'))}`);

first five message keys, Alice and Bob agree: #0 52808d97  #1 5ee9147c  #2 2caa5f1a  #3 a00f5eec  #4 22b31b27

A thief copies each device after 5 of 8 messages and tries every key it can get:
design                       before the theft   after the theft
one key for every message    5 of 5             3 of 3
chain, used keys kept        5 of 5             3 of 3
chain, used keys deleted     0 of 5             3 of 3

unchanged: hello, one bit changed: refused, same bytes as message 1: refused

"""A forward-secrecy audit for a message-key ratchet: run a design for eight messages, copy its state after message 5,
and look for the key of any earlier message in that copy or in what the copy can derive. The chain key is a lab value.
Press Edit, add a design of your own to DESIGNS, and run it again."""
import hashlib
import hmac

ROOT_CK = bytes([0x5A] * 32)
SENT, THEFT_AFTER, FORWARD_STEPS = 8, 5, 1000


def step(ck: bytes) -> tuple[bytes, bytes]:
    """The Double Ratchet's KDF_CK with HMAC-SHA-256: (next chain key, message key)."""
    return hmac.digest(ck, b"\x02", "sha256"), hmac.digest(ck, b"\x01", "sha256")


def fingerprint(key: bytes) -> str:
    return hashlib.sha256(key).hexdigest()[:8]


class OneKey:
    """Every message under the root key."""
    def __init__(self): self.key = ROOT_CK
    def send(self): return self.key
    def state(self): return [self.key]
    def advance(self, state): return state  # nothing to run forward


class KeepsUsedKeys:
    """A real chain, but every message key is kept for retries."""
    def __init__(self): self.ck, self.used = ROOT_CK, []
    def send(self):
        self.ck, mk = step(self.ck)
        self.used.append(mk)
        return mk
    def state(self): return [self.ck, *self.used]
    def advance(self, state): return state[:1]


class KeepsPreviousChainKey:
    """Deletes message keys but keeps the chain key of the message before, to retry the last send."""
    def __init__(self): self.ck, self.prev = ROOT_CK, b""
    def send(self):
        self.prev = self.ck
        self.ck, mk = step(self.ck)
        return mk
    def state(self): return [self.ck, self.prev]
    def advance(self, state): return state


class DeletesAsItGoes:
    """The ratchet: only the next chain key is kept."""
    def __init__(self): self.ck = ROOT_CK
    def send(self):
        self.ck, mk = step(self.ck)
        return mk
    def state(self): return [self.ck]
    def advance(self, state): return state


DESIGNS = [OneKey, KeepsUsedKeys, KeepsPreviousChainKey, DeletesAsItGoes]


def audit(design) -> list[int]:
    """Numbers of the messages sent before the theft whose keys the stolen copy holds or can derive."""
    d = design()
    keys, stolen = [], []
    for n in range(SENT):
        if n == THEFT_AFTER:
            stolen = d.state()
        keys.append(d.send())
    earlier: dict[bytes, set[int]] = {}
    for n in range(THEFT_AFTER):
        earlier.setdefault(keys[n], set()).add(n)  # one key may belong to several messages
    found: set[int] = set()
    for value in stolen:  # keys sitting in the copy
        found |= earlier.get(value, set())
    for ck in d.advance(stolen):  # keys the copy can reach by running a chain key forward
        probe = ck
        for _ in range(FORWARD_STEPS):
            probe, mk = step(probe)
            found |= earlier.get(mk, set())
    return sorted(found)


first = []
ck = ROOT_CK
for n in range(5):
    ck, mk = step(ck)
    first.append(f"#{n} {fingerprint(mk)}")
print("first five message keys: " + "  ".join(first))
print()
for design in DESIGNS:
    leaked = audit(design)
    noun = "message" if len(leaked) == 1 else "messages"
    verdict = "PASS" if not leaked else f"FAIL: can recover the keys of {noun} " + ", ".join(map(str, leaked))
    print(f"{design.__name__:<22} {verdict}")

first five message keys: #0 52808d97  #1 5ee9147c  #2 2caa5f1a  #3 a00f5eec  #4 22b31b27

OneKey                 FAIL: can recover the keys of messages 0, 1, 2, 3, 4
KeepsUsedKeys          FAIL: can recover the keys of messages 0, 1, 2, 3, 4
KeepsPreviousChainKey  FAIL: can recover the keys of message 4
DeletesAsItGoes        PASS

// A Double Ratchet on X25519 (sections 3.2 to 3.5 of the Double Ratchet specification, without header encryption):
// Alice and Bob talk for seven messages, and a thief copies Alice's state twice, once after message 1 and once after
// message 4. A thief clone of Alice then tries to read every message. Keys come from fixed lab seeds.
import { createCipheriv, createDecipheriv, createHmac, createPrivateKey, createPublicKey, diffieHellman, hkdfSync } from 'node:crypto';

const PKCS8 = Buffer.from('302e020100300506032b656e04220420', 'hex');
const SPKI = Buffer.from('302a300506032b656e032100', 'hex');
const privateKeyOf = (seed) => createPrivateKey({ key: Buffer.concat([PKCS8, seed]), format: 'der', type: 'pkcs8' });
const publicOf = (seed) => Buffer.from(createPublicKey(privateKeyOf(seed)).export({ type: 'spki', format: 'der' }).subarray(-32));
const dh = (seed, pub) => diffieHellman({ privateKey: privateKeyOf(seed), publicKey: createPublicKey({ key: Buffer.concat([SPKI, pub]), format: 'der', type: 'spki' }) });
const hmac = (key, byte) => createHmac('sha256', key).update(Buffer.from([byte])).digest();
const kdfCk = (ck) => ({ mk: hmac(ck, 0x01), ck: hmac(ck, 0x02) });
const kdfRk = (rk, out) => {
  const okm = Buffer.from(hkdfSync('sha256', out, rk, 'msc double ratchet demo', 64));
  return { rk: okm.subarray(0, 32), ck: okm.subarray(32) };
};

function seal(mk, text, ad) {
  const okm = Buffer.from(hkdfSync('sha256', mk, Buffer.alloc(32), 'msc ratchet demo v1', 44));
  const c = createCipheriv('aes-256-gcm', okm.subarray(0, 32), okm.subarray(32, 44));
  c.setAAD(ad);
  return Buffer.concat([c.update(text, 'utf8'), c.final(), c.getAuthTag()]);
}
function open(mk, blob, ad) {
  const okm = Buffer.from(hkdfSync('sha256', mk, Buffer.alloc(32), 'msc ratchet demo v1', 44));
  const d = createDecipheriv('aes-256-gcm', okm.subarray(0, 32), okm.subarray(32, 44));
  d.setAAD(ad);
  d.setAuthTag(blob.subarray(-16));
  try {
    return Buffer.concat([d.update(blob.subarray(0, -16)), d.final()]).toString('utf8');
  } catch {
    return null;
  }
}

// Each party draws its fresh ratchet keys from its own stream of seeds. A thief has a different stream.
const stream = (label) => {
  let i = 0;
  return () => createHmac('sha256', label).update(String(i++)).digest();
};

class Party {
  constructor(label, next) {
    this.label = label;
    this.next = next;
    this.s = { dhs: null, dhr: null, rk: null, cks: null, ckr: null, ns: 0, nr: 0, pn: 0 };
  }
  static initAlice(next, sk, bobPub) {
    const p = new Party('alice', next);
    p.s.dhs = next();
    p.s.dhr = bobPub;
    ({ rk: p.s.rk, ck: p.s.cks } = kdfRk(sk, dh(p.s.dhs, bobPub)));
    return p;
  }
  static initBob(next, sk, bobSeed) {
    const p = new Party('bob', next);
    p.s.dhs = bobSeed;
    p.s.rk = sk;
    return p;
  }
  static stolen(state, next) {
    const p = new Party('thief', next);
    p.s = { ...state };
    return p;
  }
  snapshot() {
    return { ...this.s };
  }
  send(text) {
    const { mk, ck } = kdfCk(this.s.cks);
    this.s.cks = ck; // the old chain key is gone
    const header = { pub: publicOf(this.s.dhs), n: this.s.ns++ };
    return { header, blob: seal(mk, text, Buffer.concat([header.pub, Buffer.from([header.n])])) };
  }
  ratchet(pub) {
    this.s.pn = this.s.ns;
    this.s.ns = 0;
    this.s.nr = 0;
    this.s.dhr = pub;
    ({ rk: this.s.rk, ck: this.s.ckr } = kdfRk(this.s.rk, dh(this.s.dhs, pub)));
    this.s.dhs = this.next(); // a fresh key pair; the old private key is overwritten
    ({ rk: this.s.rk, ck: this.s.cks } = kdfRk(this.s.rk, dh(this.s.dhs, pub)));
  }
  receive(msg) {
    if (!this.s.dhr || !msg.header.pub.equals(this.s.dhr)) this.ratchet(msg.header.pub);
    const { mk, ck } = kdfCk(this.s.ckr);
    this.s.ckr = ck;
    this.s.nr++;
    return open(mk, msg.blob, Buffer.concat([msg.header.pub, Buffer.from([msg.header.n])]));
  }
  // What a copy of Alice can do with a message Alice herself sent: run the sending chain forward, if the key is the same.
  readOwn(msg) {
    if (!this.s.dhs || !publicOf(this.s.dhs).equals(msg.header.pub) || msg.header.n < this.s.ns) return null;
    let ck = this.s.cks;
    let mk;
    for (let n = this.s.ns; n <= msg.header.n; n++) ({ mk, ck } = kdfCk(ck));
    return open(mk, msg.blob, Buffer.concat([msg.header.pub, Buffer.from([msg.header.n])]));
  }
}

const SK = Buffer.alloc(32, 0x42);
const bobSeed = stream('bob')(); // Bob's first ratchet key, known to Alice like a signed prekey
const alice = Party.initAlice(stream('alice'), SK, publicOf(bobSeed));
const bob = Party.initBob(stream('bob.next'), SK, bobSeed);

const script = [
  ['alice', 'hello'],
  ['alice', 'are you there'],
  ['bob', 'yes, I am here'],
  ['bob', 'send the file'],
  ['alice', 'sending now'],
  ['bob', 'got it'],
  ['alice', 'thanks'],
];
const transcript = [];
const thefts = { 1: null, 4: null };
script.forEach(([who, text], i) => {
  const sender = who === 'alice' ? alice : bob;
  const receiver = who === 'alice' ? bob : alice;
  const msg = { from: who, header: null, blob: null, text };
  Object.assign(msg, sender.send(text));
  if (receiver.receive(msg) !== text) throw new Error(`message ${i} did not arrive`);
  transcript.push(msg);
  if (i in thefts) thefts[i] = alice.snapshot();
});

// A thief with Alice's copy tries every message: the earlier ones with a fresh clone each, the later ones in order.
function attempt(theftIndex) {
  const state = thefts[theftIndex];
  const result = [];
  const thief = Party.stolen(state, stream(`thief ${theftIndex}`));
  transcript.forEach((msg, i) => {
    let text;
    if (i <= theftIndex) {
      const fresh = Party.stolen(state, stream(`thief ${theftIndex}`));
      text = msg.from === 'alice' ? fresh.readOwn(msg) : fresh.receive(msg);
    } else {
      text = msg.from === 'alice' ? thief.readOwn(msg) : thief.receive(msg);
    }
    result.push(text === msg.text);
  });
  return result;
}
const first = attempt(1);
const second = attempt(4);

// Names for the ratchet public keys, in the order they first appear: A0, A1, A2 for Alice, B0 (never used to send), B1, B2 for Bob.
const names = new Map([[publicOf(bobSeed).toString('hex'), 'B0']]);
for (const who of ['alice', 'bob']) {
  let count = who === 'alice' ? 0 : 1;
  for (const msg of transcript) {
    const hex = msg.header.pub.toString('hex');
    if (msg.from === who && !names.has(hex)) names.set(hex, `${who === 'alice' ? 'A' : 'B'}${count++}`);
  }
}

console.log('message  sender  ratchet key  thief with the copy after message 1   thief with the copy after message 4');
transcript.forEach((msg, i) => {
  const key = names.get(msg.header.pub.toString('hex'));
  const col = (r, at) => `${i <= at ? '(before) ' : ''}${r[i] ? 'READS' : 'refused'}`;
  console.log(`${String(i).padEnd(8)} ${msg.from.padEnd(7)} ${key.padEnd(12)} ${col(first, 1).padEnd(38)}${col(second, 4)}`);
});

message  sender  ratchet key  thief with the copy after message 1   thief with the copy after message 4
0        alice   A0           (before) refused                      (before) refused
1        alice   A0           (before) refused                      (before) refused
2        bob     B1           READS                                 (before) refused
3        bob     B1           READS                                 (before) refused
4        alice   A1           refused                               (before) refused
5        bob     B2           refused                               READS
6        alice   A2           refused                               refused

// The X3DH key agreement of the Signal specification (four Diffie-Hellman values, then HKDF), then the PQXDH change that
// puts a KEM secret behind them (ML-KEM-1024), and a bundle check that a client must make. Bob's prekey signature is made
// with a separate Ed25519 key here; Signal's XEdDSA signs with the X25519 identity key itself. X25519 keys come from fixed
// lab seeds; the ML-KEM keys are random, so the example prints only what does not depend on them.
import {
  createCipheriv, createDecipheriv, createHash, createPrivateKey, createPublicKey, decapsulate, diffieHellman, encapsulate,
  generateKeyPairSync, hkdfSync, sign, verify,
} from 'node:crypto';

const seed = (label) => createHash('sha256').update(`msc x3dh demo ${label}`).digest();
const X_PKCS8 = Buffer.from('302e020100300506032b656e04220420', 'hex');
const X_SPKI = Buffer.from('302a300506032b656e032100', 'hex');
const E_PKCS8 = Buffer.from('302e020100300506032b657004220420', 'hex');
const xPrivate = (label) => createPrivateKey({ key: Buffer.concat([X_PKCS8, seed(label)]), format: 'der', type: 'pkcs8' });
const xPublic = (label) => Buffer.from(createPublicKey(xPrivate(label)).export({ type: 'spki', format: 'der' }).subarray(-32));
const dh = (privLabel, pub) => diffieHellman({ privateKey: xPrivate(privLabel), publicKey: createPublicKey({ key: Buffer.concat([X_SPKI, pub]), format: 'der', type: 'spki' }) });

const F = Buffer.alloc(32, 0xff); // keeps the KDF input apart from signature inputs, as the specification says
function kdf(parts, info) {
  return Buffer.from(hkdfSync('sha256', Buffer.concat([F, ...parts]), Buffer.alloc(32), info, 32));
}

// Bob's prekey bundle, as the server hands it out.
const bobSigning = createPrivateKey({ key: Buffer.concat([E_PKCS8, seed('bob signing')]), format: 'der', type: 'pkcs8' });
const bobVerify = createPublicKey(bobSigning);
const pq = generateKeyPairSync('ml-kem-1024');
const pqRaw = pq.publicKey.export({ format: 'raw-public' });
const bundle = {
  ik: xPublic('bob identity'),
  spk: xPublic('bob signed prekey'),
  opk: xPublic('bob one-time prekey'),
  pqpk: pqRaw,
};
bundle.spkSig = sign(null, bundle.spk, bobSigning);
bundle.pqSig = sign(null, bundle.pqpk, bobSigning);

function checkBundle(b) {
  if (!verify(null, b.spk, bobVerify, b.spkSig)) throw new Error('signed prekey does not verify');
  if (b.pqpk && !verify(null, b.pqpk, bobVerify, b.pqSig)) throw new Error('post-quantum prekey does not verify');
}

// Alice's side: ephemeral key EK, then DH1..DH4 and, for PQXDH, the KEM secret.
function aliceAgrees(b, { verifyBundle = true, useOpk = true, usePq = false } = {}) {
  if (verifyBundle) checkBundle(b);
  const dhs = [dh('alice identity', b.spk), dh('alice ephemeral', b.ik), dh('alice ephemeral', b.spk)];
  if (useOpk) dhs.push(dh('alice ephemeral', b.opk));
  let ct = null;
  if (usePq) {
    const out = encapsulate(createPublicKey({ key: b.pqpk, format: 'raw-public', asymmetricKeyType: 'ml-kem-1024' }));
    ct = out.ciphertext;
    dhs.push(out.sharedKey);
  }
  return { sk: kdf(dhs, usePq ? 'msc-demo_CURVE25519_SHA-256_ML-KEM-1024' : 'msc-demo'), ct };
}

// Bob's side from Alice's first message: her identity key, her ephemeral key and which prekeys she used.
function bobAgrees({ useOpk = true, ct = null } = {}) {
  const ikA = xPublic('alice identity');
  const ekA = xPublic('alice ephemeral');
  const dhs = [dh('bob signed prekey', ikA), dh('bob identity', ekA), dh('bob signed prekey', ekA)];
  if (useOpk) dhs.push(dh('bob one-time prekey', ekA));
  if (ct) dhs.push(decapsulate(pq.privateKey, ct));
  return kdf(dhs, ct ? 'msc-demo_CURVE25519_SHA-256_ML-KEM-1024' : 'msc-demo');
}

const eq = (a, b) => a.equals(b);
const x3dh = aliceAgrees(bundle);
console.log(`X3DH, four DH values:        both sides derive the same 32 bytes: ${eq(x3dh.sk, bobAgrees())}`);
const noOpk = aliceAgrees(bundle, { useOpk: false });
console.log(`no one-time prekey, three DH: same on both sides: ${eq(noOpk.sk, bobAgrees({ useOpk: false }))}, same as with four: ${eq(noOpk.sk, x3dh.sk)}`);
const pqxdh = aliceAgrees(bundle, { usePq: true });
console.log(`PQXDH, plus ML-KEM-1024:     same on both sides: ${eq(pqxdh.sk, bobAgrees({ ct: pqxdh.ct }))}, same as X3DH: ${eq(pqxdh.sk, x3dh.sk)}`);
console.log(`ML-KEM-1024 sizes:           public key ${pqRaw.length} bytes, ciphertext ${pqxdh.ct.length} bytes, secret 32 bytes`);

// The initial message is sealed under a key from SK, with both identity keys as associated data.
const ad = Buffer.concat([xPublic('alice identity'), bundle.ik]);
const okm = Buffer.from(hkdfSync('sha256', x3dh.sk, Buffer.alloc(32), 'msc-demo initial message', 44));
const c = createCipheriv('aes-256-gcm', okm.subarray(0, 32), okm.subarray(32));
c.setAAD(ad);
const blob = Buffer.concat([c.update('hello bob', 'utf8'), c.final(), c.getAuthTag()]);
function openInitial(sk, associated) {
  const k = Buffer.from(hkdfSync('sha256', sk, Buffer.alloc(32), 'msc-demo initial message', 44));
  const d = createDecipheriv('aes-256-gcm', k.subarray(0, 32), k.subarray(32));
  d.setAAD(associated);
  d.setAuthTag(blob.subarray(-16));
  try {
    return Buffer.concat([d.update(blob.subarray(0, -16)), d.final()]).toString('utf8');
  } catch {
    return 'refused';
  }
}
console.log(`\ninitial message at Bob:       ${openInitial(bobAgrees(), ad)}; with another identity in the associated data: ${openInitial(bobAgrees(), Buffer.concat([xPublic('mallory identity'), bundle.ik]))}`);

// The check a client must make: a bundle whose signed prekey was swapped must be refused. A server that swaps the prekey
// for one it holds, and withholds the one-time prekey, can derive the key of the session if Bob's identity key leaks later.
const swapped = { ...bundle, spk: xPublic('swapped prekey') }; // the old signature does not cover this key
for (const verifyBundle of [false, true]) {
  try {
    const r = aliceAgrees(swapped, { verifyBundle, useOpk: false });
    const laterLeak = [dh('swapped prekey', xPublic('alice identity')), dh('bob identity', xPublic('alice ephemeral')), dh('swapped prekey', xPublic('alice ephemeral'))];
    console.log(`client ${verifyBundle ? 'checks' : 'skips'} the signature: accepts the swapped prekey; if Bob's identity key leaks later, its holder derives the same SK: ${eq(r.sk, kdf(laterLeak, 'msc-demo'))}`);
  } catch (e) {
    console.log(`client ${verifyBundle ? 'checks' : 'skips'} the signature: refuses the swapped prekey (${e.message})`);
  }
}

X3DH, four DH values:        both sides derive the same 32 bytes: true
no one-time prekey, three DH: same on both sides: true, same as with four: false
PQXDH, plus ML-KEM-1024:     same on both sides: true, same as X3DH: false
ML-KEM-1024 sizes:           public key 1568 bytes, ciphertext 1568 bytes, secret 32 bytes

initial message at Bob:       hello bob; with another identity in the associated data: refused
client skips the signature: accepts the swapped prekey; if Bob's identity key leaks later, its holder derives the same SK: true
client checks the signature: refuses the swapped prekey (signed prekey does not verify)

"""What one key update costs in an MLS ratchet tree, counted in HPKE ciphertexts (RFC 9420 sections 4.1, 7.5 and 12.1.3).

The model keeps the rules that decide the count: a perfect binary tree whose leaves hold the members, a node is blank
until a commit passes through it, the resolution of a blank node is the resolutions of its two children, and a commit
encrypts each new path secret once for every node in the resolution of the child that is not on the committer's path.
It leaves out unmerged leaves, the Welcome message and the tree growing; pairwise sessions need N - 1 ciphertexts."""
from math import ceil, log2


class Tree:
    def __init__(self, members: int):
        self.members = members
        self.depth = max(1, ceil(log2(members)))
        self.size = 2**self.depth  # leaves; nodes are numbered like a heap: the root is 1, the leaves are size.. 2*size-1
        self.keyed = [False] * (2 * self.size)  # a leaf: a member is there; an inner node: it holds a key
        for i in range(members):
            self.keyed[self.size + i] = True

    def resolution(self, node: int) -> int:
        """How many nodes must be encrypted to so that everyone under this node is reached."""
        if self.keyed[node]:
            return 1
        if node >= self.size:  # a blank leaf: nobody is there
            return 0
        return self.resolution(2 * node) + self.resolution(2 * node + 1)

    def cost(self, leaf: int) -> int:
        """Ciphertexts in the update path of a commit by the member at this leaf."""
        node, total = self.size + leaf, 0
        while node > 1:
            total += self.resolution(node ^ 1)  # the child that is not on the committer's path
            node >>= 1
        return total

    def commit(self, leaf: int) -> int:
        """The same commit, applied: every node on the path whose other side is not empty gets a new key."""
        total = self.cost(leaf)
        node = self.size + leaf
        while node > 1:
            self.keyed[node >> 1] = self.resolution(node ^ 1) > 0
            node >>= 1
        return total

    def key_everything(self) -> None:
        """The state after every member has committed once: each node above a member holds a key."""
        for node in range(self.size - 1, 0, -1):
            self.keyed[node] = self.keyed[2 * node] or self.keyed[2 * node + 1]

    def remove(self, leaf: int) -> None:
        """A removed member's leaf and every node above it are blanked (RFC 9420 section 12.1.3)."""
        node = self.size + leaf
        while node >= 1:
            self.keyed[node] = False
            node >>= 1


print(f"{'members':>8} {'depth':>5} {'pairwise':>9} {'first update':>13} {'update, tree keyed':>19}")
for n in (8, 100, 1000, 10000, 100000):
    fresh, ready = Tree(n), Tree(n)
    ready.key_everything()
    print(f"{n:>8,} {fresh.depth:>5} {n - 1:>9,} {fresh.cost(0):>13,} {ready.cost(0):>19,}")

# 1,000 members with a keyed tree. A quarter of them leave, one commit each, by the member next to the leaver.
t = Tree(1000)
t.key_everything()
leavers = list(range(4, 1000, 4))
commit_costs = []
for leaver in leavers:
    t.remove(leaver)
    commit_costs.append(t.commit(leaver + 1))
present = [m for m in range(1000) if m not in set(leavers)]
update_costs = [t.cost(m) for m in present]
print(f"\n{len(leavers)} of 1,000 members leave, one commit each: those commits cost {min(commit_costs)} to {max(commit_costs)} ciphertexts")
print(f"then an update by each of the {len(present)} remaining members would cost {min(update_costs)} to {max(update_costs)}, {sum(update_costs) / len(update_costs):.1f} on average")

 members depth  pairwise  first update  update, tree keyed
       8     3         7             7                   3
     100     7        99            99                   7
   1,000    10       999           999                  10
  10,000    14     9,999         9,999                  14
 100,000    17    99,999        99,999                  17

249 of 1,000 members leave, one commit each: those commits cost 7 to 9 ciphertexts
then an update by each of the 751 remaining members would cost 7 to 10, 9.7 on average