Merge pull request #546 from cryspen/franziskus/ml-kem-rand

Add randomised APIs to ml-kem
cryspen · Sep 3, 2024 · 482620d · 482620d
2 parents 9680093 + 08491c7
commit 482620d
Show file tree

Hide file tree

Showing 16 changed files with 234 additions and 88 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/libcrux-ecdh/Cargo.toml b/libcrux-ecdh/Cargo.toml
@@ -22,3 +22,6 @@ hex = { version = "0.4.3", features = ["serde"] }
 serde_json = { version = "1.0" }
 serde = { version = "1.0", features = ["derive"] }
 pretty_env_logger = "0.5"
+
+[lints.rust]
+unexpected_cfgs = { level = "warn", check-cfg = ['cfg(adx)', 'cfg(bmi2)'] }
diff --git a/libcrux-kem/src/kem.rs b/libcrux-kem/src/kem.rs
@@ -215,10 +215,14 @@ pub struct X25519MlKem768Draft00PublicKey {
 impl X25519MlKem768Draft00PublicKey {
     pub fn decode(bytes: &[u8]) -> Result<Self, Error> {
         Ok(Self {
-            mlkem: MlKem768PublicKey::try_from(&bytes[32..])
-                .ok()
-                .and_then(mlkem768::validate_public_key)
-                .ok_or(Error::InvalidPublicKey)?,
+            mlkem: {
+                let key = MlKem768PublicKey::try_from(&bytes[32..])
+                    .map_err(|_| Error::InvalidPublicKey)?;
+                if !mlkem768::validate_public_key(&key) {
+                    return Err(Error::InvalidPublicKey);
+                }
+                key
+            },
             x25519: bytes[0..32]
                 .try_into()
                 .map_err(|_| Error::InvalidPublicKey)?,
@@ -241,10 +245,14 @@ pub struct XWingKemDraft02PublicKey {
 impl XWingKemDraft02PublicKey {
     pub fn decode(bytes: &[u8]) -> Result<Self, Error> {
         Ok(Self {
-            pk_m: MlKem768PublicKey::try_from(&bytes[0..1184])
-                .ok()
-                .and_then(mlkem768::validate_public_key)
-                .ok_or(Error::InvalidPublicKey)?,
+            pk_m: {
+                let key = MlKem768PublicKey::try_from(&bytes[0..1184])
+                    .map_err(|_| Error::InvalidPublicKey)?;
+                if !mlkem768::validate_public_key(&key) {
+                    return Err(Error::InvalidPublicKey);
+                }
+                key
+            },
             pk_x: bytes[1184..]
                 .try_into()
                 .map_err(|_| Error::InvalidPublicKey)?,
@@ -654,16 +662,22 @@ impl PublicKey {
                 .try_into()
                 .map(Self::P256)
                 .map_err(|_| Error::InvalidPublicKey),
-            Algorithm::MlKem512 => MlKem512PublicKey::try_from(bytes)
-                .ok()
-                .and_then(mlkem512::validate_public_key)
-                .map(Self::MlKem512)
-                .ok_or(Error::InvalidPublicKey),
-            Algorithm::MlKem768 => MlKem768PublicKey::try_from(bytes)
-                .ok()
-                .and_then(mlkem768::validate_public_key)
-                .map(Self::MlKem768)
-                .ok_or(Error::InvalidPublicKey),
+            Algorithm::MlKem512 => {
+                let key =
+                    MlKem512PublicKey::try_from(bytes).map_err(|_| Error::InvalidPublicKey)?;
+                if !mlkem512::validate_public_key(&key) {
+                    return Err(Error::InvalidPublicKey);
+                }
+                Ok(Self::MlKem512(key))
+            }
+            Algorithm::MlKem768 => {
+                let key =
+                    MlKem768PublicKey::try_from(bytes).map_err(|_| Error::InvalidPublicKey)?;
+                if !mlkem768::validate_public_key(&key) {
+                    return Err(Error::InvalidPublicKey);
+                }
+                Ok(Self::MlKem768(key))
+            }
             Algorithm::X25519MlKem768Draft00 => {
                 X25519MlKem768Draft00PublicKey::decode(bytes).map(Self::X25519MlKem768Draft00)
             }
@@ -678,11 +692,14 @@ impl PublicKey {
             Algorithm::XWingKyberDraft02 => {
                 XWingKemDraft02PublicKey::decode(bytes).map(Self::XWingKyberDraft02)
             }
-            Algorithm::MlKem1024 => MlKem1024PublicKey::try_from(bytes)
-                .ok()
-                .and_then(mlkem1024::validate_public_key)
-                .map(Self::MlKem1024)
-                .ok_or(Error::InvalidPublicKey),
+            Algorithm::MlKem1024 => {
+                let key =
+                    MlKem1024PublicKey::try_from(bytes).map_err(|_| Error::InvalidPublicKey)?;
+                if !mlkem1024::validate_public_key(&key) {
+                    return Err(Error::InvalidPublicKey);
+                }
+                Ok(Self::MlKem1024(key))
+            }
             _ => Err(Error::UnsupportedAlgorithm),
         }
     }

diff --git a/libcrux-kem/tests/ml_kem_wycheproof_early.rs b/libcrux-kem/tests/ml_kem_wycheproof_early.rs
@@ -91,7 +91,7 @@ macro_rules! impl_known_answer_test {
                     assert_eq!(my_shared_secret.as_ref(), expected_shared_secret);
                 } else {
                     if comment == "Public key not reduced" {
-                        assert!($validate_pk(<$pk>::from(pk)).is_none());
+                        assert!(!$validate_pk(&<$pk>::from(pk)));
                     }
                 }
             }

diff --git a/libcrux-ml-kem/Cargo.toml b/libcrux-ml-kem/Cargo.toml
@@ -22,7 +22,7 @@ exclude = [
 bench = false # so libtest doesn't eat the arguments to criterion
 
 [dependencies]
-rand_core = { version = "0.6" }
+rand = { version = "0.8", optional = true }
 libcrux-platform = { version = "0.0.2-alpha.3", path = "../sys/platform" }
 libcrux-sha3 = { version = "0.0.2-alpha.3", path = "../libcrux-sha3" }
 libcrux-intrinsics = { version = "0.0.2-alpha.3", path = "../libcrux-intrinsics" }
@@ -34,7 +34,7 @@ hax-lib = { version = "0.1.0-alpha.1", git = "https://github.com/hacspec/hax/" }
 
 [features]
 # By default all variants and std are enabled.
-default = ["std", "mlkem512", "mlkem768", "mlkem1024"]
+default = ["std", "mlkem512", "mlkem768", "mlkem1024", "rand"]
 
 # Hardware features can be force enabled.
 # It is not recommended to use these. This crate performs CPU feature detection
@@ -56,6 +56,10 @@ kyber = []
 # Code that is not yet verified
 pre-verification = []
 
+# APIs that sample their own randomness
+rand = ["dep:rand"]
+
+# std support
 std = []
 
 [dev-dependencies]

diff --git a/libcrux-ml-kem/benches/ml-kem.rs b/libcrux-ml-kem/benches/ml-kem.rs
@@ -2,8 +2,7 @@ use std::hint::black_box;
 use std::time::Duration;
 
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
-use rand_core::OsRng;
-use rand_core::RngCore;
+use rand::{rngs::OsRng, RngCore};
 
 use libcrux_ml_kem::{mlkem1024, mlkem512, mlkem768};
 

diff --git a/libcrux-ml-kem/src/hash_functions.rs b/libcrux-ml-kem/src/hash_functions.rs
@@ -173,6 +173,7 @@ pub(crate) mod portable {
 }
 
 /// A SIMD256 implementation of [`Hash`] for AVX2
+#[cfg(feature = "simd256")]
 pub(crate) mod avx2 {
     use super::*;
     use libcrux_sha3::{

diff --git a/libcrux-ml-kem/src/ind_cca/multiplexing.rs b/libcrux-ml-kem/src/ind_cca/multiplexing.rs
@@ -13,7 +13,7 @@ use instantiations::avx2::{
 #[cfg(feature = "simd128")]
 use instantiations::neon::{
     decapsulate as decapsulate_neon, encapsulate as encapsulate_neon,
-    generate_keypair as generate_keypair_neon, validate_public_key as validate_public_key_neon,
+    generate_keypair as generate_keypair_neon,
 };
 
 #[cfg(not(feature = "simd256"))]