Skip to content

Add ML-KEM-1024-P384 and pure ML-KEM-768/1024#105

Open
emil-wire wants to merge 6 commits into
rozbb:mainfrom
emil-wire:main
Open

Add ML-KEM-1024-P384 and pure ML-KEM-768/1024#105
emil-wire wants to merge 6 commits into
rozbb:mainfrom
emil-wire:main

Conversation

@emil-wire

Copy link
Copy Markdown

Hello @rozbb!

This PR adds the pure ML-KEM-768 and ML-KEM-1024 KEMs and the MLKEM1024-P384 hybrid KEM (per draft-ietf-hpke-pq-04, draft-irtf-cfrg-concrete-hybrid-kems-03, and draft-irtf-cfrg-hybrid-kems-11), closely mirroring the existing mlkem768p256.rs.
A few choices below aren't obviously "the one right way" so I'd rather check with you. None of these block the KEMs from working, KATs pass against the official upstream vectors. I've made a judgement call on some of them to get stuff working, but am happy to revert or modify the approach.

Pure ML-KEM keys re-expand on every use (no cached expanded key)

The seed-only private-key representation d‖z (Nsk=64) is mandated by hpke-pq-04 §3 and is not negotiable, but storing only the seed in memory is a choice I've made - so every decap/sk_to_pk re-runs the Keccak-heavy key expansion.
This should be fine for one-shot HPKE and MLS use (our use case, it drove this decision) and keeps the resident secret small (~64 B vs ~2.4–3.2 KB or so), but it's worth revisiting at some point with an internal expanded key cache if a high throughput "same-key-many-decaps" workload ever appears.

Feature flags

mlkem768 and mlkem1024 are independent features, following the per-KEM style of the existing flags. Happy to make mlkem1024 imply mlkem768 instead if you'd prefer fewer combinations to test.
Also, sha3 is now a non-optional dependency (previously feature-gated). It's needed unconditionally by the SHAKE KDFs and ML-KEM derive_keypair. Should I re-gate it behind a PQ flag or similar?

README changes

Minor thing, I moved the post-quantum KEMs/KDFs out of the "RFC 9180 primitives" checklist into a new "Draft Post-quantum extensions" section, so the "complies with RFC 9180" statement stays accurate (X-Wing and MLKEM768-P256 were also listed under it, I moved them).
Draft links are pinned to specific versions so it's easier to reason about when a bump is due or which specific version the implementation follows.

GenerateKeyPair conformance

draft-ietf-hpke-pq-04 defines GenerateKeyPair to sample the private-key seed directly:

  • pure ML-KEM: dk = random(64)
  • hybrids: seed = random(Nseed); return DeriveKeyPair(seed) (the generic hybrid DeriveKeyPair, no LabeledDerive)

The crate's shared gen_keypair_with_rng instead calls derive_keypair(random(Nsk)), and derive_keypair applies SHAKE256.LabeledDerive.
So a serialized private key from GenerateKeyPair is LabeledDerive(random(Nsk)) rather than random(Nsk) - it diverges from the draft's GenerateKeyPair under a fixed RNG.

To be clear: this is KAT-safe (the vectors exercise DeriveKeyPair, which is correct) and is AFAIK not a security issue (the output is still uniform). More of a conformance/interop divergence thing - it comes from the shared default, so it applies to X-Wing and MLKEM768-P256 as well, not just my changes.

How would you like to handle it?
I could override gen_keypair_with_rng for all PQ KEMs to sample the seed directly to be conformant and internally consistent, but would require some changes to X-Wing/MLKEM768-P256 files.
Or just leave as it is and just document it as purposeful simplification?

I'd prefer the former, but have no strong opinions on this.

No published vectors for ML-KEM + SHAKE256 combo

Pure ML-KEM (kem 0x0041/0x0042) + SHAKE256 (kdf 0x0011) has no published vector.
The hpke-pq vectors only pin pure ML-KEM under HKDF (A.2/A.3) and TurboSHAKE256 (A.13); the only SHAKE256 vector is X-Wing. (A.7 and A.11 are titled "SHAKE256" in draft-04 but their bodies use SHAKE128 / kdf 0x0010 - checked against the vector files, seems to be a bug in the spec?)

I decided to pin the two halves separately rather than invent a vector: X-Wing covers the SHAKE256 KDF, the HKDF vectors cover ML-KEM encap/decap (plus the direct decap check in test_case).

Unsupported Auth / AuthPSK modes

The PQ KEMs are Auth = no and currently panic if an authenticated mode is used, mirroring the existing X-Wing/MLKEM768-P256 behaviour. A colleague suggested returning EncapError/DecapError (or rejecting the mode in setup) instead, so a mode choice doesn't trigger a library panic. I'm glad to switch, but it would be a breaking crate-wide behavior change - so keep panics for consistency, or move all PQ KEMs to errors?

kat feature pulling aes/chacha/getrandom - was this on purpose?

cargo test --no-default-features --features kat doesn't build, because kat enables the KEMs but not the AEADs/getrandom that the tests use. This predates this PR, and the feature's comment says it "implies … all other features", so I added the three. If you left them out on purpose (e.g. to keep a leaner kat matrix), let me know and I'll drop it.

Dead-code warnings when building without a DH KEM

I found a minor bug in the existing code you might want to know about: a PQ-only build (no x25519/p256/p384/p521) emits a few dead-code warnings - unused import: dhkem::* (kem.rs) and MAX_PUBKEY_SIZE/DhError/DhKeyExchange (dhkex.rs). This predates this PR (a mlkem768p256-only build triggers them too), so I left it alone - not sure what the consequences of changing this would be.

Zeroizing private-seed temporaries - apply to MLKEM768-P256 too?

Expanding a private key calls from_seed(&seed.into()), where .into() makes a short-lived copy of the secret seed bytes that isn't zeroized (if I'm not mistaken, please correct me). I tightened this in the new KEMs (bind the converted buffer, then zeroize it after use), but the existing MLKEM768-P256 still uses the original pattern (mlkem768p256.rs), so the PQ KEMs are now inconsistent on this point. It's a minor think (a brief stack copy that's dropped without wiping), but I'd rather not diverge without letting you know. Should I apply the same hardening to MLKEM768-P256, or revert mine to the existing convention?

Benchmarks for the PQ KEMs in a separate PR?

The current benches/ only cover the classical DH ciphersuites (P-256, X25519); no PQ/hybrid KEM is benchmarked today (including X-Wing or MLKEM768-P256). PQ keygen/encap/decap have a quite different performance profile, so benches could be interesting, but I didn't want to assume. Would you like me to add benchmarks for the new KEMs (and optionally the existing PQ ones) in a later PR (might take a bit of time though)?

Comments

I left a bunch of comments in the code. I used them as reference and sanity checks during implementation and initially wanted to remove them - but felt it might be sensible to leave them in to make reviewing, changing (if the drafts change or are ratified) and auditing easier.


Happy to make any change you deem necessary.

Thanks for taking a look and your time,
Emil

Wire Security Engineering Team

@rozbb

rozbb commented Jun 25, 2026

Copy link
Copy Markdown
Owner

Amazing, thank you! I may take a little bit to get through this but I am immensely grateful!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants