ml-kem: use MaybeBox for optional heap offload

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 = []

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<P>
 where
     P: KemParams,
 {
-    dk_pke: DecryptionKey<P>,
+    /// Decryption key.
+    dk_pke: MaybeBox<DecryptionKey<P>>,
+
+    /// Associated encapsulation key.
     ek: EncapsulationKey<P>,
-    d: Option<B32>,
-    z: B32,
+
+    /// Seed this key was initialized from.
+    d: Option<MaybeBox<B32>>,
+
+    /// Random string used during the implicit rejection process.
+    z: MaybeBox<B32>,
 }
 
 impl<P> DecapsulationKey<P>
@@ -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<P>) -> Result<Self, InvalidKey> {
         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<Seed> {
-        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<P> {
         &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<P> {
         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)
     }
 }

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<P>
 where
     P: KemParams,
 {
-    ek_pke: EncryptionKey<P>,
+    ek_pke: MaybeBox<EncryptionKey<P>>,
     h: B32,
 }
 
@@ -40,16 +41,21 @@ 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<P>, SharedKey) {
         let (K, r) = G(&[m, &self.h]);
         let c = self.ek_pke.encrypt(m, &r);
         (c, K)
     }
 
     /// Convert from an `EncryptionKey`.
+    #[inline]
     pub(crate) fn from_encryption_key(ek_pke: EncryptionKey<P>) -> Self {
         let h = H(ek_pke.to_bytes());
-        Self { ek_pke, h }
+        Self {
+            ek_pke: MaybeBox::new(ek_pke),
+            h,
+        }
     }
 
     /// Borrow the encryption key.