mlkem: Cache hash of encapsulation key (#464)

* mlkem: Cache hash of encapsulation key to avoid re-computing for multiple encap() calls

* Update CHANGELOG

* mlkem: Test KeyPair internal ek caches

* nit

Author: brycx

CHANGELOG.md

@@ -4,9 +4,11 @@
 
 **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).
-- Make `hazardous::kem::x25519_hkdf_sha256::DhKem` available in `[no_std]` context [#458](https://github.com/orion-rs/orion/pull/458).
+- Make `hazardous::kem::x25519_hkdf_sha256::DhKem` available in `#![no_std]` context [#458](https://github.com/orion-rs/orion/pull/458).
 - Add support for HPKE (RFC 9180) [#458](https://github.com/orion-rs/orion/pull/458).
 - Switch to source-based code coverage [#462](https://github.com/orion-rs/orion/pull/462).
+- ML-KEM (internal): Cache hash of encapsulation key to save computation on multiple `encap()` operations [#464](https://github.com/orion-rs/orion/pull/464).
+- ML-KEM (internal): Cache encapsulation key within decapsulation key, to avoid re-computation after generation of decapsulation key [#447](https://github.com/orion-rs/orion/pull/447).
 
 ### 0.17.9
 

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

@@ -32,11 +32,11 @@ pub(crate) mod serialization;
 /// Sampling ring elements from seeds.
 pub(crate) mod sampling;
 
-use crate::errors::UnknownCryptoError;
 use crate::hazardous::hash::sha3::sha3_256::Sha3_256;
 use crate::hazardous::hash::sha3::sha3_512::Sha3_512;
 use crate::hazardous::hash::sha3::shake256;
 use crate::hazardous::kem::ml_kem::Seed;
+use crate::{errors::UnknownCryptoError, hazardous::hash::sha3::sha3_256::SHA3_256_OUTSIZE};
 use core::marker::PhantomData;
 use fe::*;
 use re::*;
@@ -323,6 +323,7 @@ impl PkeParameters for MlKem1024Internal {
 /// ML-KEM encapsulation key.
 pub(crate) struct EncapKey<const K: usize, const ENCODED_SIZE: usize, Pke: PkeParameters> {
     pub(crate) bytes: [u8; ENCODED_SIZE],
+    h_ek: [u8; SHA3_256_OUTSIZE],
     t_hat: [RingElementNTT; K],
     mat_a: [[RingElementNTT; K]; K],
     _phantom: PhantomData<Pke>,
@@ -371,8 +372,12 @@ impl<const K: usize, const ENCODED_SIZE: usize, Pke: PkeParameters> EncapKey<K,
             }
         }
 
+        // Cache hash of the bytes so we don't need to re-compute for every encap().
+        let h_ek = Sha3_256::digest(slice)?;
+
         Ok(Self {
             bytes: slice.try_into().unwrap(), // NOTE: Should never panic if encapsulation_key_check() succeeds.
+            h_ek: h_ek.value,
             t_hat,
             mat_a,
             _phantom: PhantomData,
@@ -500,7 +505,7 @@ impl<const K: usize, const ENCODED_SIZE: usize, Pke: PkeParameters> EncapKey<K,
         }
 
         // Step 1: (K, r) ← G(m‖H(ek))
-        let (k, r) = g(&[m, Sha3_256::digest(&self.bytes).unwrap().as_ref()]);
+        let (k, r) = g(&[m, self.h_ek.as_ref()]);
 
         // Step 2: c ← K-PKE.Encrypt(ek, m, r)
         self.encrypt(m, r.as_ref(), c)?;
@@ -793,6 +798,10 @@ impl<Pke: PkeParameters> KeyPairInternal<Pke> {
         let idx = ENCODED_SIZE_EK - rho.len();
         dk.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;
+
         // Step 20
         for (re, dk_part) in dk
             .s_hat
@@ -820,6 +829,7 @@ impl<Pke: PkeParameters> KeyPairInternal<Pke> {
     > {
         let ek = EncapKey::<K, ENCODED_SIZE_EK, Pke> {
             bytes: [0u8; ENCODED_SIZE_EK],
+            h_ek: [0u8; SHA3_256_OUTSIZE],
             t_hat: [RingElementNTT::zero(); K],
             mat_a: [[RingElementNTT::zero(); K]; K],
             _phantom: PhantomData,
@@ -917,6 +927,20 @@ 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() {