fixup kyber tests for wasm

cryspen · Sep 18, 2023 · f798b5a · f798b5a
1 parent f7f59fc
commit f798b5a
Show file tree

Hide file tree

Showing 5 changed files with 61 additions and 28 deletions.
diff --git a/examples/kyber768_encapsulate.rs b/examples/kyber768_encapsulate.rs
@@ -1,13 +1,20 @@
-use libcrux::digest;
-use libcrux::drbg::Drbg;
 use libcrux::kem;
 
+#[cfg(not(target_arch = "wasm32"))]
+use libcrux::drbg;
+#[cfg(target_arch = "wasm32")]
+use rand_core::OsRng;
+
 fn main() {
-    let mut drbg = Drbg::new(digest::Algorithm::Sha256).unwrap();
-    let (_secret_key, public_key) = kem::key_gen(kem::Algorithm::Kyber768, &mut drbg).unwrap();
+    #[cfg(not(target_arch = "wasm32"))]
+    let mut rng = drbg::Drbg::new(libcrux::digest::Algorithm::Sha256).unwrap();
+    #[cfg(target_arch = "wasm32")]
+    let mut rng = OsRng;
+
+    let (_secret_key, public_key) = kem::key_gen(kem::Algorithm::Kyber768, &mut rng).unwrap();
 
     for _i in 0..100000 {
         let (_shared_secret, _ciphertext) =
-            kem::encapsulate(kem::Algorithm::Kyber768, &public_key, &mut drbg).unwrap();
+            kem::encapsulate(kem::Algorithm::Kyber768, &public_key, &mut rng).unwrap();
     }
 }
diff --git a/examples/kyber768_generate_keypair.rs b/examples/kyber768_generate_keypair.rs
@@ -1,11 +1,17 @@
-use libcrux::digest;
-use libcrux::drbg::Drbg;
 use libcrux::kem;
 
+#[cfg(not(target_arch = "wasm32"))]
+use libcrux::drbg;
+#[cfg(target_arch = "wasm32")]
+use rand_core::OsRng;
+
 fn main() {
-    let mut drbg = Drbg::new(digest::Algorithm::Sha256).unwrap();
+    #[cfg(not(target_arch = "wasm32"))]
+    let mut rng = drbg::Drbg::new(libcrux::digest::Algorithm::Sha256).unwrap();
+    #[cfg(target_arch = "wasm32")]
+    let mut rng = OsRng;
 
     for _i in 0..100000 {
-        let (_secret_key, _public_key) = kem::key_gen(kem::Algorithm::Kyber768, &mut drbg).unwrap();
+        let (_secret_key, _public_key) = kem::key_gen(kem::Algorithm::Kyber768, &mut rng).unwrap();
     }
 }
diff --git a/src/hacl.rs b/src/hacl.rs
@@ -13,7 +13,7 @@ pub(crate) mod aesgcm;
 pub(crate) mod blake2;
 pub(crate) mod chacha20_poly1305;
 pub(crate) mod curve25519;
-#[cfg(not(target_arch ="wasm32"))]
+#[cfg(not(target_arch = "wasm32"))]
 pub(crate) mod drbg;
 pub(crate) mod ed25519;
 pub(crate) mod hkdf;

diff --git a/tests/kyber768.rs b/tests/kyber768.rs
@@ -1,22 +1,30 @@
 use libcrux::{
     digest::{self, sha3_256, shake256},
-    drbg::{Drbg, RngCore},
     kem::{self, Algorithm},
 };
 
+#[cfg(not(target_arch = "wasm32"))]
+use libcrux::drbg::{self, RngCore};
+#[cfg(target_arch = "wasm32")]
+use rand_core::{OsRng, RngCore};
+
 const SHARED_SECRET_SIZE: usize = 32;
 const SECRET_KEY_SIZE: usize = 2400;
 const CIPHERTEXT_SIZE: u32 = 1088;
 
 const SECRET_KEY_REJECTION_VALUE_POSITION: usize = SECRET_KEY_SIZE - SHARED_SECRET_SIZE;
 
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
 #[test]
 fn consistency() {
-    let mut drbg = Drbg::new(digest::Algorithm::Sha256).unwrap();
+    #[cfg(not(target_arch = "wasm32"))]
+    let mut rng = drbg::Drbg::new(libcrux::digest::Algorithm::Sha256).unwrap();
+    #[cfg(target_arch = "wasm32")]
+    let mut rng = OsRng;
 
-    if let Ok((secret_key, public_key)) = kem::key_gen(Algorithm::Kyber768, &mut drbg) {
+    if let Ok((secret_key, public_key)) = kem::key_gen(Algorithm::Kyber768, &mut rng) {
         if let Ok((shared_secret, ciphertext)) =
-            kem::encapsulate(Algorithm::Kyber768, &public_key, &mut drbg)
+            kem::encapsulate(Algorithm::Kyber768, &public_key, &mut rng)
         {
             let shared_secret_decapsulated =
                 kem::decapsulate(Algorithm::Kyber768, &ciphertext, &secret_key).unwrap();
@@ -29,21 +37,25 @@ fn consistency() {
     // failing.
 }
 
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
 #[test]
 fn modified_ciphertext() {
-    let mut drbg = Drbg::new(digest::Algorithm::Sha256).unwrap();
+    #[cfg(not(target_arch = "wasm32"))]
+    let mut rng = drbg::Drbg::new(libcrux::digest::Algorithm::Sha256).unwrap();
+    #[cfg(target_arch = "wasm32")]
+    let mut rng = OsRng;
 
-    let random_u32 = drbg.next_u32();
+    let random_u32 = rng.next_u32();
     let mut random_byte: u8 = (random_u32 & 0xFF).try_into().unwrap();
     if random_byte == 0 {
         random_byte += 1;
     }
 
     let ciphertext_position: usize = (random_u32 % CIPHERTEXT_SIZE).try_into().unwrap();
 
-    if let Ok((secret_key, public_key)) = kem::key_gen(Algorithm::Kyber768, &mut drbg) {
+    if let Ok((secret_key, public_key)) = kem::key_gen(Algorithm::Kyber768, &mut rng) {
         if let Ok((shared_secret, mut ciphertext)) =
-            kem::encapsulate(Algorithm::Kyber768, &public_key, &mut drbg)
+            kem::encapsulate(Algorithm::Kyber768, &public_key, &mut rng)
         {
             ciphertext[ciphertext_position] ^= random_byte;
             let shared_secret_decapsulated =
@@ -69,11 +81,15 @@ fn compute_implicit_rejection_shared_secret(
     shake256(&to_hash)
 }
 
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
 #[test]
 fn modified_secret_key() {
-    let mut drbg = Drbg::new(digest::Algorithm::Sha256).unwrap();
+    #[cfg(not(target_arch = "wasm32"))]
+    let mut rng = drbg::Drbg::new(libcrux::digest::Algorithm::Sha256).unwrap();
+    #[cfg(target_arch = "wasm32")]
+    let mut rng = OsRng;
 
-    let random_u32 = drbg.next_u32();
+    let random_u32 = rng.next_u32();
 
     let mut random_byte: u8 = (random_u32 & 0xFF).try_into().unwrap();
     if random_byte == 0 {
@@ -82,9 +98,9 @@ fn modified_secret_key() {
 
     let secret_key_position: usize = ((random_u32 >> 8) % (SECRET_KEY_SIZE as u32 - 32)) as usize;
 
-    if let Ok((mut secret_key, public_key)) = kem::key_gen(Algorithm::Kyber768, &mut drbg) {
+    if let Ok((mut secret_key, public_key)) = kem::key_gen(Algorithm::Kyber768, &mut rng) {
         if let Ok((shared_secret, ciphertext)) =
-            kem::encapsulate(Algorithm::Kyber768, &public_key, &mut drbg)
+            kem::encapsulate(Algorithm::Kyber768, &public_key, &mut rng)
         {
             secret_key[secret_key_position] ^= random_byte;
             let shared_secret_decapsulated =
@@ -107,11 +123,15 @@ fn modified_secret_key() {
     // failing.
 }
 
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
 #[test]
 fn modified_ciphertext_and_implicit_rejection_value() {
-    let mut drbg = Drbg::new(digest::Algorithm::Sha256).unwrap();
+    #[cfg(not(target_arch = "wasm32"))]
+    let mut rng = drbg::Drbg::new(libcrux::digest::Algorithm::Sha256).unwrap();
+    #[cfg(target_arch = "wasm32")]
+    let mut rng = OsRng;
 
-    let random_u32 = drbg.next_u32();
+    let random_u32 = rng.next_u32();
 
     let mut random_byte_for_ciphertext: u8 = (random_u32 & 0xFF).try_into().unwrap();
     if random_byte_for_ciphertext == 0 {
@@ -120,7 +140,7 @@ fn modified_ciphertext_and_implicit_rejection_value() {
 
     let ciphertext_position: usize = ((random_u32 >> 8) % CIPHERTEXT_SIZE).try_into().unwrap();
 
-    let random_u32 = drbg.next_u32();
+    let random_u32 = rng.next_u32();
 
     let mut random_byte_for_secret_key: u8 = (random_u32 & 0xFF).try_into().unwrap();
     if random_byte_for_secret_key == 0 {
@@ -131,9 +151,9 @@ fn modified_ciphertext_and_implicit_rejection_value() {
         .try_into()
         .unwrap();
 
-    if let Ok((mut secret_key, public_key)) = kem::key_gen(Algorithm::Kyber768, &mut drbg) {
+    if let Ok((mut secret_key, public_key)) = kem::key_gen(Algorithm::Kyber768, &mut rng) {
         if let Ok((_, mut ciphertext)) =
-            kem::encapsulate(Algorithm::Kyber768, &public_key, &mut drbg)
+            kem::encapsulate(Algorithm::Kyber768, &public_key, &mut rng)
         {
             ciphertext[ciphertext_position] ^= random_byte_for_ciphertext;
             let shared_secret_decapsulated =

diff --git a/tests/p256.rs b/tests/p256.rs
@@ -1,8 +1,8 @@
 #[cfg(not(target_arch = "wasm32"))]
 use libcrux::drbg;
+use libcrux::ecdh::{self, key_gen};
 #[cfg(target_arch = "wasm32")]
 use rand_core::OsRng;
-use libcrux::ecdh::{self, key_gen};
 
 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
 #[test]