Adding TripleDhKem key exchange (#404)

.github/workflows/main.yml

@@ -15,6 +15,7 @@ jobs:
         backend_feature:
           -
           - --features ristretto255
+          - --features ristretto255,kem
           - --features curve25519
           - --features ecdsa
           - --features ed25519
@@ -73,6 +74,7 @@ jobs:
         backend_feature:
           -
           - --features ristretto255
+          - --features ristretto255,kem
           - --features curve25519
           - --features ecdsa
           - --features ed25519
@@ -167,6 +169,7 @@ jobs:
       matrix:
         backend_feature:
           - --features ristretto255
+          - --features ristretto255,kem
           -
     steps:
       - name: Checkout sources
@@ -204,7 +207,7 @@ jobs:
         uses: actions-rs/cargo@v1
         with:
           command: clippy
-          args: --all-targets --features argon2,std,curve25519,ecdsa,ed25519 -- -D warnings
+          args: --all-targets --features argon2,std,curve25519,ecdsa,ed25519,kem -- -D warnings
 
   format:
     name: cargo fmt
@@ -233,7 +236,7 @@ jobs:
           RUSTDOCFLAGS: -D warnings
         with:
           command: doc
-          args: --no-deps --document-private-items --features argon2,std,curve25519,ecdsa,ed25519
+          args: --no-deps --document-private-items --features argon2,std,curve25519,ecdsa,ed25519,kem
 
   taplo:
     name: Taplo

Cargo.toml

@@ -18,6 +18,7 @@ curve25519 = ["dep:curve25519-dalek"]
 default = ["ristretto255", "serde"]
 ecdsa = ["dep:ecdsa", "dep:rfc6979"]
 ed25519 = ["dep:curve25519-dalek", "dep:ed25519-dalek"]
+kem = ["dep:ml-kem"]
 ristretto255 = ["dep:curve25519-dalek", "voprf/ristretto255-ciphersuite"]
 serde = [
   "dep:serde",
@@ -53,6 +54,9 @@ elliptic-curve = { version = "0.13", features = ["hash2curve", "sec1"] }
 generic-array = "=0.14.7" # pinned to avoid deprecation warnings
 hkdf = "0.12"
 hmac = "0.12"
+ml-kem = { version = "0.2", default-features = false, features = [
+  "zeroize",
+], optional = true }
 rand = { version = "0.8", default-features = false }
 rfc6979 = { version = "0.4", optional = true }
 serde = { version = "1", default-features = false, features = [

src/errors.rs

@@ -24,6 +24,8 @@ pub enum InternalError {
     HmacError,
     /// Computing the key stretching function failed
     KsfError,
+    /// Error while performing a KEM operation
+    KemError,
     /** This error occurs when the envelope seal open hmac check fails
     HMAC check in seal open failed. */
     SealOpenHmacError,
@@ -43,6 +45,7 @@ impl Debug for InternalError {
             Self::HkdfError => f.debug_tuple("HkdfError").finish(),
             Self::HmacError => f.debug_tuple("HmacError").finish(),
             Self::KsfError => f.debug_tuple("KsfError").finish(),
+            Self::KemError => f.debug_tuple("KemError").finish(),
             Self::SealOpenHmacError => f.debug_tuple("SealOpenHmacError").finish(),
             Self::IncompatibleEnvelopeModeError => {
                 f.debug_tuple("IncompatibleEnvelopeModeError").finish()

src/key_exchange/mod.rs

@@ -13,6 +13,8 @@ pub mod group;
 pub(crate) mod shared;
 pub mod sigma_i;
 pub mod tripledh;
+#[cfg(feature = "kem")]
+pub mod tripledh_kem;
 
 use core::iter;
 use core::ops::Add;

src/key_exchange/tripledh_kem.rs

@@ -0,0 +1,696 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+// This source code is dual-licensed under either the MIT license found in the
+// LICENSE-MIT file in the root directory of this source tree or the Apache
+// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
+// of this source tree. You may select, at your option, one of the above-listed
+// licenses.
+
+//! TripleDH-KEM is a variant of the OPAQUE Triple Diffie-Hellman handshake in
+//! which the client supplies a KEM public key in KE1 and the server performs a
+//! KEM encapsulation in KE2 instead of relying solely on the final Diffie-
+//! Hellman hop. The server bundles the KEM ciphertext alongside the classic
+//! `TripleDH` payload, both parties absorb the ciphertext into the transcript
+//! and mix the encapsulated shared secret with the three Diffie-Hellman
+//! products when deriving handshake keys, and the client decapsulates during
+//! KE3 to recover that shared secret before validating the server MAC. This
+//! file contains the data model and trait glue that layer
+//! the generic `ml-kem` abstractions into the existing OPAQUE key-exchange
+//! pipeline.
+
+use core::convert::TryFrom;
+use core::fmt::Debug;
+use core::marker::PhantomData;
+use core::ops::Add;
+
+use derive_where::derive_where;
+use digest::core_api::BlockSizeUser;
+use digest::{Digest, Output};
+use generic_array::sequence::Concat;
+use generic_array::typenum::{IsLess, Le, NonZero, Sum, U256};
+use generic_array::{ArrayLength, GenericArray};
+use ml_kem::kem::{Decapsulate, Encapsulate};
+use ml_kem::{
+    Ciphertext as MlKemCiphertext, Encoded, EncodedSizeUser, KemCore, SharedKey as MlKemSharedKey,
+};
+use rand::{CryptoRng, RngCore};
+use subtle::{ConstantTimeEq, CtOption};
+use zeroize::{Zeroize, ZeroizeOnDrop};
+
+use super::shared::{self, Ke1Message, Ke1State, NonceLen};
+use super::{
+    Deserialize, GenerateKe1Result, GenerateKe2Result, GenerateKe3Result, KeyExchange, Serialize,
+    SerializedContext, SerializedCredentialRequest, SerializedCredentialResponse,
+    SerializedIdentifiers,
+};
+use crate::ciphersuite::{CipherSuite, KeGroup};
+use crate::errors::{InternalError, ProtocolError};
+use crate::hash::{Hash, OutputSize, ProxyHash};
+use crate::key_exchange::group::Group;
+use crate::keypair::{PrivateKey, PublicKey};
+use crate::opaque::Identifiers;
+use crate::serialization::SliceExt;
+
+/// Adapter trait that augments the `ml-kem` core traits with the metadata
+/// required by OPAQUE (e.g. fixed lengths and serialization hooks).
+pub trait KemCoreWrapper {
+    /// Public key type used for encapsulation operations.
+    type EncapsulationKey: Clone;
+
+    /// Secret key type used for decapsulation operations.
+    type DecapsulationKey: Clone + ZeroizeOnDrop;
+
+    /// Length (in bytes) of the serialized public key.
+    type EncapsulationKeyLen: ArrayLength<u8>;
+    /// Length (in bytes) of the serialized secret key.
+    type DecapsulationKeyLen: ArrayLength<u8>;
+    /// Length (in bytes) of the encapsulated ciphertext.
+    type CiphertextLen: ArrayLength<u8>;
+    /// Length (in bytes) of the shared secret output by the KEM.
+    type SharedSecretLen: ArrayLength<u8>;
+
+    /// Generates a fresh KEM key pair.
+    fn generate<R: RngCore + CryptoRng>(
+        rng: &mut R,
+    ) -> Result<(Self::DecapsulationKey, Self::EncapsulationKey), ProtocolError>;
+
+    /// Serializes the public encapsulation key.
+    fn serialize_encapsulation_key(
+        key: &Self::EncapsulationKey,
+    ) -> GenericArray<u8, Self::EncapsulationKeyLen>;
+
+    /// Deserializes the public encapsulation key, advancing the input slice.
+    fn deserialize_encapsulation_key(
+        input: &mut &[u8],
+    ) -> Result<Self::EncapsulationKey, ProtocolError>;
+
+    /// Serializes the secret decapsulation key.
+    fn serialize_decapsulation_key(
+        key: &Self::DecapsulationKey,
+    ) -> GenericArray<u8, Self::DecapsulationKeyLen>;
+
+    /// Deserializes the secret decapsulation key, advancing the input slice.
+    fn deserialize_decapsulation_key(
+        input: &mut &[u8],
+    ) -> Result<Self::DecapsulationKey, ProtocolError>;
+
+    /// Encapsulates to the given public key, returning the ciphertext and
+    /// shared secret.
+    #[allow(clippy::type_complexity)]
+    fn encapsulate<R: RngCore + CryptoRng>(
+        key: &Self::EncapsulationKey,
+        rng: &mut R,
+    ) -> Result<
+        (
+            GenericArray<u8, Self::CiphertextLen>,
+            GenericArray<u8, Self::SharedSecretLen>,
+        ),
+        ProtocolError,
+    >;
+
+    /// Decapsulates the shared secret from the provided ciphertext.
+    fn decapsulate(
+        key: &Self::DecapsulationKey,
+        encapsulated_key: &GenericArray<u8, Self::CiphertextLen>,
+    ) -> Result<GenericArray<u8, Self::SharedSecretLen>, ProtocolError>;
+}
+
+type RcEncapsulationKeyLen<K> = <<K as KemCore>::EncapsulationKey as EncodedSizeUser>::EncodedSize;
+type RcDecapsulationKeyLen<K> = <<K as KemCore>::DecapsulationKey as EncodedSizeUser>::EncodedSize;
+type RcCiphertextLen<K> = <K as KemCore>::CiphertextSize;
+type RcSharedSecretLen<K> = <K as KemCore>::SharedKeySize;
+
+impl<K> KemCoreWrapper for K
+where
+    K: KemCore,
+    K::EncapsulationKey: Encapsulate<MlKemCiphertext<K>, MlKemSharedKey<K>> + Clone,
+    K::DecapsulationKey: Decapsulate<MlKemCiphertext<K>, MlKemSharedKey<K>> + Clone + ZeroizeOnDrop,
+    RcEncapsulationKeyLen<K>: ArrayLength<u8>,
+    RcDecapsulationKeyLen<K>: ArrayLength<u8>,
+    RcCiphertextLen<K>: ArrayLength<u8>,
+    RcSharedSecretLen<K>: ArrayLength<u8>,
+{
+    type EncapsulationKey = K::EncapsulationKey;
+    type DecapsulationKey = K::DecapsulationKey;
+    type EncapsulationKeyLen = RcEncapsulationKeyLen<K>;
+    type DecapsulationKeyLen = RcDecapsulationKeyLen<K>;
+    type CiphertextLen = RcCiphertextLen<K>;
+    type SharedSecretLen = RcSharedSecretLen<K>;
+
+    fn generate<R: RngCore + CryptoRng>(
+        rng: &mut R,
+    ) -> Result<(Self::DecapsulationKey, Self::EncapsulationKey), ProtocolError> {
+        Ok(K::generate(rng))
+    }
+
+    fn serialize_encapsulation_key(
+        key: &Self::EncapsulationKey,
+    ) -> GenericArray<u8, Self::EncapsulationKeyLen> {
+        GenericArray::clone_from_slice(key.as_bytes().as_slice())
+    }
+
+    fn deserialize_encapsulation_key(
+        input: &mut &[u8],
+    ) -> Result<Self::EncapsulationKey, ProtocolError> {
+        let bytes: GenericArray<u8, RcEncapsulationKeyLen<K>> =
+            input.take_array("kem encapsulation key")?;
+        let encoded = Encoded::<K::EncapsulationKey>::try_from(bytes.as_slice())
+            .map_err(|_| ProtocolError::SerializationError)?;
+        Ok(K::EncapsulationKey::from_bytes(&encoded))
+    }
+
+    fn serialize_decapsulation_key(
+        key: &Self::DecapsulationKey,
+    ) -> GenericArray<u8, Self::DecapsulationKeyLen> {
+        GenericArray::clone_from_slice(key.as_bytes().as_slice())
+    }
+
+    fn deserialize_decapsulation_key(
+        input: &mut &[u8],
+    ) -> Result<Self::DecapsulationKey, ProtocolError> {
+        let bytes: GenericArray<u8, RcDecapsulationKeyLen<K>> =
+            input.take_array("kem decapsulation key")?;
+        let encoded = Encoded::<K::DecapsulationKey>::try_from(bytes.as_slice())
+            .map_err(|_| ProtocolError::SerializationError)?;
+        Ok(K::DecapsulationKey::from_bytes(&encoded))
+    }
+
+    fn encapsulate<R: RngCore + CryptoRng>(
+        key: &Self::EncapsulationKey,
+        rng: &mut R,
+    ) -> Result<
+        (
+            GenericArray<u8, Self::CiphertextLen>,
+            GenericArray<u8, Self::SharedSecretLen>,
+        ),
+        ProtocolError,
+    > {
+        key.encapsulate(rng)
+            .map(|(ciphertext, shared)| {
+                (
+                    GenericArray::clone_from_slice(ciphertext.as_slice()),
+                    GenericArray::clone_from_slice(shared.as_slice()),
+                )
+            })
+            .map_err(|_| ProtocolError::LibraryError(InternalError::KemError))
+    }
+
+    fn decapsulate(
+        key: &Self::DecapsulationKey,
+        encapsulated_key: &GenericArray<u8, Self::CiphertextLen>,
+    ) -> Result<GenericArray<u8, Self::SharedSecretLen>, ProtocolError> {
+        let ciphertext = MlKemCiphertext::<K>::try_from(encapsulated_key.as_slice())
+            .map_err(|_| ProtocolError::SerializationError)?;
+        key.decapsulate(&ciphertext)
+            .map(|shared| GenericArray::clone_from_slice(shared.as_slice()))
+            .map_err(|_| ProtocolError::LibraryError(InternalError::KemError))
+    }
+}
+/// Triple Diffie-Hellman-style key exchange that offloads the second hop to a
+/// generic KEM.
+#[derive(Clone, Debug)]
+pub struct TripleDhKem<G, H, K>(PhantomData<(G, H, K)>);
+
+/// Client state combining the classic `TripleDH` state with a KEM secret key.
+#[cfg_attr(
+    feature = "serde",
+    derive(serde::Deserialize, serde::Serialize),
+    serde(bound(
+        deserialize = "Ke1State<G>: serde::Deserialize<'de>, K::DecapsulationKey: \
+                       serde::Deserialize<'de>",
+        serialize = "Ke1State<G>: serde::Serialize, K::DecapsulationKey: serde::Serialize",
+    ))
+)]
+#[derive_where(Clone, ZeroizeOnDrop)]
+#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; Ke1State<G>, K::DecapsulationKey)]
+pub struct KemKe1State<G: Group, K: KemCoreWrapper> {
+    dh_state: Ke1State<G>,
+    kem_decapsulation_key: K::DecapsulationKey,
+}
+
+/// Client message including the ephemeral Diffie-Hellman component alongside a
+/// serialized KEM public key.
+#[cfg_attr(
+    feature = "serde",
+    derive(serde::Deserialize, serde::Serialize),
+    serde(bound(
+        deserialize = "Ke1Message<G>: serde::Deserialize<'de>",
+        serialize = "Ke1Message<G>: serde::Serialize",
+    ))
+)]
+#[derive_where(Clone, ZeroizeOnDrop)]
+#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; Ke1Message<G>)]
+pub struct KemKe1Message<G: Group, K: KemCoreWrapper> {
+    dh_message: Ke1Message<G>,
+    kem_encapsulation_key: GenericArray<u8, K::EncapsulationKeyLen>,
+}
+
+/// Server state mirrors the `TripleDH` state and carries the client’s KEM
+/// public key for later use.
+#[cfg_attr(
+    feature = "serde",
+    derive(serde::Deserialize, serde::Serialize),
+    serde(bound = "")
+)]
+#[derive_where(Clone, ZeroizeOnDrop)]
+#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
+pub struct KemKe2State<K: KemCoreWrapper, H: Hash>
+where
+    H::Core: ProxyHash,
+    <H::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<H::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+{
+    base_state: super::tripledh::Ke2State<H>,
+    kem_encapsulation_key: GenericArray<u8, K::EncapsulationKeyLen>,
+    server_kem_ciphertext: GenericArray<u8, K::CiphertextLen>,
+}
+
+/// Server builder placeholder capturing the data needed to finish the KEM
+/// exchange.
+#[derive_where(Clone)]
+pub struct KemKe2Builder<G: Group, H: Hash, K: KemCoreWrapper>
+where
+    H::Core: ProxyHash,
+    <H::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<H::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+{
+    server_nonce: GenericArray<u8, NonceLen>,
+    transcript_hasher: H,
+    client_e_pk: PublicKey<G>,
+    server_e_pk: PublicKey<G>,
+    shared_secret_1: GenericArray<u8, G::PkLen>,
+    shared_secret_3: GenericArray<u8, G::PkLen>,
+    kem_encapsulation_key: GenericArray<u8, K::EncapsulationKeyLen>,
+    kem_ciphertext: GenericArray<u8, K::CiphertextLen>,
+    kem_shared_secret: GenericArray<u8, K::SharedSecretLen>,
+}
+
+/// Server message bundles the `TripleDH` payload with the KEM encapsulation.
+#[cfg_attr(
+    feature = "serde",
+    derive(serde::Deserialize, serde::Serialize),
+    serde(bound(
+        deserialize = "super::tripledh::Ke2Message<G, H>: serde::Deserialize<'de>",
+        serialize = "super::tripledh::Ke2Message<G, H>: serde::Serialize",
+    ))
+)]
+#[derive_where(Clone, ZeroizeOnDrop)]
+#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; super::tripledh::Ke2Message<G, H>)]
+pub struct KemKe2Message<G: Group, H: Hash, K: KemCoreWrapper>
+where
+    H::Core: ProxyHash,
+    <H::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<H::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+{
+    dh_message: super::tripledh::Ke2Message<G, H>,
+    kem_ciphertext: GenericArray<u8, K::CiphertextLen>,
+}
+
+/// Third message remains the same as `TripleDH`.
+pub type KemKe3Message<H> = super::tripledh::Ke3Message<H>;
+
+impl<G, H, K> Drop for KemKe2Builder<G, H, K>
+where
+    G: Group,
+    H: Hash,
+    H::Core: ProxyHash,
+    <H::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<H::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+    K: KemCoreWrapper,
+{
+    fn drop(&mut self) {
+        self.server_nonce.zeroize();
+        self.transcript_hasher.reset();
+        self.shared_secret_1.zeroize();
+        self.shared_secret_3.zeroize();
+        self.kem_shared_secret.zeroize();
+        self.kem_ciphertext.zeroize();
+    }
+}
+
+impl<G, H, K> ZeroizeOnDrop for KemKe2Builder<G, H, K>
+where
+    G: Group,
+    H: Hash,
+    H::Core: ProxyHash,
+    <H::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<H::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+    K: KemCoreWrapper,
+{
+}
+
+impl<G, H, K> KeyExchange for TripleDhKem<G, H, K>
+where
+    G: Group + 'static,
+    G::Sk: shared::DiffieHellman<G>,
+    H: Hash,
+    H::Core: ProxyHash,
+    <H::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<H::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+    K: KemCoreWrapper,
+    NonceLen: Add<K::EncapsulationKeyLen>,
+    Sum<NonceLen, K::EncapsulationKeyLen>: ArrayLength<u8>,
+{
+    type Group = G;
+    type Hash = H;
+
+    type KE1State = KemKe1State<G, K>;
+    type KE2State<CS: CipherSuite> = KemKe2State<K, H>;
+    type KE1Message = KemKe1Message<G, K>;
+    type KE2Builder<'a, CS: CipherSuite<KeyExchange = Self>> = KemKe2Builder<G, H, K>;
+    type KE2BuilderData<'a, CS: 'static + CipherSuite> = (
+        &'a PublicKey<G>,
+        &'a GenericArray<u8, K::EncapsulationKeyLen>,
+    );
+    type KE2BuilderInput<CS: CipherSuite> = GenericArray<u8, G::PkLen>;
+    type KE2Message = KemKe2Message<G, H, K>;
+    type KE3Message = KemKe3Message<H>;
+
+    fn generate_ke1<R: RngCore + CryptoRng>(
+        rng: &mut R,
+    ) -> Result<GenerateKe1Result<Self>, ProtocolError> {
+        let base = super::tripledh::TripleDh::<G, H>::generate_ke1(rng)?;
+        let (kem_secret, kem_public) = K::generate(rng)?;
+        let kem_encapsulation_key = K::serialize_encapsulation_key(&kem_public);
+
+        Ok(GenerateKe1Result {
+            state: KemKe1State {
+                dh_state: base.state,
+                kem_decapsulation_key: kem_secret,
+            },
+            message: KemKe1Message {
+                dh_message: base.message,
+                kem_encapsulation_key,
+            },
+        })
+    }
+
+    fn ke2_builder<'a, CS: CipherSuite<KeyExchange = Self>, R: RngCore + CryptoRng>(
+        rng: &mut R,
+        credential_request: SerializedCredentialRequest<CS>,
+        ke1_message: Self::KE1Message,
+        credential_response: SerializedCredentialResponse<CS>,
+        client_s_pk: PublicKey<G>,
+        identifiers: SerializedIdentifiers<'_, KeGroup<CS>>,
+        context: SerializedContext<'a>,
+    ) -> Result<Self::KE2Builder<'a, CS>, ProtocolError> {
+        let shared::Ke2BuilderCommon {
+            server_nonce,
+            transcript_hasher,
+            client_e_pk,
+            server_e_pk,
+            shared_secret_1,
+            shared_secret_3,
+        } = shared::ke2_builder_common::<G, H, CS, R>(
+            rng,
+            credential_request,
+            ke1_message.dh_message.clone(),
+            credential_response,
+            client_s_pk,
+            identifiers,
+            context,
+        )?;
+
+        let mut kem_bytes_slice: &[u8] = ke1_message.kem_encapsulation_key.as_slice();
+        let encapsulation_key = K::deserialize_encapsulation_key(&mut kem_bytes_slice)?;
+        let (kem_ciphertext, kem_shared_secret) = K::encapsulate(&encapsulation_key, rng)?;
+
+        let mut transcript_hasher = transcript_hasher;
+        transcript_hasher.update(ke1_message.kem_encapsulation_key.as_slice());
+        transcript_hasher.update(kem_ciphertext.as_slice());
+
+        Ok(KemKe2Builder {
+            server_nonce,
+            transcript_hasher,
+            client_e_pk,
+            server_e_pk,
+            shared_secret_1,
+            shared_secret_3,
+            kem_encapsulation_key: ke1_message.kem_encapsulation_key.clone(),
+            kem_ciphertext,
+            kem_shared_secret,
+        })
+    }
+
+    fn ke2_builder_data<'a, CS: 'static + CipherSuite<KeyExchange = Self>>(
+        builder: &'a Self::KE2Builder<'_, CS>,
+    ) -> Self::KE2BuilderData<'a, CS> {
+        (&builder.client_e_pk, &builder.kem_encapsulation_key)
+    }
+
+    fn generate_ke2_input<CS: CipherSuite<KeyExchange = Self>, R: CryptoRng + RngCore>(
+        builder: &Self::KE2Builder<'_, CS>,
+        _: &mut R,
+        server_s_sk: &PrivateKey<G>,
+    ) -> Self::KE2BuilderInput<CS> {
+        server_s_sk.ke_diffie_hellman(&builder.client_e_pk)
+    }
+
+    fn build_ke2<CS: CipherSuite<KeyExchange = Self>>(
+        mut builder: Self::KE2Builder<'_, CS>,
+        shared_secret_2: Self::KE2BuilderInput<CS>,
+    ) -> Result<GenerateKe2Result<CS>, ProtocolError> {
+        let transcript_digest = builder.transcript_hasher.clone().finalize();
+        let derived_keys = shared::derive_keys::<H>(
+            [
+                builder.shared_secret_1.as_slice(),
+                shared_secret_2.as_slice(),
+                builder.shared_secret_3.as_slice(),
+                builder.kem_shared_secret.as_slice(),
+            ]
+            .into_iter(),
+            &transcript_digest,
+        )?;
+
+        let (mac, expected_mac) = shared::compute_ke2_macs(
+            &mut builder.transcript_hasher,
+            &derived_keys,
+            &transcript_digest,
+        )?;
+
+        Ok(GenerateKe2Result {
+            state: KemKe2State {
+                base_state: super::tripledh::Ke2State {
+                    session_key: derived_keys.session_key.clone(),
+                    expected_mac,
+                },
+                kem_encapsulation_key: builder.kem_encapsulation_key.clone(),
+                server_kem_ciphertext: builder.kem_ciphertext.clone(),
+            },
+            message: KemKe2Message {
+                dh_message: super::tripledh::Ke2Message {
+                    server_nonce: builder.server_nonce,
+                    server_e_pk: builder.server_e_pk.clone(),
+                    mac,
+                },
+                kem_ciphertext: builder.kem_ciphertext.clone(),
+            },
+            #[cfg(test)]
+            handshake_secret: derived_keys.handshake_secret,
+            #[cfg(test)]
+            km2: derived_keys.km2,
+        })
+    }
+
+    fn generate_ke3<CS: CipherSuite<KeyExchange = Self>, R: CryptoRng + RngCore>(
+        _rng: &mut R,
+        credential_request: SerializedCredentialRequest<CS>,
+        ke1_message: Self::KE1Message,
+        credential_response: SerializedCredentialResponse<CS>,
+        ke1_state: &Self::KE1State,
+        ke2_message: Self::KE2Message,
+        server_s_pk: PublicKey<G>,
+        client_s_sk: PrivateKey<G>,
+        identifiers: SerializedIdentifiers<'_, KeGroup<CS>>,
+        context: SerializedContext<'_>,
+    ) -> Result<GenerateKe3Result<Self>, ProtocolError> {
+        let mut transcript_hasher = shared::transcript(
+            &context,
+            &identifiers,
+            &credential_request,
+            &ke1_message.dh_message.to_iter(),
+            &credential_response,
+            ke2_message.dh_message.server_nonce,
+            &ke2_message.dh_message.server_e_pk.serialize(),
+        );
+        transcript_hasher.update(ke1_message.kem_encapsulation_key.as_slice());
+        transcript_hasher.update(ke2_message.kem_ciphertext.as_slice());
+
+        let shared_secret_1 = ke1_state
+            .dh_state
+            .client_e_sk
+            .ke_diffie_hellman(&ke2_message.dh_message.server_e_pk);
+        let shared_secret_2 = ke1_state
+            .dh_state
+            .client_e_sk
+            .ke_diffie_hellman(&server_s_pk);
+        let shared_secret_3 = client_s_sk.ke_diffie_hellman(&ke2_message.dh_message.server_e_pk);
+        let kem_shared_secret = K::decapsulate(
+            &ke1_state.kem_decapsulation_key,
+            &ke2_message.kem_ciphertext,
+        )?;
+
+        let (derived_keys, client_mac) = shared::finalize_ke3_transcript(
+            &mut transcript_hasher,
+            [
+                shared_secret_1.as_slice(),
+                shared_secret_2.as_slice(),
+                shared_secret_3.as_slice(),
+                kem_shared_secret.as_slice(),
+            ]
+            .into_iter(),
+            &ke2_message.dh_message.mac,
+        )?;
+
+        Ok(GenerateKe3Result {
+            session_key: derived_keys.session_key,
+            message: super::tripledh::Ke3Message { mac: client_mac },
+            #[cfg(test)]
+            handshake_secret: derived_keys.handshake_secret,
+            #[cfg(test)]
+            km3: derived_keys.km3,
+        })
+    }
+
+    fn finish_ke<CS: CipherSuite>(
+        ke2_state: &Self::KE2State<CS>,
+        ke3_message: Self::KE3Message,
+        _identifiers: Identifiers<'_>,
+        _context: SerializedContext<'_>,
+    ) -> Result<Output<Self::Hash>, ProtocolError> {
+        CtOption::new(
+            ke2_state.base_state.session_key.clone(),
+            ke2_state.base_state.expected_mac.ct_eq(&ke3_message.mac),
+        )
+        .into_option()
+        .ok_or(ProtocolError::InvalidLoginError)
+    }
+}
+
+/// Serialization logic will be implemented once the concrete KEM wiring is in
+/// place.
+impl<G: Group, K: KemCoreWrapper> Deserialize for KemKe1State<G, K> {
+    fn deserialize_take(input: &mut &[u8]) -> Result<Self, ProtocolError> {
+        Ok(Self {
+            dh_state: Ke1State::<G>::deserialize_take(input)?,
+            kem_decapsulation_key: K::deserialize_decapsulation_key(input)?,
+        })
+    }
+}
+
+impl<G: Group, K: KemCoreWrapper> Serialize for KemKe1State<G, K>
+where
+    Ke1State<G>: Serialize,
+    <Ke1State<G> as Serialize>::Len: Add<K::DecapsulationKeyLen>,
+    Sum<<Ke1State<G> as Serialize>::Len, K::DecapsulationKeyLen>: ArrayLength<u8>,
+{
+    type Len = Sum<<Ke1State<G> as Serialize>::Len, K::DecapsulationKeyLen>;
+
+    fn serialize(&self) -> GenericArray<u8, Self::Len> {
+        self.dh_state
+            .serialize()
+            .concat(K::serialize_decapsulation_key(&self.kem_decapsulation_key))
+    }
+}
+
+impl<G: Group, K: KemCoreWrapper> Deserialize for KemKe1Message<G, K> {
+    fn deserialize_take(input: &mut &[u8]) -> Result<Self, ProtocolError> {
+        Ok(Self {
+            dh_message: Ke1Message::<G>::deserialize_take(input)?,
+            kem_encapsulation_key: input.take_array("kem encapsulation key")?,
+        })
+    }
+}
+
+impl<G: Group, K: KemCoreWrapper> Serialize for KemKe1Message<G, K>
+where
+    Ke1Message<G>: Serialize,
+    <Ke1Message<G> as Serialize>::Len: Add<K::EncapsulationKeyLen>,
+    Sum<<Ke1Message<G> as Serialize>::Len, K::EncapsulationKeyLen>: ArrayLength<u8>,
+{
+    type Len = Sum<<Ke1Message<G> as Serialize>::Len, K::EncapsulationKeyLen>;
+
+    fn serialize(&self) -> GenericArray<u8, Self::Len> {
+        self.dh_message
+            .serialize()
+            .concat(self.kem_encapsulation_key.clone())
+    }
+}
+
+impl<K: KemCoreWrapper, H: Hash> Deserialize for KemKe2State<K, H>
+where
+    H::Core: ProxyHash,
+    <H::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<H::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+{
+    fn deserialize_take(input: &mut &[u8]) -> Result<Self, ProtocolError> {
+        Ok(Self {
+            base_state: super::tripledh::Ke2State::<H>::deserialize_take(input)?,
+            kem_encapsulation_key: input.take_array("kem encapsulation key")?,
+            server_kem_ciphertext: input.take_array("kem ciphertext")?,
+        })
+    }
+}
+
+impl<K: KemCoreWrapper, H: Hash> Serialize for KemKe2State<K, H>
+where
+    H::Core: ProxyHash,
+    <H::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<H::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+    super::tripledh::Ke2State<H>: Serialize,
+    <super::tripledh::Ke2State<H> as Serialize>::Len: Add<K::EncapsulationKeyLen>,
+    Sum<<super::tripledh::Ke2State<H> as Serialize>::Len, K::EncapsulationKeyLen>:
+        ArrayLength<u8> + Add<K::CiphertextLen>,
+    Sum<
+        Sum<<super::tripledh::Ke2State<H> as Serialize>::Len, K::EncapsulationKeyLen>,
+        K::CiphertextLen,
+    >: ArrayLength<u8>,
+{
+    type Len = Sum<
+        Sum<<super::tripledh::Ke2State<H> as Serialize>::Len, K::EncapsulationKeyLen>,
+        K::CiphertextLen,
+    >;
+
+    fn serialize(&self) -> GenericArray<u8, Self::Len> {
+        self.base_state
+            .serialize()
+            .concat(self.kem_encapsulation_key.clone())
+            .concat(self.server_kem_ciphertext.clone())
+    }
+}
+
+impl<G: Group, H: Hash, K: KemCoreWrapper> Deserialize for KemKe2Message<G, H, K>
+where
+    H::Core: ProxyHash,
+    <H::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<H::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+{
+    fn deserialize_take(input: &mut &[u8]) -> Result<Self, ProtocolError> {
+        Ok(Self {
+            dh_message: super::tripledh::Ke2Message::<G, H>::deserialize_take(input)?,
+            kem_ciphertext: input.take_array("kem ciphertext")?,
+        })
+    }
+}
+
+impl<G: Group, H: Hash, K: KemCoreWrapper> Serialize for KemKe2Message<G, H, K>
+where
+    H::Core: ProxyHash,
+    <H::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<H::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+    NonceLen: Add<G::PkLen>,
+    Sum<NonceLen, G::PkLen>: ArrayLength<u8> + Add<OutputSize<H>>,
+    Sum<Sum<NonceLen, G::PkLen>, OutputSize<H>>: ArrayLength<u8>,
+    super::tripledh::Ke2Message<G, H>: Serialize,
+    <super::tripledh::Ke2Message<G, H> as Serialize>::Len: Add<K::CiphertextLen>,
+    <<super::tripledh::Ke2Message<G, H> as Serialize>::Len as Add<K::CiphertextLen>>::Output:
+        ArrayLength<u8>,
+{
+    type Len = Sum<<super::tripledh::Ke2Message<G, H> as Serialize>::Len, K::CiphertextLen>;
+
+    fn serialize(&self) -> GenericArray<u8, Self::Len> {
+        self.dh_message
+            .serialize()
+            .concat(self.kem_ciphertext.clone())
+    }
+}

src/lib.rs

@@ -706,6 +706,38 @@
 //! # Ok::<(), ProtocolError>(())
 //! ```
 //!
+//! ## `TripleDhKem` Key Exchange
+//!
+//! `TripleDhKem` extends the default [`TripleDh`] handshake by mixing a
+//! KEM shared secret into the transcript alongside the three Diffie-Hellman
+//! products. This hybrid exchange offers a post-quantum (PQ) upgrade path while
+//! preserving backwards-compatible session keys and transcript bindings.
+//!
+//! Note that this key exchange does not provide full PQ security
+//! for OPAQUE, as the OPRF used in the protocol is still not PQ-secure.
+//! Morever, we only use ephemeral KEM keys to provide confidentiality,
+//! implicitly delegating the authentication guarantees to the static (and
+//! classical) Diffie-Hellman keys. This is intended as a stopgap measure to
+//! provide some level of confidentiality against passive quantum attackers (as
+//! opposed to active ones). See [RFC 9807, Appendix B](https://www.rfc-editor.org/rfc/rfc9807.html#appendix-B)
+//! for a more detailed discussion of the security guarantees.
+//!
+//! This can be enabled with the `kem` feature to compile the integration with
+//! the [`ml-kem`](https://docs.rs/ml-kem/latest/ml_kem/) crate and instantiating
+//! the ciphersuite as follows:
+//!
+//! ```ignore
+//! use opaque_ke::CipherSuite;
+//!
+//! struct KemSuite;
+//!
+//! impl CipherSuite for KemSuite {
+//!     type OprfCs = opaque_ke::Ristretto255;
+//!     type KeyExchange = opaque_ke::TripleDhKem<opaque_ke::Ristretto255, sha2::Sha512, opaque_ke::ml_kem::MlKem768>;
+//!     type Ksf = opaque_ke::ksf::Identity;
+//! }
+//! ```
+//!
 //! ## Custom Identifiers
 //!
 //! Typically when applications use OPAQUE to authenticate a client to a server,
@@ -1314,6 +1346,10 @@
 //! - The `ed25519` feature enables using [`Ed25519`]s with [`PureEddsa`] and
 //!   [`HashEddsa`] for [`SigmaI`]s signature algorithm.
 //!
+//! - The `kem` feature enables the [`TripleDhKem`] key exchange, adding support
+//!   for KEM-based handshakes backed by the `ml-kem` crate. Disabling the
+//!   feature removes those types and their associated tests from the build.
+//!
 //! [`alloc`]: https://doc.rust-lang.org/alloc
 //! [curve25519-dalek]: https://docs.rs/curve25519-dalek/4/curve25519_dalek/index.html#backends
 
@@ -1362,6 +1398,8 @@ pub use crate::key_exchange::sigma_i::ecdsa::Ecdsa;
 pub use crate::key_exchange::sigma_i::hash_eddsa::HashEddsa;
 pub use crate::key_exchange::sigma_i::pure_eddsa::PureEddsa;
 pub use crate::key_exchange::tripledh::TripleDh;
+#[cfg(feature = "kem")]
+pub use crate::key_exchange::tripledh_kem::TripleDhKem;
 pub use crate::messages::{
     CredentialFinalization, CredentialFinalizationLen, CredentialRequest, CredentialRequestLen,
     CredentialResponse, CredentialResponseLen, RegistrationRequest, RegistrationRequestLen,

src/tests/full_test.rs

@@ -16,6 +16,8 @@ use std::{format, println, vec};
 use digest::Output;
 use generic_array::typenum::{Sum, Unsigned};
 use generic_array::{ArrayLength, GenericArray};
+#[cfg(feature = "kem")]
+use ml_kem::MlKem768;
 use rand::SeedableRng;
 use rand::rngs::OsRng;
 use rand_chacha::ChaCha20Rng;
@@ -29,6 +31,8 @@ use crate::errors::*;
 use crate::hash::OutputSize;
 use crate::key_exchange::group::Group;
 use crate::key_exchange::shared::NonceLen;
+#[cfg(feature = "kem")]
+use crate::key_exchange::tripledh_kem::TripleDhKem;
 use crate::key_exchange::{
     Deserialize, Ke1MessageLen, Ke1StateLen, Ke2MessageLen, KeyExchange, Serialize,
 };
@@ -156,6 +160,27 @@ macro_rules! triple_dh_ciphersuites {
     };
 }
 
+#[cfg(feature = "kem")]
+macro_rules! triple_dh_kem_ciphersuites {
+    ($macro:ident!$par:tt) => {
+        oprf_ciphersuites!(
+            $macro!$par => [
+                #[cfg(feature = "ristretto255")] [
+                    TripleDhKemRistretto255,
+                    TripleDhKem<crate::Ristretto255, sha2::Sha512, MlKem768>,
+                ],
+                [TripleDhKemP256, TripleDhKem<p256::NistP256, sha2::Sha256, MlKem768>],
+                [TripleDhKemP384, TripleDhKem<p384::NistP384, sha2::Sha384, MlKem768>],
+                [TripleDhKemP521, TripleDhKem<p521::NistP521, sha2::Sha512, MlKem768>],
+                #[cfg(feature = "curve25519")] [
+                    TripleDhKemCurve25519,
+                    TripleDhKem<crate::Curve25519, sha2::Sha512, MlKem768>,
+                ],
+            ]
+        );
+    };
+}
+
 macro_rules! sigma_i_ciphersuites {
     ($macro:ident!$par:tt) => {
         sigma_i_ciphersuites!(
@@ -192,6 +217,8 @@ macro_rules! sigma_i_ciphersuites {
 }
 
 triple_dh_ciphersuites!(ciphersuite_types!());
+#[cfg(feature = "kem")]
+triple_dh_kem_ciphersuites!(ciphersuite_types!());
 sigma_i_ciphersuites!(ciphersuite_types!());
 
 pub struct TestVectorParameters {
@@ -741,6 +768,8 @@ fn generate_test_vectors() -> Result<(), ProtocolError> {
     );
 
     triple_dh_ciphersuites!(generate!(output));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(generate!(output));
     sigma_i_ciphersuites!(generate!(output));
 
     if let Ok(path) = std::env::var("FULL_TEST_VECTORS_FILE") {
@@ -776,6 +805,8 @@ fn test_registration_request() -> Result<(), ProtocolError> {
     }
 
     triple_dh_ciphersuites!(run_all!(inner));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(run_all!(inner));
     sigma_i_ciphersuites!(run_all!(inner));
 
     Ok(())
@@ -808,6 +839,8 @@ fn test_serialization() -> Result<(), ProtocolError> {
     }
 
     triple_dh_ciphersuites!(run_all!(inner));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(run_all!(inner));
     sigma_i_ciphersuites!(run_all!(inner));
 
     Ok(())
@@ -847,6 +880,8 @@ fn test_registration_response() -> Result<(), ProtocolError> {
     }
 
     triple_dh_ciphersuites!(run_all!(inner));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(run_all!(inner));
     sigma_i_ciphersuites!(run_all!(inner));
 
     Ok(())
@@ -896,6 +931,8 @@ fn test_registration_upload() -> Result<(), ProtocolError> {
     }
 
     triple_dh_ciphersuites!(run_all!(inner));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(run_all!(inner));
     sigma_i_ciphersuites!(run_all!(inner));
 
     Ok(())
@@ -926,6 +963,8 @@ fn test_password_file() -> Result<(), ProtocolError> {
     }
 
     triple_dh_ciphersuites!(run_all!(inner));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(run_all!(inner));
     sigma_i_ciphersuites!(run_all!(inner));
 
     Ok(())
@@ -969,6 +1008,8 @@ fn test_credential_request() -> Result<(), ProtocolError> {
     }
 
     triple_dh_ciphersuites!(run_all!(inner));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(run_all!(inner));
     sigma_i_ciphersuites!(run_all!(inner));
 
     Ok(())
@@ -1037,6 +1078,8 @@ fn test_credential_response() -> Result<(), ProtocolError> {
     }
 
     triple_dh_ciphersuites!(run_all!(inner));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(run_all!(inner));
     sigma_i_ciphersuites!(run_all!(inner));
 
     Ok(())
@@ -1089,6 +1132,8 @@ fn test_credential_finalization() -> Result<(), ProtocolError> {
     }
 
     triple_dh_ciphersuites!(run_all!(inner));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(run_all!(inner));
     sigma_i_ciphersuites!(run_all!(inner));
 
     Ok(())
@@ -1124,6 +1169,8 @@ fn test_server_login_finish() -> Result<(), ProtocolError> {
     }
 
     triple_dh_ciphersuites!(run_all!(inner));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(run_all!(inner));
     sigma_i_ciphersuites!(run_all!(inner));
 
     Ok(())
@@ -1201,6 +1248,12 @@ fn test_complete_flow_success() -> Result<(), ProtocolError> {
         b"good password",
         b"good password"
     ));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(run_all!(
+        test_complete_flow,
+        b"good password",
+        b"good password"
+    ));
     sigma_i_ciphersuites!(run_all!(
         test_complete_flow,
         b"good password",
@@ -1216,6 +1269,12 @@ fn test_complete_flow_fail() -> Result<(), ProtocolError> {
         b"good password",
         b"bad password"
     ));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(run_all!(
+        test_complete_flow,
+        b"good password",
+        b"bad password"
+    ));
     sigma_i_ciphersuites!(run_all!(
         test_complete_flow,
         b"good password",
@@ -1259,6 +1318,8 @@ fn test_scalar_always_nonzero() -> Result<(), ProtocolError> {
     }
 
     triple_dh_ciphersuites!(run_all!(inner));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(run_all!(inner));
     sigma_i_ciphersuites!(run_all!(inner));
 
     Ok(())
@@ -1304,6 +1365,8 @@ fn test_reflected_value_error_registration() -> Result<(), ProtocolError> {
     }
 
     triple_dh_ciphersuites!(run_all!(inner));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(run_all!(inner));
     sigma_i_ciphersuites!(run_all!(inner));
 
     Ok(())
@@ -1364,6 +1427,8 @@ fn test_reflected_value_error_login() -> Result<(), ProtocolError> {
     }
 
     triple_dh_ciphersuites!(run_all!(inner));
+    #[cfg(feature = "kem")]
+    triple_dh_kem_ciphersuites!(run_all!(inner));
     sigma_i_ciphersuites!(run_all!(inner));
 
     Ok(())

src/tests/mod.rs

@@ -8,6 +8,7 @@
 
 mod full_test;
 #[rustfmt::skip]
+#[allow(dead_code)]
 mod full_test_vectors;
 pub mod mock_rng;
 mod parser;