From 68feb8b93e5be377387fd49ecfaccbf9937da696 Mon Sep 17 00:00:00 2001 From: mmalenic Date: Mon, 11 May 2026 20:40:14 +1000 Subject: [PATCH] refactor: avoid storing whole decapsulation key for mlkem --- Cargo.lock | 130 ++++++++++++++++++++++++++++++++--------------------- Cargo.toml | 4 +- src/kex.rs | 67 ++++++++++++++------------- 3 files changed, 116 insertions(+), 85 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0bcecd8..c9a1204 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -342,7 +342,7 @@ version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cdd35008169921d80bc60d3d0ab416eecb028c4cd653352907921d95084790be" dependencies = [ - "hybrid-array 0.4.11", + "hybrid-array", ] [[package]] @@ -502,9 +502,9 @@ dependencies = [ [[package]] name = "cpubits" -version = "0.1.0" +version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ef0c543070d296ea414df2dd7625d1b24866ce206709d8a4a424f28377f5861" +checksum = "15b85f9c39137c3a891689859392b1bd49812121d0d61c9caf00d46ed5ce06ae" [[package]] name = "cpufeatures" @@ -571,9 +571,9 @@ checksum = "42a0d26b245348befa0c121944541476763dcc46ede886c88f9d12e1697d27c3" dependencies = [ "cpubits", "ctutils", - "hybrid-array 0.4.11", + "hybrid-array", "num-traits", - "rand_core 0.10.1", + "rand_core 0.10.0", "subtle", "zeroize", ] @@ -594,8 +594,8 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77727bb15fa921304124b128af125e7e3b968275d1b108b379190264f4423710" dependencies = [ - "hybrid-array 0.4.11", - "rand_core 0.10.1", + "hybrid-array", + "rand_core 0.10.0", ] [[package]] @@ -804,15 +804,15 @@ dependencies = [ [[package]] name = "ecdsa" -version = "0.17.0-rc.17" +version = "0.17.0-rc.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc4bf51f0534ed6e59a0f2f26272b64ba55c470133f8424c2adfd1c4d59d9988" +checksum = "54fb064faabbee66e1fc8e5c5a9458d4269dc2d8b638fe86a425adb2510d1a96" dependencies = [ "der 0.8.0", "digest 0.11.2", "elliptic-curve 0.14.0-rc.32", - "rfc6979 0.5.0-rc.5", - "signature 3.0.0-rc.10", + "rfc6979 0.5.0", + "signature 3.0.0", "zeroize", ] @@ -874,8 +874,8 @@ dependencies = [ "crypto-bigint 0.7.3", "crypto-common 0.2.1", "digest 0.11.2", - "hybrid-array 0.4.11", - "rand_core 0.10.1", + "hybrid-array", + "rand_core 0.10.0", "rustcrypto-ff", "rustcrypto-group", "sec1 0.8.1", @@ -1588,19 +1588,11 @@ checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4" [[package]] name = "hybrid-array" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2d35805454dc9f8662a98d6d61886ffe26bd465f5960e0e55345c70d5c0d2a9" -dependencies = [ - "typenum", -] - -[[package]] -name = "hybrid-array" -version = "0.4.11" +version = "0.4.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08d46837a0ed51fe95bd3b05de33cd64a1ee88fc797477ca48446872504507c5" +checksum = "3944cf8cf766b40e2a1a333ee5e9b563f854d5fa49d6a8ca2764e97c6eddb214" dependencies = [ + "ctutils", "subtle", "typenum", "zeroize", @@ -1691,14 +1683,24 @@ dependencies = [ "cpufeatures 0.2.12", ] +[[package]] +name = "keccak" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e24a010dd405bd7ed803e5253182815b41bf2e6a80cc3bfc066658e03a198aa" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", +] + [[package]] name = "kem" -version = "0.3.0-pre.0" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b8645470337db67b01a7f966decf7d0bafedbae74147d33e641c67a91df239f" +checksum = "01737161ba802849cfd486b5bd209d38ba4943494c249a8126005170c7621edd" dependencies = [ - "rand_core 0.6.4", - "zeroize", + "crypto-common 0.2.1", + "rand_core 0.10.0", ] [[package]] @@ -1716,7 +1718,7 @@ dependencies = [ "pico-args", "regex", "regex-syntax", - "sha3", + "sha3 0.10.8", "string_cache", "term", "unicode-xid", @@ -1834,14 +1836,27 @@ dependencies = [ [[package]] name = "ml-kem" -version = "0.2.1" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97befee0c869cb56f3118f49d0f9bb68c9e3f380dec23c1100aedc4ec3ba239a" +checksum = "5e15f3e5b957493873e396a66914e83e616b6afe335cdef7efe5c6e1216aba66" dependencies = [ - "hybrid-array 0.2.3", + "hybrid-array", "kem", - "rand_core 0.6.4", - "sha3", + "module-lattice", + "rand_core 0.10.0", + "sha3 0.11.0", + "zeroize", +] + +[[package]] +name = "module-lattice" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c61b87c9683ab7cb1c6871d261ad5479b6b10ceb52c4352aaca3b5d35a8febe" +dependencies = [ + "ctutils", + "hybrid-array", + "num-traits", "zeroize", ] @@ -1881,10 +1896,11 @@ dependencies = [ [[package]] name = "num-bigint-dig" -version = "0.8.6" +version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e661dda6640fad38e827a6d4a310ff4763082116fe217f279885c97f511bb0b7" +checksum = "dc84195820f291c7697304f3cbdadd1cb7199c0efc917ff5eafd71225c136151" dependencies = [ + "byteorder", "lazy_static", "libm", "num-integer", @@ -2024,7 +2040,7 @@ version = "0.14.0-rc.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b97e3bf0465157ae90975ff52dbeb1362ba618924878c9f74c25baa27a65f9a" dependencies = [ - "ecdsa 0.17.0-rc.17", + "ecdsa 0.17.0-rc.18", "elliptic-curve 0.14.0-rc.32", "primefield", "primeorder 0.14.0-rc.9", @@ -2269,7 +2285,7 @@ checksum = "1b52e6ee42db392378a95622b463c9740631171d1efce43fa445a569c1600cb6" dependencies = [ "crypto-bigint 0.7.3", "crypto-common 0.2.1", - "rand_core 0.10.1", + "rand_core 0.10.0", "rustcrypto-ff", "subtle", "zeroize", @@ -2364,9 +2380,9 @@ dependencies = [ [[package]] name = "rand_core" -version = "0.10.1" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" +checksum = "0c8d0fd677905edcbeedbf2edb6494d676f0e98d54d5cf9bda0b061cb8fb8aba" [[package]] name = "redox_syscall" @@ -2427,9 +2443,9 @@ dependencies = [ [[package]] name = "rfc6979" -version = "0.5.0-rc.5" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23a3127ee32baec36af75b4107082d9bd823501ec14a4e016be4b6b37faa74ae" +checksum = "5236ce872cac07e0fb3969b0cbf468c7d2f37d432f1b627dcb7b8d34563fb0c3" dependencies = [ "hmac 0.13.0", "subtle", @@ -2554,7 +2570,7 @@ version = "0.14.0-rc.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fd2a8adb347447693cd2ba0d218c4b66c62da9b0a5672b17b981e4291ec65ff6" dependencies = [ - "rand_core 0.10.1", + "rand_core 0.10.0", "subtle", ] @@ -2564,7 +2580,7 @@ version = "0.14.0-rc.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "369f9b61aa45933c062c9f6b5c3c50ab710687eca83dd3802653b140b43f85ed" dependencies = [ - "rand_core 0.10.1", + "rand_core 0.10.0", "rustcrypto-ff", "subtle", ] @@ -2627,7 +2643,7 @@ dependencies = [ "base16ct 1.0.0", "ctutils", "der 0.8.0", - "hybrid-array 0.4.11", + "hybrid-array", "subtle", "zeroize", ] @@ -2718,7 +2734,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "75872d278a8f37ef87fa0ddbda7802605cb18344497949862c0d4dcb291eba60" dependencies = [ "digest 0.10.7", - "keccak", + "keccak 0.1.6", +] + +[[package]] +name = "sha3" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be176f1a57ce4e3d31c1a166222d9768de5954f811601fb7ca06fc8203905ce1" +dependencies = [ + "digest 0.11.2", + "keccak 0.2.0", ] [[package]] @@ -2742,12 +2768,12 @@ dependencies = [ [[package]] name = "signature" -version = "3.0.0-rc.10" +version = "3.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f1880df446116126965eeec169136b2e0251dba37c6223bcc819569550edea3" +checksum = "28d567dcbaf0049cb8ac2608a76cd95ff9e4412e1899d389ee400918ca7537f5" dependencies = [ "digest 0.11.2", - "rand_core 0.10.1", + "rand_core 0.10.0", ] [[package]] @@ -2976,7 +3002,7 @@ dependencies = [ "ctr", "curve25519-dalek", "digest 0.10.7", - "ecdsa 0.17.0-rc.17", + "ecdsa 0.17.0-rc.18", "ed25519-dalek", "embedded-io", "getrandom", @@ -3250,9 +3276,9 @@ dependencies = [ [[package]] name = "typenum" -version = "1.20.0" +version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "40ce102ab67701b8526c123c1bab5cbe42d7040ccfd0f64af1a385808d2f43de" +checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" [[package]] name = "ufmt-write" diff --git a/Cargo.toml b/Cargo.toml index 238c0cf..fd03291 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -55,7 +55,7 @@ subtle = { version = "2.4", default-features = false } ed25519-dalek = { version = "2.1", default-features = false, features = ["zeroize", "rand_core"] } x25519-dalek = { version = "2.0", default-features = false, features = ["zeroize"] } curve25519-dalek = { version = "4.1", default-features = false, features = ["zeroize"] } -ml-kem = { version = "0.2.1", default-features = false, features = ["zeroize"], optional = true } +ml-kem = { version = "0.3.2", default-features = false, features = ["zeroize"], optional = true } rsa = { version = "0.9", default-features = false, optional = true, features = ["sha2"] } # TODO: getrandom feature is a workaround for missing ssh-key dependency with rsa. fixed in pending 0.6 ssh-key = { version = "0.6", default-features = false, optional = true, features = ["getrandom"] } @@ -67,7 +67,7 @@ p256 = { version = "0.14.0-rc.9", optional = true, default-features = false, fea ecdsa = { version = "0.17.0-rc.17", default-features = false, optional = true , features = ["hazmat", "algorithm"]} [features] -default = [] +default = ["mlkem"] std = ["snafu/std", "ssh-key/alloc", "larger", "mlkem"] backtrace = ["snafu/backtrace"] rsa = ["dep:rsa", "ssh-key/rsa"] diff --git a/src/kex.rs b/src/kex.rs index 108dd4e..1d8e3b5 100644 --- a/src/kex.rs +++ b/src/kex.rs @@ -15,8 +15,8 @@ use core::{fmt, marker::PhantomData}; use digest::Digest; #[cfg(feature = "mlkem")] use ml_kem::{ - kem::{Decapsulate, Encapsulate, EncapsulationKey, Kem}, - Ciphertext, EncodedSizeUser, KemCore, MlKem768, MlKem768Params, + kem::{Decapsulate, EncapsulationKey}, + Ciphertext, DecapsulationKey, Key, KeyExport, MlKem768, Seed, B32, }; use rand_core::OsRng; use sha2::Sha256; @@ -60,8 +60,6 @@ const fixed_options_hostsig: &[&str] = &[ SSH_NAME_ED25519, #[cfg(feature = "rsa")] SSH_NAME_RSA_SHA256, - #[cfg(feature = "ecdsa256")] - SSH_NAME_ECDSA256, ]; const fixed_options_cipher: &[&str] = &[SSH_NAME_CHAPOLY, SSH_NAME_AES256_CTR]; @@ -679,7 +677,7 @@ impl SharedSecret { Self::KexCurve25519(k) => k.pubkey(), #[cfg(feature = "mlkem")] Self::KexMlkemX25519(k) => { - mlkem_bytes = k.init_pubkey_arr_client(); + mlkem_bytes = k.init_pubkey_arr_client()?; &mlkem_bytes } }; @@ -914,15 +912,15 @@ impl KexCurve25519 { #[cfg(feature = "mlkem")] pub(crate) struct KexMlkemX25519 { ecdh: KexCurve25519, - // Initialised in `new()`, cleared after deriving the secret - mlkem_ours: Option< as KemCore>::DecapsulationKey>, + // 64-byte seed for deterministic mlkem key regeneration. + mlkem_seed: Option, } #[cfg(feature = "mlkem")] impl core::fmt::Debug for KexMlkemX25519 { fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result { f.debug_struct("KexMlkemX25519") - .field("ours", &if self.mlkem_ours.is_some() { "Some" } else { "None" }) + .field("seed", &if self.mlkem_seed.is_some() { "Some" } else { "None" }) .field("ecdh", &self.ecdh) .finish() } @@ -940,37 +938,40 @@ impl KexMlkemX25519 { Self::MLKEM768_CIPHERTEXT_SIZE + Self::X25519_PUBKEY_SIZE; const _CHECK0: () = assert!( - Self::MLKEM768_PUBKEY_SIZE - == size_of::>>() - ); - const _CHECK1: () = assert!( - Self::MLKEM768_CIPHERTEXT_SIZE == size_of::>() + Self::MLKEM768_PUBKEY_SIZE == size_of::>>() ); + const _CHECK1: () = + assert!(Self::MLKEM768_CIPHERTEXT_SIZE == size_of::>()); fn new() -> Result { - Ok(Self { ecdh: KexCurve25519::new()?, mlkem_ours: None }) + Ok(Self { ecdh: KexCurve25519::new()?, mlkem_seed: None }) } - /// Generates the publickey for a sent kexdhinit - fn init_pubkey_arr_client(&mut self) -> [u8; Self::PUBLICKEY_SIZE] { - debug_assert!(self.mlkem_ours.is_none()); - // TODO does this construct in-place? - let (dk, _ek) = MlKem768::generate(&mut rand_core::OsRng); + /// Generates the publickey for a sent kexdhinit.The key is re-derived from the seed during + /// decapsulation to avoid storing it in memory. + fn init_pubkey_arr_client(&mut self) -> Result<[u8; Self::PUBLICKEY_SIZE]> { + debug_assert!(self.mlkem_seed.is_none()); + + let mut seed = Seed::default(); + random::fill_random(&mut seed)?; + + let dk = DecapsulationKey::from_seed(seed); let pubkey = self.pubkey_client(dk.encapsulation_key()); - self.mlkem_ours = Some(dk); - pubkey + + self.mlkem_seed = Some(seed); + Ok(pubkey) } fn pubkey_client( &mut self, - ek: &EncapsulationKey, + ek: &EncapsulationKey, ) -> [u8; Self::PUBLICKEY_SIZE] { let mut out = [0u8; Self::PUBLICKEY_SIZE]; // Concatenate pq and ecdh. // C_INIT = C_PK2 || C_PK1. C_PK2 pq kem, C_PK1 ecdh let (pq, ec) = out.split_at_mut(Self::MLKEM768_PUBKEY_SIZE); let pq: &mut [u8; Self::MLKEM768_PUBKEY_SIZE] = pq.try_into().unwrap(); - *pq = ek.as_bytes().into(); + *pq = ek.to_bytes().into(); ec.copy_from_slice(self.ecdh.pubkey()); out } @@ -990,16 +991,17 @@ impl KexMlkemX25519 { .ok_or_else(|| error::BadKex.build())?; let ek = pq_in.try_into().map_err(|_| error::BadKex.build())?; - let ek = EncapsulationKey::::from_bytes(ek); + let ek = EncapsulationKey::::new(ek) + .map_err(|_| error::BadKex.build())?; // S_REPLY = S_CT2 || S_PK1. S_CT2 pq kem, S_PK1 ecdh let (pq, ec) = ct_out.split_at_mut(Self::MLKEM768_CIPHERTEXT_SIZE); let pq: &mut [u8; Self::MLKEM768_CIPHERTEXT_SIZE] = pq.try_into().unwrap(); - let enc = ek - .encapsulate(&mut rand_core::OsRng) - .map_err(|_| error::BadKex.build())?; - let (ct, pq_secret) = enc; - // TODO: check if this is another stack copy. + + let mut m = B32::default(); + random::fill_random(&mut m)?; + let (ct, pq_secret) = ek.encapsulate_deterministic(&m); + *pq = ct.into(); ec.copy_from_slice(self.ecdh.pubkey()); @@ -1019,8 +1021,11 @@ impl KexMlkemX25519 { let ct: &Ciphertext = pq_in.try_into().map_err(|_| error::BadKex.build())?; - let dk = self.mlkem_ours.take().trap()?; - let pq_secret = dk.decapsulate(ct).map_err(|_| error::BadKex.build())?; + + // Re-derive the DecapsulationKey from the seed + let seed = self.mlkem_seed.take().trap()?; + let dk = DecapsulationKey::from_seed(seed); + let pq_secret = dk.decapsulate(ct); let ek = dk.encapsulation_key(); let c_pk = self.pubkey_client(ek);