A KEM Is Not Diffie-Hellman: Rewiring Protocol Code for ML-KEM

Most protocol code written in the last twenty years contains a quiet assumption. Somewhere there is a call that looks like agree(my secret, your public value), both sides make the same call, and both arrive at the same shared secret. Module-Lattice-Based Key-Encapsulation Mechanism (ML-KEM), the post-quantum standard published by the United States National Institute of Standards and Technology (NIST) as Federal Information Processing Standard (FIPS) 203, does not offer that call. It offers encapsulate(your public key), which returns a ciphertext and a secret, and decapsulate(my private key, ciphertext), which returns the same secret. The difference looks cosmetic, and that is exactly why it causes trouble. Replace one call with the other inside a working protocol and the code will usually compile, the happy path may even pass its tests, and yet the message order, the place where state lives, the way identity is proven and the way failures show up have all changed underneath. This article walks through each of those changes and ends with a checklist for rewiring protocol code safely.

Two shapes for the same job

Diffie-Hellman (DH), and its elliptic-curve form, Elliptic-Curve Diffie-Hellman (ECDH), is symmetric in shape. Each party picks a secret number, publishes a matching public value, and combines its own secret with the other party's public value. Neither side is special. The two public values can even be sent at the same time, in either order, because each side's contribution is independent of the other's.

A key-encapsulation mechanism (KEM) is asymmetric in shape. One party owns a key pair: a public encapsulation key and a private decapsulation key. The other party, the encapsulator, takes that public key and runs encapsulation, which manufactures a fresh random secret and a ciphertext that carries it. Only the ciphertext travels back. The key owner runs decapsulation on the ciphertext and recovers the same secret. One side generates keys; the other side generates the secret. The roles are not interchangeable.

There is a revealing one-way relationship between the two shapes. Hybrid Public Key Encryption (HPKE), standardized as Request for Comments (RFC) 9180, defines its whole design around a KEM interface, and it turns Diffie-Hellman into a KEM with a construction called DHKEM: to encapsulate, generate a throwaway Diffie-Hellman key pair, combine its secret with the recipient's public key, and send the throwaway public value as the ciphertext. So any Diffie-Hellman can be dressed up as a KEM. The reverse does not work. A KEM cannot be dressed up as Diffie-Hellman, because it has no operation that combines two parties' long-term keys. Every difference in the rest of this article follows from that asymmetry.

The practical consequence is that a protocol written around agree() has quietly relied on properties that a KEM does not have, and a protocol written around encapsulate() needs decisions that a Diffie-Hellman protocol never had to make. Treat the change as a protocol redesign with a small blast radius, not as a library upgrade.

Message order and who keeps state

Because the encapsulator needs the key owner's public key before it can do anything, the message order is fixed: key first, ciphertext second. In a Diffie-Hellman design, both parties could fire off their public values in the first flight. In a KEM design, whoever will encapsulate must wait until it has received, or already holds, the other side's encapsulation key.

The second change is where state lives. The key owner generates a key pair, sends the public half, and must then keep the private decapsulation key alive until the ciphertext arrives. That is memory per pending session, a timeout, a rule for what happens if the reply never comes, and a duty to erase the key afterwards. The encapsulator, by contrast, can be almost stateless: it receives a key, encapsulates immediately, sends the ciphertext and already holds the shared secret.

Transport Layer Security (TLS) version 1.3, the protocol under HTTPS, shows the choice in practice. In the hybrid post-quantum key exchange standardized in RFC 10024, the client generates the ML-KEM key pair and places the encapsulation key in its first message; the server encapsulates and returns the ciphertext. The server never generates an ML-KEM key pair at all, and the client is the side that must hold a decapsulation key across the round trip. That suits a busy server, which stays cheap and stateless, and places the memory and timing burden on the side that started the conversation.

When you design your own protocol, choosing who owns the key is therefore a real lever. Make the side that can afford state the key owner. For a service that receives bursts of connections, that is usually the caller. For a device that wakes up briefly and must send one message, it is usually the receiver, whose long-term encapsulation key the device already knows.

Authentication no longer falls out of the keys

Diffie-Hellman has a property many protocols lean on without naming it: two long-term keys can be combined directly. If Alice combines her long-term secret with Bob's long-term public key, and Bob does the reverse, they reach a shared value that only those two key holders could compute. This static-static combination authenticates both parties implicitly, without a single signature. Designs in the Noise protocol framework and Signal's Extended Triple Diffie-Hellman (X3DH) handshake mix such combinations into their key derivation to prove identity.

A KEM has no static-static mode: a key can only be encapsulated to, and two key pairs cannot be combined. How the post-quantum extension of HPKE copes with this, and why its KEMs give up HPKE's authenticated modes, is analysed in the Research article 'ML-KEM Inside HPKE, X-Wing and SSH: What Each Construction Assumes of the KEM'.

That leaves three honest options. First, sign: keep the KEM for secrecy and add a digital signature from a long-term signing key, which is what TLS does with certificates. Second, use a pre-shared key (PSK), a secret both sides already hold, mixed into the key derivation. Third, authenticate with KEMs: each party encapsulates to the other's long-term KEM key and both secrets feed the final keys. The third option avoids signatures but costs an extra message flight in an interactive protocol, because each side must first learn the other's key and send a ciphertext to it; the research protocol known as KEMTLS explores exactly this trade.

Signal's post-quantum handshake, PQXDH, illustrates a fourth, hybrid path. It keeps the classical Diffie-Hellman steps that provide authentication and adds an ML-KEM encapsulation to a signed post-quantum prekey, mixing the KEM secret into the derived keys. Confidentiality gains post-quantum protection; authentication still rests on the classical keys and the signature on the prekey. That is a deliberate, documented choice, not an accident of swapping primitives.

For teams that already use HPKE's authenticated modes, the migration is a concrete piece of work. RFC 9180 numbers its four modes: base (0x00), pre-shared key (0x01), auth (0x02) and auth with pre-shared key (0x03). Search code and configuration for the last two first, because those are exactly the call sites that lose sender authentication when the KEM changes, and list them with the owner of each.

Each call site then gets an explicit replacement. Where sender and receiver can be provisioned with a shared secret in advance, move to mode 0x01 and provision a pre-shared key and its identifier; RFC 9180 requires that key to carry at least 32 bytes of entropy. Where they cannot, keep base mode and have the sender sign the encapsulated value, the info string and the ciphertext with its long-term signing key, and have the receiver verify that signature before it opens anything.

The failure to watch for is silent. A receiver that drops to base mode without adding a signature or a pre-shared key still decrypts every message; it simply no longer knows who sent it, and no existing test fails. Add a negative test in which an unregistered sender seals a well-formed message that the receiver must reject, version the message format so old authenticated-mode messages cannot be confused with new ones, and log the authenticated sender of every accepted message, so a gap shows up in monitoring rather than in an incident.

Failures that do not announce themselves

In much classical code, a malformed key-exchange message produces an error at the point of processing. ML-KEM deliberately does not. Its decapsulation uses implicit rejection: when a ciphertext of the right length fails the internal validity check, decapsulation still returns a secret, but a pseudorandom one derived from a hidden value and the ciphertext, unrelated to anything the sender holds. The reason is security: an error message would tell an attacker which of its crafted ciphertexts were valid, and that signal is the raw material of chosen-ciphertext attacks. The Go standard library documents the behaviour directly: decapsulation returns an error only when the ciphertext is the wrong length, and a correct-length but invalid ciphertext produces a shared key that does not match the sender's.

For protocol code this moves the failure. Decapsulation succeeds; the session keys are simply wrong. The mismatch becomes visible only later, when the first authenticated message fails to verify under an authenticated-encryption-with-associated-data (AEAD) cipher, or when an explicit key-confirmation check fails. TLS 1.3 catches it at the Finished message, which proves both sides derived the same keys from the same transcript.

Three habits follow. Do not try to detect a 'bad ciphertext' yourself, and do not log decapsulation as failed or succeeded, because any such signal recreates the oracle the standard went to some trouble to remove. Bind the ciphertext, the encapsulation key and the protocol context into the transcript or the key derivation function (KDF), so that the derived keys depend on exactly what was exchanged. And confirm keys explicitly before trusting the channel with anything that matters, then report a single generic handshake failure if confirmation fails.

Inputs worth checking before you trust them

Implicit rejection protects the key owner from bad ciphertexts, but FIPS 203 also defines explicit checks for the other inputs, and a protocol should make sure they actually run. Before encapsulating to a peer's encapsulation key, the standard calls for a type check, that the key has exactly the right length, and a modulus check, that decoding the key and re-encoding it gives back the same bytes, which confirms that every number inside it lies below the scheme's modulus of 3329. A key that fails either check must not be used.

On the decapsulation side, the standard checks the ciphertext length and, when a private decapsulation key is imported from outside, its length and a hash check: the private key carries a copy of the public key and a hash of it, and the two must agree. That check catches keys that were corrupted or stitched together from mismatched parts.

Libraries differ in whether they run these checks automatically, at import, at use, or not at all, and in whether they reject or silently accept a key of the wrong length. Read the documentation of the one you call, and if it is silent, add the checks at the protocol boundary where peer keys arrive.

Reuse, lifetime and forward secrecy

Ephemeral Diffie-Hellman taught a reflex: generate fresh key pairs for every session so that stealing a long-term key later reveals nothing about past traffic, a property called forward secrecy. The same reflex applies to ML-KEM, with one nuance that matters for engineering budgets.

ML-KEM is designed to be secure against chosen-ciphertext attack, often written IND-CCA, which means one key pair can safely receive many ciphertexts, including ones crafted by an attacker. Reusing a key pair across several sessions is therefore not an algorithm break. It is a forward-secrecy decision: if that private key is ever stolen, every session that used it is exposed. Generating a fresh key pair costs a little computation on the key owner and nothing on the network that a reused key would save, so the default should remain one key pair per session.

Some designs deliberately use long-lived KEM keys: a recipient that publishes an encapsulation key so that senders can seal messages to it without a round trip, or a storage system that wraps data keys to a service key. Those designs give up forward secrecy for convenience, so they need key rotation, a lifetime, and, where it matters, an additional ratchet or ephemeral layer on top. Whatever the lifetime, erase decapsulation keys and shared secrets from memory as soon as they are no longer needed.

Where AI agents meet this

AI agents, software that plans and takes actions by calling tools, increasingly talk to other agents, tool servers and message queues over channels that are not plain TLS. Encrypted tool results, sealed task payloads placed on a queue for another agent, and custom handshakes inside agent frameworks are all places where someone once wrote an agree() call. Much of that code is now written or modified by AI coding assistants.

A coding assistant asked to 'replace X25519 with ML-KEM' in such a codebase will often produce something that compiles and passes a happy-path test. It may keep a static-static step that no longer exists, silently dropping sender authentication; it may move the private key to the wrong side so that a busy service now holds state for every caller; or it may add an error branch around decapsulation that reintroduces the very signal implicit rejection removes. None of these shows up as a failing build.

Two review rules prevent most of the damage. First, treat any change from a Diffie-Hellman call to a KEM call in protocol code as a design change that needs the questions in this article answered in the pull request, not as a dependency bump. Second, for sealed agent-to-agent messages, prefer a standard KEM-based construction, such as HPKE with a post-quantum KEM, over hand-rolled code, and where the old design relied on HPKE's authenticated mode, add a signature or a pre-shared key explicitly.

A rewiring checklist for practitioners

The decisions below are the ones a practitioner actually faces when moving protocol code from Diffie-Hellman to ML-KEM, whether alone or in a hybrid with a classical exchange. Answer each one in writing before the code changes, and the change becomes reviewable.

The short code sketch shows the shape the protocol must now follow. It is deliberately generic: real libraries name the calls differently, but every one of them returns the ciphertext and the secret together on one side and needs the private key and ciphertext on the other.

owner:        (ek, dk) = KeyGen();  send(ek)
encapsulator: check(ek);  (ct, ss) = Encapsulate(ek);  send(ct)
owner:        ss = Decapsulate(dk, ct);  erase(dk)
both:         keys = KDF(ss, transcript(ek, ct, context))
both:         confirm(keys)  # the only place a failure appears

1. Decide who owns the KEM key pair and who encapsulates, based on which side can afford to hold state.
2. Fix the message order: the encapsulation key must reach the encapsulator before anything else depends on the secret.
3. Replace any static-static Diffie-Hellman authentication with a signature, a pre-shared key, or KEM-based authentication with its extra flight.
4. Run the encapsulation-key check and the length checks at the boundary where peer input arrives.
5. Bind the ciphertext, keys and context into the key derivation, and confirm keys before trusting the channel.
6. Report one generic failure; never log or branch on decapsulation outcomes.
7. Use fresh key pairs per session unless a documented design needs long-lived keys, and erase secrets after use.

Glossary

Key-encapsulation mechanism (KEM)

A scheme in which one party encapsulates a fresh secret to another party's public key, producing a ciphertext that only the matching private key can decapsulate.

Diffie-Hellman (DH)

A key-agreement method in which both parties combine their own secret with the other's public value to reach the same shared secret.

Static-static combination

Combining two parties' long-term Diffie-Hellman keys, which authenticates both implicitly and has no equivalent in a KEM.

Hybrid Public Key Encryption (HPKE)

An IETF standard, RFC 9180, for sealing messages to a recipient's public key, built around a KEM interface.

Implicit rejection

ML-KEM's behaviour of returning an unrelated pseudorandom secret, rather than an error, when a correctly sized ciphertext is invalid.

Key confirmation

A protocol step that proves both sides derived the same keys, such as the Finished message in TLS 1.3.

Forward secrecy

The property that stealing a long-term key later does not expose sessions that were completed earlier.

Pre-shared key (PSK)

A secret both parties hold in advance and mix into the key derivation to authenticate each other.

Encapsulation key check

The FIPS 203 length and modulus test that a peer's public key must pass before anyone encapsulates to it.

Negative test

A test that feeds a deliberately invalid or unauthorised input and passes only if the system rejects it.

Further reading

1. FIPS 203, Module-Lattice-Based Key-Encapsulation Mechanism Standard (NIST)
2. RFC 9180, Hybrid Public Key Encryption (IETF)
3. The PQXDH Key Agreement Protocol (Signal)
4. ML-KEM Inside HPKE, X-Wing and SSH (Research article, this site)