mlkem: Move caching of encapsulation key (#465)

* mlkem: Move caching of encapsulation key up to individual DecapsulationKey newtypes to avoid storing it in its enitrety twice in memory

* 0.17.10

Author: brycx

CHANGELOG.md

@@ -1,6 +1,6 @@
 ### 0.17.10
 
-**Date:** TBD.
+**Date:** April 12, 2025.
 
 **Changelog:**
 - Add `encap_deterministic()` and `auth_encap_deterministic()` to `DhKem` in `hazardous::kem::x25519_hkdf_sha256::DhKem` [#458](https://github.com/orion-rs/orion/pull/458).

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "orion"
-version = "0.17.9"
+version = "0.17.10"
 authors = ["brycx <brycx@protonmail.com>"]
 description = "Usable, easy and safe pure-Rust crypto"
 keywords = ["cryptography", "crypto", "aead", "pqc", "kem"]

src/hazardous/ecc/x25519.rs

@@ -630,7 +630,6 @@ pub fn key_agreement(
 
 #[cfg(test)]
 mod public {
-    use super::FieldElement;
     use crate::hazardous::ecc::x25519::{
         key_agreement, PrivateKey, PublicKey, Scalar, SharedKey, BASEPOINT, PRIVATE_KEY_SIZE,
         PUBLIC_KEY_SIZE,
@@ -696,6 +695,7 @@ mod public {
     #[cfg(feature = "safe_api")]
     // format! is only available with std
     fn test_privatekey_debug_impl() {
+        use super::FieldElement;
         let value = format!("{:?}", [1u64, 0u64, 0u64, 0u64, 0u64,].as_ref());
         let test_debug_contents = format!("{:?}", FieldElement::one());
         assert!(test_debug_contents.contains(&value));

src/hazardous/kem/ml_kem/internal/mod.rs

@@ -86,7 +86,7 @@ pub fn g(c: &[&[u8]]) -> ([u8; 32], Zeroizing<[u8; 32]>) {
 }
 
 /// Internal PKE-related function, for generalizing over the three different PKE parameter-sets.
-pub(crate) trait PkeParameters: Clone {
+pub(crate) trait PkeParameters {
     const N: usize = 256;
     const K: usize;
     const ETA_1: usize;
@@ -523,7 +523,7 @@ pub(crate) struct DecapKey<
 > {
     pub(crate) bytes: [u8; ENCODED_SIZE_DK],
     s_hat: [RingElementNTT; K],
-    ek: EncapKey<K, ENCODED_SIZE_EK, Pke>,
+    _phantom: PhantomData<Pke>,
 }
 
 impl<
@@ -617,15 +617,10 @@ impl<
             ByteSerialization::decode_12(dk_part, &mut s_hat_poly.coefficients);
         }
 
-        // Save the encapsulation key, such that it doesn't need to be re-computed in MLKEM-decap_internal().
-        let ek = EncapKey::<K, ENCODED_SIZE_EK, Pke>::from_slice(
-            &slice[ENCODE_SIZE_POLY * K..(768 * K) + 32],
-        )?;
-
         Ok(Self {
             bytes: slice.try_into().unwrap(), // NOTE: Should never panic if decapsulation_key_check() succeeds.
             s_hat,
-            ek,
+            _phantom: PhantomData,
         })
     }
 
@@ -682,11 +677,13 @@ impl<
     /// - decapsulation key dk ∈ B^{768k+96}.
     /// - ciphertext c ∈ B^{32*(d_u*k+d_v)}.
     /// - shared secret K ∈ B^{32}.
-    pub(crate) fn mlkem_decap_internal(
+    pub(crate) fn mlkem_decap_internal_with_ek(
         &self,
         c: &[u8],
         c_prime: &mut [u8],
+        ek: &EncapKey<K, ENCODED_SIZE_EK, Pke>,
     ) -> Result<[u8; 32], UnknownCryptoError> {
+        debug_assert_eq!(self.get_encapsulation_key_bytes(), ek.as_ref());
         debug_assert_eq!(c.len(), Pke::CIPHERTEXT_SIZE);
 
         // Step 1:
@@ -712,9 +709,8 @@ impl<
         xof.squeeze(k_bar.as_mut())?;
 
         // Step 8:
-        debug_assert_eq!(self.get_encapsulation_key_bytes(), self.ek.as_ref());
         debug_assert_eq!(self.get_encapsulation_key_bytes(), ek_pke);
-        self.ek.encrypt(&m, r.as_ref(), c_prime)?;
+        ek.encrypt(&m, r.as_ref(), c_prime)?;
 
         // Step 9:
         let ct_choice = c.ct_ne(c_prime);
@@ -727,6 +723,19 @@ impl<
         Ok(k)
     }
 
+    #[cfg(feature = "safe_api")] // used in from_keys which requires safe_api
+    pub(crate) fn mlkem_decap_internal(
+        &self,
+        c: &[u8],
+        c_prime: &mut [u8],
+    ) -> Result<[u8; 32], UnknownCryptoError> {
+        // In this case we aren't provided a cached encapsulation key.
+        let ek =
+            EncapKey::<K, ENCODED_SIZE_EK, Pke>::from_slice(self.get_encapsulation_key_bytes())?;
+
+        self.mlkem_decap_internal_with_ek(c, c_prime, &ek)
+    }
+
     pub(crate) fn unprotected_as_bytes(&self) -> &[u8] {
         &self.bytes
     }
@@ -745,6 +754,7 @@ impl<Pke: PkeParameters> KeyPairInternal<Pke> {
     /// k \in [2, 3, 4]
     fn keygen<const K: usize, const ENCODED_SIZE_EK: usize, const ENCODED_SIZE_DK: usize>(
         d: &[u8],
+        ek: &mut EncapKey<K, ENCODED_SIZE_EK, Pke>,
         dk: &mut DecapKey<K, ENCODED_SIZE_EK, ENCODED_SIZE_DK, Pke>,
     ) -> Result<(), UnknownCryptoError> {
         let (rho, sigma) = g(&[d, &[Pke::K as u8]]);
@@ -753,7 +763,7 @@ impl<Pke: PkeParameters> KeyPairInternal<Pke> {
         // Steps 3..7
         for i in 0..Pke::K {
             for j in 0..Pke::K {
-                dk.ek.mat_a[i][j] = sample_ntt(&rho, &[j as u8, i as u8])?;
+                ek.mat_a[i][j] = sample_ntt(&rho, &[j as u8, i as u8])?;
             }
         }
 
@@ -773,34 +783,33 @@ impl<Pke: PkeParameters> KeyPairInternal<Pke> {
 
         for i in 0..Pke::K {
             dk.s_hat[i] = to_ntt(&s[i]);
-            dk.ek.t_hat[i] = to_ntt(&e[i]);
+            ek.t_hat[i] = to_ntt(&e[i]);
         }
 
         s.zeroize();
 
         // t ← A ∘ ŝ + ê
         for i in 0..Pke::K {
             for j in 0..Pke::K {
-                dk.ek.t_hat[i] += dk.ek.mat_a[i][j] * dk.s_hat[j];
+                ek.t_hat[i] += ek.mat_a[i][j] * dk.s_hat[j];
             }
         }
 
         // Step 19
-        for (re, ek_part) in dk
-            .ek
+        for (re, ek_part) in ek
             .t_hat
             .iter()
-            .zip(dk.ek.bytes.chunks_exact_mut(ENCODE_SIZE_POLY))
+            .zip(ek.bytes.chunks_exact_mut(ENCODE_SIZE_POLY))
         {
             ByteSerialization::encode_12(&re.coefficients, ek_part);
         }
 
         let idx = ENCODED_SIZE_EK - rho.len();
-        dk.ek.bytes[idx..].copy_from_slice(&rho);
+        ek.bytes[idx..].copy_from_slice(&rho);
 
         // Cache hash of ek so we don't need to re-compute for every encap().
-        let h_ek = Sha3_256::digest(&dk.ek.bytes)?;
-        dk.ek.h_ek = h_ek.value;
+        let h_ek = Sha3_256::digest(&ek.bytes)?;
+        ek.h_ek = h_ek.value;
 
         // Step 20
         for (re, dk_part) in dk
@@ -827,7 +836,7 @@ impl<Pke: PkeParameters> KeyPairInternal<Pke> {
         ),
         UnknownCryptoError,
     > {
-        let ek = EncapKey::<K, ENCODED_SIZE_EK, Pke> {
+        let mut encap_key = EncapKey::<K, ENCODED_SIZE_EK, Pke> {
             bytes: [0u8; ENCODED_SIZE_EK],
             h_ek: [0u8; SHA3_256_OUTSIZE],
             t_hat: [RingElementNTT::zero(); K],
@@ -839,28 +848,27 @@ impl<Pke: PkeParameters> KeyPairInternal<Pke> {
         let mut decap_key = DecapKey::<K, ENCODED_SIZE_EK, ENCODED_SIZE_DK, Pke> {
             bytes: [0u8; ENCODED_SIZE_DK],
             s_hat: [RingElementNTT::zero(); K],
-            ek,
+            _phantom: PhantomData,
         };
 
         // Step 1 + 2. (ekPKE, dkPKE) ← K-PKE.KeyGen(d)
-        Self::keygen(&seed.unprotected_as_bytes()[..32], &mut decap_key)?;
+        Self::keygen(
+            &seed.unprotected_as_bytes()[..32],
+            &mut encap_key,
+            &mut decap_key,
+        )?;
 
         // Step 3. dk ← (dkPKE‖ek‖H(ek)‖z)
-        let ek_bytes = decap_key.ek.as_ref();
         decap_key.bytes[(ENCODE_SIZE_POLY * K)..(ENCODE_SIZE_POLY * K) + Pke::EK_SIZE]
-            .copy_from_slice(ek_bytes);
+            .copy_from_slice(&encap_key.bytes);
         decap_key.bytes
             [(ENCODE_SIZE_POLY * K) + Pke::EK_SIZE..(ENCODE_SIZE_POLY * K) + Pke::EK_SIZE + 32]
-            .copy_from_slice(Sha3_256::digest(ek_bytes).unwrap().as_ref());
+            .copy_from_slice(Sha3_256::digest(&encap_key.bytes).unwrap().as_ref());
         decap_key.bytes[(ENCODE_SIZE_POLY * K) + Pke::EK_SIZE + 32
             ..(ENCODE_SIZE_POLY * K) + Pke::EK_SIZE + 32 + 32]
             .copy_from_slice(&seed.unprotected_as_bytes()[32..64]);
 
-        let encap_key = decap_key.ek.clone(); // TODO: Can we maybe get rid of this clone? Probably some internal API changes propagating upwards
-        debug_assert_eq!(
-            decap_key.get_encapsulation_key_bytes(),
-            decap_key.ek.as_ref()
-        );
+        debug_assert_eq!(decap_key.get_encapsulation_key_bytes(), encap_key.as_ref());
 
         Ok((encap_key, decap_key))
     }
@@ -927,20 +935,6 @@ mod tests {
     use crate::hazardous::kem::ml_kem::mlkem512::KeyPair as MlKem512KeyPair;
     use crate::hazardous::kem::ml_kem::mlkem768::KeyPair as MlKem768KeyPair;
 
-    #[test]
-    fn test_keypair_dk_ek_match_internal() {
-        let seed = Seed::from_slice(&[128u8; 64]).unwrap();
-
-        let kp = MlKem512KeyPair::try_from(&seed).unwrap();
-        assert_eq!(kp.public().value, kp.private().value.ek);
-
-        let kp = MlKem768KeyPair::try_from(&seed).unwrap();
-        assert_eq!(kp.public().value, kp.private().value.ek);
-
-        let kp = MlKem1024KeyPair::try_from(&seed).unwrap();
-        assert_eq!(kp.public().value, kp.private().value.ek);
-    }
-
     #[test]
     #[cfg(feature = "safe_api")]
     fn test_seed_and_dk_mismatch() {

src/hazardous/kem/ml_kem/mlkem1024.rs

@@ -110,7 +110,6 @@ impl_from_trait!(Ciphertext, MlKem1024Internal::CIPHERTEXT_SIZE);
 /// A keypair of ML-KEM-1024 keys, that are derived from a given seed.
 pub struct KeyPair {
     seed: Seed,
-    ek: EncapsulationKey,
     dk: DecapsulationKey,
 }
 
@@ -124,8 +123,10 @@ impl KeyPair {
 
         Ok(Self {
             seed,
-            ek: EncapsulationKey { value: ek },
-            dk: DecapsulationKey { value: dk },
+            dk: DecapsulationKey {
+                value: dk,
+                cached_ek: EncapsulationKey { value: ek },
+            },
         })
     }
 
@@ -147,8 +148,10 @@ impl KeyPair {
 
         Ok(Self {
             seed: Seed::from_slice(seed.unprotected_as_bytes()).unwrap(),
-            ek: EncapsulationKey { value: ek },
-            dk: DecapsulationKey { value: dk },
+            dk: DecapsulationKey {
+                value: dk,
+                cached_ek: EncapsulationKey { value: ek },
+            },
         })
     }
 
@@ -160,7 +163,7 @@ impl KeyPair {
 
     /// Get the public [EncapsulationKey] corresponding to this keypair.
     pub fn public(&self) -> &EncapsulationKey {
-        &self.ek
+        &self.dk.cached_ek
     }
 
     /// Get the private [DecapsulationKey] used to generate this keypair. In order to store the private
@@ -178,8 +181,10 @@ impl TryFrom<&Seed> for KeyPair {
 
         Ok(Self {
             seed: Seed::from_slice(value.unprotected_as_bytes()).unwrap(),
-            ek: EncapsulationKey { value: ek },
-            dk: DecapsulationKey { value: dk },
+            dk: DecapsulationKey {
+                value: dk,
+                cached_ek: EncapsulationKey { value: ek },
+            },
         })
     }
 }
@@ -188,6 +193,10 @@ impl TryFrom<&Seed> for KeyPair {
 /// A type to represent the `DecapsulationKey` that ML-KEM-1024 produces.
 pub struct DecapsulationKey {
     pub(crate) value: DecapKey<4, 1568, 3168, MlKem1024Internal>,
+    // NOTE(brycx): This is simply a cache of the encapsulation key, so we avoid recomputing it
+    // on decap() operations. This is not a part of PartialEq, AsRef<> implementations or other logic
+    // pertaining to the `DecapsulationKey`, serving a purely internal purpose.
+    pub(crate) cached_ek: EncapsulationKey,
 }
 
 impl PartialEq<&[u8]> for DecapsulationKey {
@@ -200,17 +209,25 @@ impl PartialEq<&[u8]> for DecapsulationKey {
 impl DecapsulationKey {
     /// Instantiate a [DecapsulationKey] with only key-checks from FIPS-203, section 7.3. Not MAL-BIND-K-CT secure.
     pub fn unchecked_from_slice(slice: &[u8]) -> Result<Self, UnknownCryptoError> {
+        let dk_unchecked =
+            DecapKey::<4, 1568, 3168, MlKem1024Internal>::unchecked_from_slice(slice)?;
+        let ek_unchecked =
+            EncapsulationKey::from_slice(dk_unchecked.get_encapsulation_key_bytes())?;
+
         Ok(Self {
-            value: DecapKey::<4, 1568, 3168, MlKem1024Internal>::unchecked_from_slice(slice)?,
+            value: dk_unchecked,
+            cached_ek: ek_unchecked,
         })
     }
 
     /// Perform decapsulation of a [Ciphertext].
     pub fn decap(&self, c: &Ciphertext) -> Result<SharedSecret, UnknownCryptoError> {
         let mut c_prime_buf = [0u8; MlKem1024Internal::CIPHERTEXT_SIZE];
-        let mut k_internal = self
-            .value
-            .mlkem_decap_internal(c.as_ref(), &mut c_prime_buf)?;
+        let mut k_internal = self.value.mlkem_decap_internal_with_ek(
+            c.as_ref(),
+            &mut c_prime_buf,
+            &self.cached_ek.value,
+        )?;
         let k = SharedSecret::from_slice(&k_internal)?;
         k_internal.zeroize();
 
@@ -329,7 +346,7 @@ mod tests {
             let kp = KeyPair::try_from(&Seed::from_slice(seed).unwrap()).unwrap();
 
             Ok((
-                kp.ek.as_ref().to_vec(),
+                kp.dk.cached_ek.as_ref().to_vec(),
                 kp.dk.value.unprotected_as_bytes().to_vec(),
             ))
         }
@@ -349,11 +366,41 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_keypair_dk_ek_match_internal() {
+        let seed = Seed::from_slice(&[128u8; 64]).unwrap();
+        let kp = KeyPair::try_from(&seed).unwrap();
+        assert_eq!(kp.public(), &kp.private().cached_ek);
+    }
+
+    #[test]
+    #[cfg(feature = "safe_api")]
+    fn test_dk_cached_ek() {
+        let seed = Seed::from_slice(&[128u8; 64]).unwrap();
+        let kp = KeyPair::try_from(&seed).unwrap();
+        let (ss_pubapi, ct_pubapi) = kp.public().encap_deterministic(&[125u8; 32]).unwrap();
+        let mut c_prime = [0u8; MlKem1024Internal::CIPHERTEXT_SIZE];
+        // This call re-computes encap key internally from the bytes a decapkey would store.
+        let ss_privapi = kp
+            .private()
+            .value
+            .mlkem_decap_internal(ct_pubapi.as_ref(), &mut c_prime)
+            .unwrap();
+        assert_eq!(ss_privapi.as_ref(), ss_pubapi.unprotected_as_bytes());
+        assert_eq!(
+            MlKem1024::decap(kp.private(), &ct_pubapi).unwrap(),
+            ss_pubapi
+        );
+    }
+
     #[cfg(feature = "safe_api")]
     #[test]
     fn test_dk_to_ek_conversions() {
         let kp = KeyPair::generate().unwrap();
-        assert_eq!(kp.ek, EncapsulationKey::try_from(kp.private()).unwrap());
+        assert_eq!(
+            kp.dk.cached_ek,
+            EncapsulationKey::try_from(kp.private()).unwrap()
+        );
     }
 
     #[cfg(feature = "safe_api")]