add an additional MLKEM implementation, and select with a feature

Cargo.toml

@@ -38,6 +38,7 @@ required-features = ["test-utils"]
 
 [features]
 test-utils = []
+mlkem768 = []
 
 [dependencies]
 cosmian_crypto_core = { version = "10.0.0", default-features = false, features = ["ser", "sha3", "aes", "curve25519"] }

src/core/kem.rs

@@ -1,154 +1,7 @@
-use cosmian_crypto_core::bytes_ser_de::{Deserializer, Serializable, Serializer};
-use cosmian_crypto_core::{reexport::rand_core::CryptoRngCore, Secret};
-use ml_kem::{
-    array::Array,
-    kem::{Decapsulate, Encapsulate},
-    EncodedSizeUser, KemCore,
-};
-use zeroize::Zeroize;
+pub mod mlkem;
 
-use crate::traits::Kem;
-use crate::{core::SHARED_SECRET_LENGTH, Error};
+#[cfg(not(feature = "mlkem768"))]
+pub use mlkem::MlKem512 as MlKem;
 
-#[derive(Debug, PartialEq, Clone)]
-pub struct EncapsulationKey512(Box<<ml_kem::MlKem512 as KemCore>::EncapsulationKey>);
-
-impl Serializable for EncapsulationKey512 {
-    type Error = Error;
-
-    fn length(&self) -> usize {
-        800
-    }
-
-    fn write(&self, ser: &mut Serializer) -> Result<usize, Self::Error> {
-        let mut bytes = self.0.as_bytes();
-        let n = ser.write_array(&bytes)?;
-        bytes.zeroize();
-        Ok(n)
-    }
-
-    fn read(de: &mut Deserializer) -> Result<Self, Self::Error> {
-        let mut bytes = Array::from(de.read_array::<800>()?);
-        let ek = <<ml_kem::MlKem512 as KemCore>::EncapsulationKey>::from_bytes(&bytes);
-        bytes.zeroize();
-        Ok(Self(Box::new(ek)))
-    }
-}
-
-#[derive(Debug, Clone, PartialEq)]
-pub struct DecapsulationKey512(Box<<ml_kem::MlKem512 as KemCore>::DecapsulationKey>);
-
-impl DecapsulationKey512 {
-    pub fn ek(&self) -> EncapsulationKey512 {
-        EncapsulationKey512(Box::new(self.0.encapsulation_key().clone()))
-    }
-}
-
-impl Serializable for DecapsulationKey512 {
-    type Error = Error;
-
-    fn length(&self) -> usize {
-        1632
-    }
-
-    fn write(&self, ser: &mut Serializer) -> Result<usize, Self::Error> {
-        let mut bytes = self.0.as_bytes();
-        let n = ser.write_array(&bytes)?;
-        bytes.zeroize();
-        Ok(n)
-    }
-
-    fn read(de: &mut Deserializer) -> Result<Self, Self::Error> {
-        let mut bytes = Array::from(de.read_array::<1632>()?);
-        let dk = <<ml_kem::MlKem512 as KemCore>::DecapsulationKey>::from_bytes(&bytes);
-        bytes.zeroize();
-        Ok(Self(Box::new(dk)))
-    }
-}
-
-#[derive(Debug, PartialEq, Eq, Clone, Hash)]
-pub struct Encapsulation512(Box<Array<u8, <ml_kem::MlKem512 as KemCore>::CiphertextSize>>);
-
-impl Serializable for Encapsulation512 {
-    type Error = Error;
-
-    fn length(&self) -> usize {
-        768
-    }
-
-    fn write(&self, ser: &mut Serializer) -> Result<usize, Self::Error> {
-        Ok(ser.write_array(&self.0)?)
-    }
-
-    fn read(de: &mut Deserializer) -> Result<Self, Self::Error> {
-        Ok(Self(Box::new(Array::<
-            u8,
-            <ml_kem::MlKem512 as KemCore>::CiphertextSize,
-        >::from(de.read_array::<768>()?))))
-    }
-}
-
-pub struct MlKem512;
-
-impl Kem for MlKem512 {
-    type EncapsulationKey = EncapsulationKey512;
-    type DecapsulationKey = DecapsulationKey512;
-    type SessionKey = Secret<SHARED_SECRET_LENGTH>;
-
-    type Encapsulation = Encapsulation512;
-
-    type Error = Error;
-
-    fn keygen(
-        rng: &mut impl CryptoRngCore,
-    ) -> Result<(Self::DecapsulationKey, Self::EncapsulationKey), Self::Error> {
-        let (dk, ek) = <ml_kem::MlKem512 as KemCore>::generate(rng);
-        Ok((
-            DecapsulationKey512(Box::new(dk)),
-            EncapsulationKey512(Box::new(ek)),
-        ))
-    }
-
-    fn enc(
-        ek: &Self::EncapsulationKey,
-        rng: &mut impl CryptoRngCore,
-    ) -> Result<(Self::SessionKey, Self::Encapsulation), Self::Error> {
-        let (enc, mut ss) =
-            ek.0.encapsulate(rng)
-                .map_err(|e| Error::Kem(format!("{:?}", e)))?;
-        let ss = Secret::from_unprotected_bytes(ss.as_mut());
-        Ok((ss, Encapsulation512(Box::new(enc))))
-    }
-
-    fn dec(
-        dk: &Self::DecapsulationKey,
-        enc: &Self::Encapsulation,
-    ) -> Result<Self::SessionKey, Self::Error> {
-        let mut ss =
-            dk.0.decapsulate(&enc.0)
-                .map_err(|e| Self::Error::Kem(format!("{e:?}")))?;
-        let ss = Secret::from_unprotected_bytes(ss.as_mut());
-        Ok(ss)
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use cosmian_crypto_core::{
-        bytes_ser_de::test_serialization, reexport::rand_core::SeedableRng, CsRng,
-    };
-
-    use super::*;
-
-    #[test]
-    fn test_mlkem() {
-        let mut rng = CsRng::from_entropy();
-        let (dk, ek) = MlKem512::keygen(&mut rng).unwrap();
-        test_serialization(&dk).unwrap();
-        test_serialization(&ek).unwrap();
-        let (ss1, enc) = MlKem512::enc(&ek, &mut rng).unwrap();
-        test_serialization(&enc).unwrap();
-        let ss2 = MlKem512::dec(&dk, &enc).unwrap();
-        assert_eq!(ss1, ss2);
-    }
-}
+#[cfg(feature = "mlkem768")]
+pub use mlkem::MlKem768 as MlKem;

src/core/kem/mlkem.rs

@@ -0,0 +1,183 @@
+use cosmian_crypto_core::bytes_ser_de::{Deserializer, Serializable, Serializer};
+use cosmian_crypto_core::{reexport::rand_core::CryptoRngCore, Secret};
+use ml_kem::{
+    array::Array,
+    kem::{Decapsulate, Encapsulate},
+    EncodedSizeUser, KemCore,
+};
+use zeroize::Zeroize;
+
+use crate::traits::Kem;
+use crate::{core::SHARED_SECRET_LENGTH, Error};
+
+macro_rules! make_mlkem {
+    ($base: ident, $ek: ident, $ek_len: literal, $dk: ident, $dk_len: literal, $enc: ident, $enc_len:literal) => {
+        #[derive(Debug, PartialEq, Clone)]
+        pub struct $ek(Box<<ml_kem::$base as KemCore>::EncapsulationKey>);
+
+        impl Serializable for $ek {
+            type Error = Error;
+
+            fn length(&self) -> usize {
+                $ek_len
+            }
+
+            fn write(&self, ser: &mut Serializer) -> Result<usize, Self::Error> {
+                let mut bytes = self.0.as_bytes();
+                let n = ser.write_array(&bytes)?;
+                bytes.zeroize();
+                Ok(n)
+            }
+
+            fn read(de: &mut Deserializer) -> Result<Self, Self::Error> {
+                let mut bytes = Array::from(de.read_array::<$ek_len>()?);
+                let ek = <<ml_kem::$base as KemCore>::EncapsulationKey>::from_bytes(&bytes);
+                bytes.zeroize();
+                Ok(Self(Box::new(ek)))
+            }
+        }
+
+        #[derive(Debug, Clone, PartialEq)]
+        pub struct $dk(Box<<ml_kem::$base as KemCore>::DecapsulationKey>);
+
+        #[allow(dead_code)]
+        impl $dk {
+            pub fn ek(&self) -> $ek {
+                $ek(Box::new(self.0.encapsulation_key().clone()))
+            }
+        }
+
+        impl Serializable for $dk {
+            type Error = Error;
+
+            fn length(&self) -> usize {
+                $dk_len
+            }
+
+            fn write(&self, ser: &mut Serializer) -> Result<usize, Self::Error> {
+                let mut bytes = self.0.as_bytes();
+                let n = ser.write_array(&bytes)?;
+                bytes.zeroize();
+                Ok(n)
+            }
+
+            fn read(de: &mut Deserializer) -> Result<Self, Self::Error> {
+                let mut bytes = Array::from(de.read_array::<$dk_len>()?);
+                let dk = <<ml_kem::$base as KemCore>::DecapsulationKey>::from_bytes(&bytes);
+                bytes.zeroize();
+                Ok(Self(Box::new(dk)))
+            }
+        }
+
+        #[derive(Debug, PartialEq, Eq, Clone, Hash)]
+        pub struct $enc(Box<Array<u8, <ml_kem::$base as KemCore>::CiphertextSize>>);
+
+        impl Serializable for $enc {
+            type Error = Error;
+
+            fn length(&self) -> usize {
+                $enc_len
+            }
+
+            fn write(&self, ser: &mut Serializer) -> Result<usize, Self::Error> {
+                Ok(ser.write_array(&self.0)?)
+            }
+
+            fn read(de: &mut Deserializer) -> Result<Self, Self::Error> {
+                Ok(Self(Box::new(Array::<
+                    u8,
+                    <ml_kem::$base as KemCore>::CiphertextSize,
+                >::from(de.read_array::<$enc_len>()?))))
+            }
+        }
+
+        pub struct $base;
+
+        impl Kem for $base {
+            type EncapsulationKey = $ek;
+            type DecapsulationKey = $dk;
+            type SessionKey = Secret<SHARED_SECRET_LENGTH>;
+
+            type Encapsulation = $enc;
+
+            type Error = Error;
+
+            fn keygen(
+                rng: &mut impl CryptoRngCore,
+            ) -> Result<(Self::DecapsulationKey, Self::EncapsulationKey), Self::Error> {
+                let (dk, ek) = <ml_kem::$base as KemCore>::generate(rng);
+                Ok(($dk(Box::new(dk)), $ek(Box::new(ek))))
+            }
+
+            fn enc(
+                ek: &Self::EncapsulationKey,
+                rng: &mut impl CryptoRngCore,
+            ) -> Result<(Self::SessionKey, Self::Encapsulation), Self::Error> {
+                let (enc, mut ss) =
+                    ek.0.encapsulate(rng)
+                        .map_err(|e| Error::Kem(format!("{:?}", e)))?;
+                let ss = Secret::from_unprotected_bytes(ss.as_mut());
+                Ok((ss, $enc(Box::new(enc))))
+            }
+
+            fn dec(
+                dk: &Self::DecapsulationKey,
+                enc: &Self::Encapsulation,
+            ) -> Result<Self::SessionKey, Self::Error> {
+                let mut ss =
+                    dk.0.decapsulate(&enc.0)
+                        .map_err(|e| Self::Error::Kem(format!("{e:?}")))?;
+                let ss = Secret::from_unprotected_bytes(ss.as_mut());
+                Ok(ss)
+            }
+        }
+    };
+}
+
+make_mlkem!(
+    MlKem512,
+    EncapsulationKey512,
+    800,
+    DecapsulationKey512,
+    1632,
+    Encapsulation512,
+    768
+);
+
+make_mlkem!(
+    MlKem768,
+    EncapsulationKey768,
+    1184,
+    DecapsulationKey768,
+    2400,
+    Encapsulation768,
+    1088
+);
+
+#[cfg(test)]
+mod tests {
+    use cosmian_crypto_core::{
+        bytes_ser_de::test_serialization, reexport::rand_core::SeedableRng, CsRng,
+    };
+
+    use super::*;
+
+    macro_rules! test_mlkem {
+        ($base:ident, $test_name:ident) => {
+            #[test]
+            fn $test_name() {
+                let mut rng = CsRng::from_entropy();
+                let (dk, ek) = $base::keygen(&mut rng).unwrap();
+                test_serialization(&dk).unwrap();
+                test_serialization(&ek).unwrap();
+                let (ss1, enc) = $base::enc(&ek, &mut rng).unwrap();
+                test_serialization(&enc).unwrap();
+                let ss2 = $base::dec(&dk, &enc).unwrap();
+                assert_eq!(ss1, ss2);
+            }
+        };
+    }
+
+    test_mlkem!(MlKem512, test_mlkem512);
+    test_mlkem!(MlKem768, test_mlkem768);
+}