Implement hybrid EC+MLKEM groups from draft-ietf-tls-ecdhe-mlkem

Cargo.toml

@@ -28,6 +28,7 @@ ecdsa = { version = "0.17.0-rc.16", default-features = false, features = ["alloc
 ed25519-dalek = { version = "3.0.0-pre.6", default-features = false, features = ["pkcs8"] }
 getrandom = { version = "0.4", default-features = false, features = ["sys_rng"] }
 hmac = { version = "0.13.0-rc.5", default-features = false }
+ml-kem = { version = "0.3.0-rc.0", default-features = false, features = ["getrandom"] }
 p256 = { version = "0.14.0-rc.7", default-features = false, features = ["pem", "ecdsa", "ecdh"] }
 p384 = { version = "0.14.0-rc.7", default-features = false, features = ["pem", "ecdsa", "ecdh"] }
 paste = { version = "1", default-features = false }
@@ -39,6 +40,7 @@ sec1 = { version = "0.8.0-rc.13", default-features = false }
 sha2 = { version = "0.11.0-rc.5", default-features = false }
 signature = { version = "3.0.0-rc.10", default-features = false }
 x25519-dalek = { version = "3.0.0-pre.6", default-features = false }
+zeroize = { version = "1.8", default-features = false, optional = true }
 
 [features]
 default = ["std", "tls12", "zeroize"]
@@ -52,4 +54,4 @@ tls12 = ["rustls/tls12"]
 std = ["alloc", "pki-types/std", "rustls/std"]
 # TODO: go through all of these to ensure to_vec etc. impls are exposed
 alloc = ["pki-types/alloc", "aead/alloc", "ed25519-dalek/alloc"]
-zeroize = ["ed25519-dalek/zeroize", "x25519-dalek/zeroize"]
+zeroize = ["ed25519-dalek/zeroize", "x25519-dalek/zeroize", "ml-kem/zeroize", "dep:zeroize"]

src/kx.rs

@@ -7,6 +7,11 @@ use getrandom::rand_core::UnwrapErr;
 use paste::paste;
 use rustls::crypto;
 
+mod hybrid;
+mod mlkem;
+
+pub use hybrid::{SecP256r1MLKEM768, SecP384r1MLKEM1024, X25519MLKEM768};
+
 #[derive(Debug)]
 pub struct X25519;
 
@@ -109,4 +114,56 @@ macro_rules! impl_kx {
 impl_kx! {SecP256R1, rustls::NamedGroup::secp256r1, p256::ecdh::EphemeralSecret, p256::PublicKey}
 impl_kx! {SecP384R1, rustls::NamedGroup::secp384r1, p384::ecdh::EphemeralSecret, p384::PublicKey}
 
-pub const ALL_KX_GROUPS: &[&dyn SupportedKxGroup] = &[&X25519, &SecP256R1, &SecP384R1];
+pub const ALL_KX_GROUPS: &[&dyn SupportedKxGroup] = &[
+    &X25519,
+    &SecP256R1,
+    &SecP384R1,
+    &X25519MLKEM768,
+    &SecP256r1MLKEM768,
+    &SecP384r1MLKEM1024,
+];
+
+#[cfg(test)]
+mod test {
+
+    // Make sure every key exchange algorithm can round-trip with itself.
+    #[test]
+    fn kx_roundtrip() -> Result<(), rustls::Error> {
+        for kx in super::ALL_KX_GROUPS {
+            let client_state = kx.start()?;
+            let server_output = kx.start_and_complete(client_state.pub_key())?;
+            let client_output = client_state.complete(&server_output.pub_key)?;
+
+            assert_eq!(server_output.group, kx.name());
+            assert_eq!(
+                server_output.secret.secret_bytes(),
+                client_output.secret_bytes()
+            );
+        }
+        Ok(())
+    }
+
+    // Make sure that the hybrid optimization works for each key
+    // exchange that provides it.
+    #[test]
+    fn kx_hybrid_optimization() -> Result<(), rustls::Error> {
+        for kx in super::ALL_KX_GROUPS {
+            let client_state = kx.start()?;
+            if let Some((grp, client_pubkey)) = client_state.hybrid_component() {
+                let server_kx = super::ALL_KX_GROUPS
+                    .iter()
+                    .find(|g| g.name() == grp)
+                    .unwrap();
+                let server_output = server_kx.start_and_complete(client_pubkey)?;
+                let client_output =
+                    client_state.complete_hybrid_component(&server_output.pub_key)?;
+                assert_eq!(server_output.group, grp);
+                assert_eq!(
+                    server_output.secret.secret_bytes(),
+                    client_output.secret_bytes()
+                );
+            }
+        }
+        Ok(())
+    }
+}

src/kx/hybrid.rs

@@ -0,0 +1,165 @@
+//! Implement the hybrid postquantum key exchanges from
+//! https://datatracker.ietf.org/doc/draft-ietf-tls-ecdhe-mlkem/ .
+//!
+//! These key exchanges work by combining the key_exchange shares from
+//! an elliptic curve key exchange and an MLKEM key exchange, and
+//! simply concatenating them.
+//!
+//! Since all of the encodings are constant-length, concatenation and
+//! splitting is trivial.
+
+use alloc::{boxed::Box, vec::Vec};
+use crypto::SupportedKxGroup as _;
+use paste::paste;
+use rustls::{crypto, NamedGroup};
+
+use super::mlkem::{MLKEM1024, MLKEM768};
+use super::{SecP256R1, SecP384R1, X25519};
+
+const SECP256R1MLKEM768_ID: u16 = 4587;
+const X25519MLKEM768_ID: u16 = 4588;
+const SECP384R1MLKEM1024_ID: u16 = 4589;
+
+/// Make a new vector by concatenating two slices.
+///
+/// Only allocates once. (This is important, since reallocating would
+/// imply that secret data could be left on the heap by the realloc
+/// call.)
+fn concat(b1: &[u8], b2: &[u8]) -> Vec<u8> {
+    let mut v = Vec::with_capacity(b1.len() + b2.len());
+    v.extend_from_slice(b1);
+    v.extend_from_slice(b2);
+    v
+}
+
+/// Replacement for slice::split_at_checked, which is not available
+/// at the current MSRV.
+fn split_at_checked(slice: &[u8], mid: usize) -> Option<(&[u8], &[u8])> {
+    if mid <= slice.len() {
+        Some(slice.split_at(mid))
+    } else {
+        None
+    }
+}
+
+fn first<A>(tup: (A, A)) -> A {
+    tup.0
+}
+fn second<A>(tup: (A, A)) -> A {
+    tup.1
+}
+
+// Positions to split the client and server keyshare components respectively
+// in the X25519MLKEM768 handshake.
+const X25519MLKEM768_CKE_SPLIT: usize = 1184;
+const X25519MLKEM768_SKE_SPLIT: usize = 1088;
+
+// Positions to split the client and server keyshare components respectively
+// in the SecP256r1MLKEM768 handshake.
+const SECP256R1MLKEM768_CKE_SPLIT: usize = 65;
+const SECP256R1MLKEM768_SKE_SPLIT: usize = 65;
+
+// Positions to split the client and server keyshare components respectively
+// in the SecP384r1MLKEM1024 handshake.
+const SECP384R1MLKEM1024_CKE_SPLIT: usize = 97;
+const SECP384R1MLKEM1024_SKE_SPLIT: usize = 97;
+
+macro_rules! hybrid_kex {
+    ($name:ident, $kex1:ty, $kex2:ty, $kex_ec:ty, $ec_member:expr) => {
+        paste! {
+            #[derive(Debug)]
+            pub struct $name;
+
+            struct [< $name KeyExchange >] {
+                // Note: This is redundant with pub_key in kx1 and kx2.
+                pub_key: Box<[u8]>,
+                kx1: Box<dyn crypto::ActiveKeyExchange>,
+                kx2: Box<dyn crypto::ActiveKeyExchange>,
+            }
+
+            impl crypto::SupportedKxGroup for $name {
+                fn name(&self) -> NamedGroup {
+                    NamedGroup::from([< $name:upper _ID >])
+                }
+
+                fn usable_for_version(&self, version: rustls::ProtocolVersion) -> bool {
+                    version == rustls::ProtocolVersion::TLSv1_3
+                }
+
+                fn start(&self) -> Result<Box<dyn crypto::ActiveKeyExchange>, rustls::Error> {
+                    let kx1 = $kex1.start()?;
+                    let kx2 = $kex2.start()?;
+                    Ok(Box::new([< $name KeyExchange >] {
+                        pub_key: concat(kx1.pub_key(), kx2.pub_key()).into(),
+                        kx1,
+                        kx2,
+                    }))
+                }
+
+                fn start_and_complete(
+                    &self,
+                    peer: &[u8],
+                ) -> Result<crypto::CompletedKeyExchange, rustls::Error> {
+                    let (kx1_pubkey, kx2_pubkey) =
+                        split_at_checked(peer, [< $name:upper _CKE_SPLIT >])
+                        .ok_or_else(|| rustls::Error::from(rustls::PeerMisbehaved::InvalidKeyShare))?;
+                    let kx1_completed = $kex1.start_and_complete(kx1_pubkey)?;
+                    let kx2_completed = $kex2.start_and_complete(kx2_pubkey)?;
+
+                    Ok(crypto::CompletedKeyExchange {
+                        group: self.name(),
+                        pub_key: concat(&kx1_completed.pub_key, &kx2_completed.pub_key).into(),
+                        secret: concat(
+                            kx1_completed.secret.secret_bytes(),
+                            kx2_completed.secret.secret_bytes(),
+                        )
+                            .into(),
+                    })
+                }
+            }
+
+            impl crypto::ActiveKeyExchange for [< $name KeyExchange >] {
+                fn group(&self) -> NamedGroup {
+                    NamedGroup::from([< $name:upper _ID >])
+                }
+
+                fn pub_key(&self) -> &[u8] {
+                    &self.pub_key
+                }
+
+                fn complete(self: Box<Self>, peer: &[u8]) -> Result<crypto::SharedSecret, rustls::Error> {
+                    let (kx1_pubkey, kx2_pubkey) =
+                        split_at_checked(peer, [< $name:upper _SKE_SPLIT >])
+                        .ok_or_else(|| rustls::Error::from(rustls::PeerMisbehaved::InvalidKeyShare))?;
+                    let secret1 = self.kx1.complete(kx1_pubkey)?;
+                    let secret2 = self.kx2.complete(kx2_pubkey)?;
+                    Ok(concat(secret1.secret_bytes(), secret2.secret_bytes()).into())
+                }
+
+                fn hybrid_component(&self) -> Option<(NamedGroup, &[u8])> {
+                    let pk = self.pub_key.split_at([< $name:upper _CKE_SPLIT >]);
+                    let ec_pk = ($ec_member)(pk);
+                    Some((
+                        $kex_ec.name(),
+                        ec_pk,
+                    ))
+                }
+
+                fn complete_hybrid_component(
+                    self: Box<Self>,
+                    peer: &[u8],
+                ) -> Result<crypto::SharedSecret, rustls::Error> {
+                    let ec_kx = ($ec_member)((self.kx1, self.kx2));
+                    ec_kx.complete(peer)
+                }
+            }
+        }
+    }
+}
+
+// Note: The EC key appears first in the SecP* groups,
+// but (for historical reasons) appears second in X25519MLKEM768.
+
+hybrid_kex! { X25519MLKEM768, MLKEM768, X25519, X25519, second }
+hybrid_kex! { SecP256r1MLKEM768, SecP256R1, MLKEM768, SecP256R1, first }
+hybrid_kex! { SecP384r1MLKEM1024, SecP384R1, MLKEM1024, SecP384R1, first }