From 8ae7b6f5c1fce3e96e39372c7de91347872c81d1 Mon Sep 17 00:00:00 2001 From: Tony Arcieri Date: Sun, 10 May 2026 12:34:33 -0600 Subject: [PATCH] ml-kem: use `MaybeBox` for optional heap offload Uses the `MaybeBox` type added in #309 to offload large values onto the heap when the `alloc` feature of the crate has been enabled. --- ml-kem/Cargo.toml | 2 +- ml-kem/src/decapsulation_key.rs | 42 +++++++++++++++++++++++---------- ml-kem/src/encapsulation_key.rs | 10 ++++++-- 3 files changed, 39 insertions(+), 15 deletions(-) diff --git a/ml-kem/Cargo.toml b/ml-kem/Cargo.toml index 80210ab8..e329eb51 100644 --- a/ml-kem/Cargo.toml +++ b/ml-kem/Cargo.toml @@ -16,7 +16,7 @@ keywords = ["crypto", "kyber", "lattice", "post-quantum"] exclude = ["tests/key-gen.rs", "tests/key-gen.json", "tests/encap-decap.rs", "tests/encap-decap.json"] [features] -alloc = ["pkcs8?/alloc"] +alloc = ["module-lattice/alloc", "pkcs8?/alloc"] getrandom = ["kem/getrandom"] hazmat = [] diff --git a/ml-kem/src/decapsulation_key.rs b/ml-kem/src/decapsulation_key.rs index a48f3975..47f25b50 100644 --- a/ml-kem/src/decapsulation_key.rs +++ b/ml-kem/src/decapsulation_key.rs @@ -12,7 +12,10 @@ use kem::{ Ciphertext, Decapsulate, Decapsulator, Generate, InvalidKey, Kem, KeyExport, KeyInit, KeySizeUser, }; -use module_lattice::ctutils::{CtEq, CtSelect}; +use module_lattice::{ + MaybeBox, + ctutils::{CtEq, CtSelect}, +}; use rand_core::{TryCryptoRng, TryRng}; #[cfg(feature = "zeroize")] @@ -25,10 +28,17 @@ pub struct DecapsulationKey

where P: KemParams, { - dk_pke: DecryptionKey

, + /// Decryption key. + dk_pke: MaybeBox>, + + /// Associated encapsulation key. ek: EncapsulationKey

, - d: Option, - z: B32, + + /// Seed this key was initialized from. + d: Option>, + + /// Random string used during the implicit rejection process. + z: MaybeBox, } impl

DecapsulationKey

@@ -45,15 +55,15 @@ where /// Initialize a [`DecapsulationKey`] from the serialized expanded key form. /// - /// Note that this form is deprecated in practice; prefer to use - /// [`DecapsulationKey::from_seed`]. See [`ExpandedKeyEncoding`] for more information. + /// Note that this form is deprecated in practice; use [`DecapsulationKey::from_seed`]. + /// See [`ExpandedKeyEncoding`] for more information. /// /// # Errors /// - Returns [`InvalidKey`] in the event the expanded key failed validation #[deprecated(since = "0.3.0", note = "use `DecapsulationKey::from_seed` instead")] pub fn from_expanded(enc: &ExpandedDecapsulationKey

) -> Result { let (dk_pke, ek_pke, h, z) = P::split_dk(enc); - let dk_pke = DecryptionKey::from_bytes(dk_pke); + let dk_pke = MaybeBox::new(DecryptionKey::from_bytes(dk_pke)); let ek_pke = EncryptionKey::from_bytes(ek_pke)?; let ek = EncapsulationKey::from_encryption_key(ek_pke); @@ -65,7 +75,7 @@ where dk_pke, ek, d: None, - z: z.clone(), + z: MaybeBox::new(z.clone()), }) } @@ -82,11 +92,13 @@ where /// - `Some` if the [`DecapsulationKey`] was initialized using `from_seed` or `generate`. /// - `None` if the [`DecapsulationKey`] was initialized from the expanded form. #[inline] + #[must_use] pub fn to_seed(&self) -> Option { - self.d.map(|d| d.concat(self.z)) + self.d.as_ref().map(|d| d.concat(*self.z)) } /// Get the [`EncapsulationKey`] which corresponds to this [`DecapsulationKey`]. + #[must_use] pub fn encapsulation_key(&self) -> &EncapsulationKey

{ &self.ek } @@ -107,7 +119,11 @@ where pub(crate) fn generate_deterministic(d: B32, z: B32) -> Self { let (dk_pke, ek_pke) = DecryptionKey::generate(&d); let ek = EncapsulationKey::from_encryption_key(ek_pke); - let d = Some(d); + + let dk_pke = MaybeBox::new(dk_pke); + let d = Some(MaybeBox::new(d)); + let z = MaybeBox::new(z); + Self { dk_pke, ek, d, z } } } @@ -130,7 +146,9 @@ where { fn drop(&mut self) { self.dk_pke.zeroize(); - self.d.zeroize(); + if let Some(d) = self.d.as_mut() { + d.zeroize(); + } self.z.zeroize(); } } @@ -258,6 +276,6 @@ where fn to_expanded_bytes(&self) -> ExpandedDecapsulationKey

{ let dk_pke = self.dk_pke.to_bytes(); let ek = self.ek.to_bytes(); - P::concat_dk(dk_pke, ek, self.ek.h(), self.z.clone()) + P::concat_dk(dk_pke, ek, self.ek.h(), *self.z) } } diff --git a/ml-kem/src/encapsulation_key.rs b/ml-kem/src/encapsulation_key.rs index f09e2460..be2520e2 100644 --- a/ml-kem/src/encapsulation_key.rs +++ b/ml-kem/src/encapsulation_key.rs @@ -7,6 +7,7 @@ use crate::{ }; use array::sizes::U32; use kem::{Ciphertext, Encapsulate, Generate}; +use module_lattice::MaybeBox; use rand_core::CryptoRng; /// An `EncapsulationKey` provides the ability to encapsulate a shared key so that it can only be @@ -16,7 +17,7 @@ pub struct EncapsulationKey

where P: KemParams, { - ek_pke: EncryptionKey

, + ek_pke: MaybeBox>, h: B32, } @@ -40,6 +41,7 @@ where /// Do NOT use this function unless you know what you're doing. If you fail to use all uniform /// random bytes even once, you can have catastrophic security failure. #[cfg_attr(not(feature = "hazmat"), doc(hidden))] + #[must_use] pub fn encapsulate_deterministic(&self, m: &B32) -> (Ciphertext

, SharedKey) { let (K, r) = G(&[m, &self.h]); let c = self.ek_pke.encrypt(m, &r); @@ -47,9 +49,13 @@ where } /// Convert from an `EncryptionKey`. + #[inline] pub(crate) fn from_encryption_key(ek_pke: EncryptionKey

) -> Self { let h = H(ek_pke.to_bytes()); - Self { ek_pke, h } + Self { + ek_pke: MaybeBox::new(ek_pke), + h, + } } /// Borrow the encryption key.