Assert that MLKEM APIs generate real keys

Author: kornelski

boring/src/mlkem.rs

@@ -21,6 +21,7 @@ use std::mem::MaybeUninit;
 use crate::cvt;
 use crate::error::ErrorStack;
 use crate::ffi;
+use crate::memcmp::eq;
 
 // CBS_init is inline in BoringSSL, so bindgen can't generate bindings for it.
 #[inline]
@@ -131,16 +132,25 @@ impl MlKemPublicKey {
     /// Encapsulates a shared secret to the given public key, returning
     /// `(ciphertext, shared_secret)`.
     pub fn encapsulate(&self) -> Result<(Vec<u8>, MlKemSharedSecret), ErrorStack> {
-        match &self.0 {
+        let mut ss = [0; SHARED_SECRET_BYTES];
+        let ct = match &self.0 {
             Either::MlKem768(pk) => {
-                let (ct, ss) = pk.encapsulate();
-                Ok((ct.to_vec(), ss))
+                let mut ct = vec![0; MlKem768PrivateKey::CIPHERTEXT_BYTES];
+                pk.encapsulate_into(ct.as_mut_slice().try_into().unwrap(), &mut ss);
+                ct
             }
             Either::MlKem1024(pk) => {
-                let (ct, ss) = pk.encapsulate();
-                Ok((ct.to_vec(), ss))
+                let mut ct = vec![0; MlKem1024PrivateKey::CIPHERTEXT_BYTES];
+                pk.encapsulate_into(ct.as_mut_slice().try_into().unwrap(), &mut ss);
+                ct
             }
+        };
+        if eq(ss, [0; SHARED_SECRET_BYTES]) {
+            return Err(ErrorStack::internal_error_str(
+                "Amazing! I've got the same combination on my luggage",
+            ));
         }
+        Ok((ct, ss))
     }
 
     /// Query public key and ciphertext length
@@ -229,14 +239,17 @@ impl MlKem768PrivateKey {
             ffi::init();
             let mut public_key_bytes: MaybeUninit<[u8; MlKem768PublicKey::PUBLIC_KEY_BYTES]> =
                 MaybeUninit::uninit();
-            let mut seed: MaybeUninit<MlKemPrivateKeySeed> = MaybeUninit::uninit();
+            let mut seed = [0; PRIVATE_KEY_SEED_BYTES];
             let mut expanded: MaybeUninit<ffi::MLKEM768_private_key> = MaybeUninit::uninit();
 
             ffi::MLKEM768_generate_key(
                 public_key_bytes.as_mut_ptr().cast(),
-                seed.as_mut_ptr().cast(),
+                seed.as_mut_ptr(),
                 expanded.as_mut_ptr(),
             );
+            if eq(seed, [0; PRIVATE_KEY_SEED_BYTES]) {
+                panic!("Amazing! I've got the same combination on my luggage");
+            }
 
             let bytes = public_key_bytes.assume_init();
 
@@ -251,7 +264,7 @@ impl MlKem768PrivateKey {
                     parsed: parsed.assume_init(),
                 }),
                 Box::new(MlKem768PrivateKey {
-                    seed: seed.assume_init(),
+                    seed,
                     expanded: expanded.assume_init(),
                 }),
             )
@@ -388,25 +401,19 @@ impl MlKem768PublicKey {
     }
 
     /// Encapsulate: returns (ciphertext, shared_secret).
-    fn encapsulate(
+    fn encapsulate_into(
         &self,
-    ) -> (
-        [u8; MlKem768PrivateKey::CIPHERTEXT_BYTES],
-        MlKemSharedSecret,
+        ciphertext: &mut [u8; MlKem768PrivateKey::CIPHERTEXT_BYTES],
+        shared_secret: &mut MlKemSharedSecret,
     ) {
         // SAFETY: buffers correctly sized, parsed key is valid
         unsafe {
             ffi::init();
-            let mut ciphertext = [0u8; MlKem768PrivateKey::CIPHERTEXT_BYTES];
-            let mut shared_secret = [0u8; SHARED_SECRET_BYTES];
-
             ffi::MLKEM768_encap(
                 ciphertext.as_mut_ptr(),
                 shared_secret.as_mut_ptr(),
                 &self.parsed,
             );
-
-            (ciphertext, shared_secret)
         }
     }
 }
@@ -446,14 +453,17 @@ impl MlKem1024PrivateKey {
             ffi::init();
             let mut public_key_bytes: MaybeUninit<[u8; MlKem1024PublicKey::PUBLIC_KEY_BYTES]> =
                 MaybeUninit::uninit();
-            let mut seed: MaybeUninit<MlKemPrivateKeySeed> = MaybeUninit::uninit();
+            let mut seed = [0; PRIVATE_KEY_SEED_BYTES];
             let mut expanded: MaybeUninit<ffi::MLKEM1024_private_key> = MaybeUninit::uninit();
 
             ffi::MLKEM1024_generate_key(
                 public_key_bytes.as_mut_ptr().cast(),
-                seed.as_mut_ptr().cast(),
+                seed.as_mut_ptr(),
                 expanded.as_mut_ptr(),
             );
+            if eq(seed, [0; PRIVATE_KEY_SEED_BYTES]) {
+                panic!("Amazing! I've got the same combination on my luggage");
+            }
 
             let bytes = public_key_bytes.assume_init();
 
@@ -468,7 +478,7 @@ impl MlKem1024PrivateKey {
                     parsed: parsed.assume_init(),
                 }),
                 Box::new(MlKem1024PrivateKey {
-                    seed: seed.assume_init(),
+                    seed,
                     expanded: expanded.assume_init(),
                 }),
             )
@@ -607,25 +617,19 @@ impl MlKem1024PublicKey {
     }
 
     /// Encapsulate: returns (ciphertext, shared_secret).
-    fn encapsulate(
+    fn encapsulate_into(
         &self,
-    ) -> (
-        [u8; MlKem1024PrivateKey::CIPHERTEXT_BYTES],
-        [u8; SHARED_SECRET_BYTES],
+        ciphertext: &mut [u8; MlKem1024PrivateKey::CIPHERTEXT_BYTES],
+        shared_secret: &mut [u8; SHARED_SECRET_BYTES],
     ) {
         // SAFETY: buffers correctly sized, parsed key is valid
         unsafe {
             ffi::init();
-            let mut ciphertext = [0u8; MlKem1024PrivateKey::CIPHERTEXT_BYTES];
-            let mut shared_secret = [0u8; SHARED_SECRET_BYTES];
-
             ffi::MLKEM1024_encap(
                 ciphertext.as_mut_ptr(),
                 shared_secret.as_mut_ptr(),
                 &self.parsed,
             );
-
-            (ciphertext, shared_secret)
         }
     }
 }
@@ -650,7 +654,9 @@ mod tests {
                 #[test]
                 fn roundtrip() {
                     let (pk, sk) = <$priv>::generate();
-                    let (ct, ss1) = pk.encapsulate();
+                    let mut ct = [0; _];
+                    let mut ss1 = [0; _];
+                    pk.encapsulate_into(&mut ct, &mut ss1);
                     let ss2 = sk.decapsulate(&ct);
                     assert_eq!(ss1, ss2);
                 }
@@ -659,7 +665,9 @@ mod tests {
                 fn seed_roundtrip() {
                     let (pk, sk) = <$priv>::generate();
                     let sk2 = <$priv>::from_seed(&sk.seed).unwrap();
-                    let (ct, ss1) = pk.encapsulate();
+                    let mut ct = [0; _];
+                    let mut ss1 = [0; _];
+                    pk.encapsulate_into(&mut ct, &mut ss1);
                     let ss2 = sk2.decapsulate(&ct);
                     assert_eq!(ss1, ss2);
                 }