Merge pull request #10 from dkg/align-tests-with-features

Align tests with features
integritychain · May 24, 2024 · 9f7e0df · 9f7e0df
2 parents a52d758 + ee24306
commit 9f7e0df
Show file tree

Hide file tree

Showing 7 changed files with 75 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -23,6 +23,7 @@ See <https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.203.ipd.pdf> for a full de
 The functionality is extremely simple to use, as demonstrated by the following example.
 
 ~~~rust
+# #[cfg(all(feature = "ml-kem-512", feature = "default-rng"))] {
 // Use the desired target parameter set.
 use fips203::ml_kem_512; // Could also be ml_kem_768 or ml_kem_1024. 
 use fips203::traits::{Decaps, Encaps, KeyGen, SerDes};
@@ -49,6 +50,7 @@ let alice_ssk_bytes = alice_dk.try_decaps(&alice_ct).unwrap();
 
 // Alice and Bob will now have the same secret key
 assert_eq!(bob_ssk_bytes, alice_ssk_bytes);
+# }
 ~~~
 
 The Rust [Documentation][docs-link] lives under each **Module** corresponding to the desired

diff --git a/src/traits.rs b/src/traits.rs
@@ -26,6 +26,7 @@ pub trait KeyGen {
     /// ```rust
     /// # use std::error::Error;
     /// # fn main() -> Result<(), Box<dyn Error>> {
+    /// # #[cfg(feature = "ml-kem-512")] {
     /// use fips203::ml_kem_512;                             // Could also be ml_kem_768 or ml_kem_1024.
     /// use fips203::traits::{KeyGen, SerDes, Decaps, Encaps};
     ///
@@ -44,6 +45,7 @@ pub trait KeyGen {
     /// let ssk1 = dk1.try_decaps(&ct1)?;                 // Party 1 runs decaps to generate the shared secret
     ///
     /// assert_eq!(ssk1, ssk2);                              // Each party has the same shared secret
+    /// # }
     /// # Ok(())}
     /// ```
     #[cfg(feature = "default-rng")]
@@ -61,6 +63,7 @@ pub trait KeyGen {
     /// ```rust
     /// # use std::error::Error;
     /// # fn main() -> Result<(), Box<dyn Error>> {
+    /// # #[cfg(feature = "ml-kem-512")] {
     /// use rand_core::OsRng;
     /// use fips203::ml_kem_512;                             // Could also be ml_kem_768 or ml_kem_1024.
     /// use fips203::traits::{KeyGen, SerDes, Decaps, Encaps};
@@ -71,7 +74,7 @@ pub trait KeyGen {
     /// let ek2_bytes = ek1_bytes;                           // Party 1 sends encaps bytes to party 2
     ///
     /// let ek2 = ml_kem_512::EncapsKey::try_from_bytes(ek2_bytes)?;  // Party 2 deserializes the encaps key
-    /// let (ssk2, ct2) = ek2.try_encaps()?;              // Party 2 generates shared secret and ciphertext
+    /// let (ssk2, ct2) = ek2.try_encaps_with_rng(&mut OsRng)?;       // Party 2 generates shared secret and ciphertext
     /// let ct2_bytes = ct2.into_bytes();                    // Party 2 serializes the ciphertext
     ///
     /// let ct1_bytes = ct2_bytes;                           // Party 2 sends the ciphertext to party 1
@@ -80,6 +83,7 @@ pub trait KeyGen {
     /// let ssk1 = dk1.try_decaps(&ct1)?;                 // Party 1 runs decaps to generate the shared secret
     ///
     /// assert_eq!(ssk1, ssk2);                              // Each party has the same shared secret
+    /// # }
     /// # Ok(())}
     /// ```
     fn try_keygen_with_rng(
@@ -94,6 +98,7 @@ pub trait KeyGen {
     /// ```rust
     /// # use std::error::Error;
     /// # fn main() -> Result<(), Box<dyn Error>> {
+    /// # #[cfg(feature = "ml-kem-512")] {
     /// use rand_core::OsRng;
     /// use fips203::ml_kem_512;                             // Could also be ml_kem_768 or ml_kem_1024.
     /// use fips203::traits::{KeyGen, SerDes, Decaps, Encaps};
@@ -103,6 +108,7 @@ pub trait KeyGen {
     /// let dk_bytes = dk.into_bytes();    // Serialize and perhaps store-then-restore decaps key
     /// assert!(ml_kem_512::KG::validate_keypair_vartime(&ek_bytes, &dk_bytes));  // Validate their correspondence
     ///
+    /// # }
     /// # Ok(())}
     /// ```
     fn validate_keypair_vartime(ek: &Self::EncapsByteArray, dk: &Self::DecapsByteArray) -> bool;
@@ -126,6 +132,7 @@ pub trait Encaps {
     /// ```rust
     /// # use std::error::Error;
     /// # fn main() -> Result<(), Box<dyn Error>> {
+    /// # #[cfg(feature = "ml-kem-512")] {
     /// use rand_core::OsRng;
     /// use fips203::ml_kem_512;                             // Could also be ml_kem_768 or ml_kem_1024.
     /// use fips203::traits::{KeyGen, SerDes, Decaps, Encaps};
@@ -145,6 +152,7 @@ pub trait Encaps {
     /// let ssk1 = dk1.try_decaps(&ct1)?;                 // Party 1 runs decaps to generate the shared secret
     ///
     /// assert_eq!(ssk1, ssk2);                              // Each party has the same shared secret
+    /// # }
     /// # Ok(())}
     /// ```
     #[cfg(feature = "default-rng")]
@@ -162,6 +170,7 @@ pub trait Encaps {
     /// ```rust
     /// # use std::error::Error;
     /// # fn main() -> Result<(), Box<dyn Error>> {
+    /// # #[cfg(feature = "ml-kem-512")] {
     /// use rand_core::OsRng;
     /// use fips203::ml_kem_512;                             // Could also be ml_kem_768 or ml_kem_1024.
     /// use fips203::traits::{KeyGen, SerDes, Decaps, Encaps};
@@ -181,6 +190,7 @@ pub trait Encaps {
     /// let ssk1 = dk1.try_decaps(&ct1)?;                 // Party 1 runs decaps to generate the shared secret
     ///
     /// assert_eq!(ssk1, ssk2);                              // Each party has the same shared secret
+    /// # }
     /// # Ok(())}
     /// ```
     fn try_encaps_with_rng(
@@ -205,6 +215,7 @@ pub trait Decaps {
     /// ```rust
     /// # use std::error::Error;
     /// # fn main() -> Result<(), Box<dyn Error>> {
+    /// # #[cfg(feature = "ml-kem-512")] {
     /// use rand_core::OsRng;
     /// use fips203::ml_kem_512;                             // Could also be ml_kem_768 or ml_kem_1024.
     /// use fips203::traits::{KeyGen, SerDes, Decaps, Encaps};
@@ -224,6 +235,7 @@ pub trait Decaps {
     /// let ssk1 = dk1.try_decaps(&ct1)?;                 // Party 1 runs decaps to generate the shared secret
     ///
     /// assert_eq!(ssk1, ssk2);                              // Each party has the same shared secret
+    /// # }
     /// # Ok(())}
     /// ```
     fn try_decaps(&self, ct: &Self::CipherText) -> Result<Self::SharedSecretKey, &'static str>;
@@ -241,6 +253,7 @@ pub trait SerDes {
     /// ```rust
     /// # use std::error::Error;
     /// # fn main() -> Result<(), Box<dyn Error>> {
+    /// # #[cfg(feature = "ml-kem-512")] {
     /// use rand_core::OsRng;
     /// use fips203::ml_kem_512;                             // Could also be ml_kem_768 or ml_kem_1024.
     /// use fips203::traits::{KeyGen, SerDes, Decaps, Encaps};
@@ -260,6 +273,7 @@ pub trait SerDes {
     /// let ssk1 = dk1.try_decaps(&ct1)?;                 // Party 1 runs decaps to generate the shared secret
     ///
     /// assert_eq!(ssk1, ssk2);                              // Each party has the same shared secret
+    /// # }
     /// # Ok(())}
     /// ```
     fn into_bytes(self) -> Self::ByteArray;
@@ -272,6 +286,7 @@ pub trait SerDes {
     /// ```rust
     /// # use std::error::Error;
     /// # fn main() -> Result<(), Box<dyn Error>> {
+    /// # #[cfg(feature = "ml-kem-512")] {
     /// use rand_core::OsRng;
     /// use fips203::ml_kem_512;                             // Could also be ml_kem_768 or ml_kem_1024.
     /// use fips203::traits::{KeyGen, SerDes, Decaps, Encaps};
@@ -291,6 +306,7 @@ pub trait SerDes {
     /// let ssk1 = dk1.try_decaps(&ct1)?;                 // Party 1 runs decaps to generate the shared secret
     ///
     /// assert_eq!(ssk1, ssk2);                              // Each party has the same shared secret
+    /// # }
     /// # Ok(())}
     /// ```
     fn try_from_bytes(ba: Self::ByteArray) -> Result<Self, &'static str>

diff --git a/tests/cctv_vectors/mod.rs b/tests/cctv_vectors/mod.rs
@@ -6,7 +6,12 @@ use hex::decode;
 use regex::Regex;
 
 use fips203::traits::{Decaps, Encaps, KeyGen, SerDes};
-use fips203::{ml_kem_1024, ml_kem_512, ml_kem_768};
+#[cfg(feature = "ml-kem-512")]
+use fips203::ml_kem_512;
+#[cfg(feature = "ml-kem-768")]
+use fips203::ml_kem_768;
+#[cfg(feature = "ml-kem-1024")]
+use fips203::ml_kem_1024;
 
 use super::TestRng;
 
@@ -42,6 +47,7 @@ fn get_intermediate_vec(
 }
 
 #[test]
+#[cfg(feature = "ml-kem-512")]
 pub fn test_intermediate_512() {
     let (d, z, ek_exp, dk_exp, m, k_exp, c_exp) =
         get_intermediate_vec("./tests/cctv_vectors/ML-KEM/intermediate/ML-KEM-512.txt");
@@ -60,6 +66,7 @@ pub fn test_intermediate_512() {
 }
 
 #[test]
+#[cfg(feature = "ml-kem-768")]
 pub fn test_intermediate_768() {
     let (d, z, ek_exp, dk_exp, m, k_exp, c_exp) =
         get_intermediate_vec("./tests/cctv_vectors/ML-KEM/intermediate/ML-KEM-768.txt");
@@ -78,6 +85,7 @@ pub fn test_intermediate_768() {
 }
 
 #[test]
+#[cfg(feature = "ml-kem-1024")]
 pub fn test_intermediate_1024() {
     let (d, z, ek_exp, dk_exp, m, k_exp, c_exp) =
         get_intermediate_vec("./tests/cctv_vectors/ML-KEM/intermediate/ML-KEM-1024.txt");
@@ -108,6 +116,7 @@ fn get_strcmp_vec(filename: &str) -> (Vec<u8>, Vec<u8>, Vec<u8>) {
 }
 
 #[test]
+#[cfg(feature = "ml-kem-512")]
 pub fn test_strcmp_512() {
     let (dk_exp, k_exp, c_exp) =
         get_strcmp_vec("./tests/cctv_vectors/ML-KEM/strcmp/ML-KEM-512.txt");
@@ -118,6 +127,7 @@ pub fn test_strcmp_512() {
 }
 
 #[test]
+#[cfg(feature = "ml-kem-768")]
 pub fn test_strcmp_768() {
     let (dk_exp, k_exp, c_exp) =
         get_strcmp_vec("./tests/cctv_vectors/ML-KEM/strcmp/ML-KEM-768.txt");
@@ -128,6 +138,7 @@ pub fn test_strcmp_768() {
 }
 
 #[test]
+#[cfg(feature = "ml-kem-1024")]
 pub fn test_strcmp_1024() {
     let (dk_exp, k_exp, c_exp) =
         get_strcmp_vec("./tests/cctv_vectors/ML-KEM/strcmp/ML-KEM-1024.txt");
@@ -138,6 +149,7 @@ pub fn test_strcmp_1024() {
 }
 
 #[test]
+#[cfg(feature = "ml-kem-512")]
 pub fn test_unlucky_512() {
     let (d, z, ek_exp, dk_exp, m, k_exp, c_exp) =
         get_intermediate_vec("./tests/cctv_vectors/ML-KEM/unluckysample/ML-KEM-512.txt");
@@ -156,6 +168,7 @@ pub fn test_unlucky_512() {
 }
 
 #[test]
+#[cfg(feature = "ml-kem-768")]
 pub fn test_unlucky_768() {
     let (d, z, ek_exp, dk_exp, m, k_exp, c_exp) =
         get_intermediate_vec("./tests/cctv_vectors/ML-KEM/unluckysample/ML-KEM-768.txt");
@@ -174,6 +187,7 @@ pub fn test_unlucky_768() {
 }
 
 #[test]
+#[cfg(feature = "ml-kem-1024")]
 pub fn test_unlucky_1024() {
     let (d, z, ek_exp, dk_exp, m, k_exp, c_exp) =
         get_intermediate_vec("./tests/cctv_vectors/ML-KEM/unluckysample/ML-KEM-1024.txt");
@@ -192,6 +206,7 @@ pub fn test_unlucky_1024() {
 }
 
 #[test]
+#[cfg(feature = "ml-kem-512")]
 fn test_modulus_512() {
     let gz = fs::read("./tests/cctv_vectors/ML-KEM/modulus/ML-KEM-512.txt.gz").unwrap();
     let mut d = GzDecoder::new(&gz[..]);
@@ -205,6 +220,7 @@ fn test_modulus_512() {
 }
 
 #[test]
+#[cfg(feature = "ml-kem-768")]
 fn test_modulus_768() {
     let gz = fs::read("./tests/cctv_vectors/ML-KEM/modulus/ML-KEM-768.txt.gz").unwrap();
     let mut d = GzDecoder::new(&gz[..]);
@@ -218,6 +234,7 @@ fn test_modulus_768() {
 }
 
 #[test]
+#[cfg(feature = "ml-kem-1024")]
 fn test_modulus_1024() {
     let gz = fs::read("./tests/cctv_vectors/ML-KEM/modulus/ML-KEM-1024.txt.gz").unwrap();
     let mut d = GzDecoder::new(&gz[..]);

diff --git a/tests/fails.rs b/tests/fails.rs
@@ -1,10 +1,12 @@
+#[cfg(feature = "ml-kem-512")]
 use fips203::ml_kem_512;
 use fips203::traits::{KeyGen, SerDes};
 use rand_chacha::rand_core::SeedableRng;
 use rand_core::RngCore;
 
 // Highlights potential validation opportunities
 #[test]
+#[cfg(feature = "ml-kem-512")]
 fn fails_512() {
     let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(123);
     for _i in 0..100 {

diff --git a/tests/integration.rs b/tests/integration.rs
@@ -1,9 +1,15 @@
 use fips203::traits::{Decaps, Encaps, KeyGen, SerDes};
-use fips203::{ml_kem_1024, ml_kem_512, ml_kem_768};
+#[cfg(feature = "ml-kem-512")]
+use fips203::ml_kem_512;
+#[cfg(feature = "ml-kem-768")]
+use fips203::ml_kem_768;
+#[cfg(feature = "ml-kem-1024")]
+use fips203::ml_kem_1024;
 use rand_chacha::rand_core::SeedableRng;
 
 
 #[test]
+#[cfg(feature = "ml-kem-512")]
 fn test_expected_flow_512() {
     let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(123);
     for _i in 0..100 {
@@ -37,6 +43,7 @@ fn test_expected_flow_512() {
 
 
 #[test]
+#[cfg(feature = "ml-kem-768")]
 fn test_expected_flow_768() {
     let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(456);
     for _i in 0..100 {
@@ -70,6 +77,7 @@ fn test_expected_flow_768() {
 
 
 #[test]
+#[cfg(feature = "ml-kem-1024")]
 fn test_expected_flow_1024() {
     let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(789);
     for _i in 0..100 {
@@ -105,6 +113,7 @@ fn test_expected_flow_1024() {
 // $ cargo test -- --ignored
 #[ignore]
 #[test]
+#[cfg(feature = "ml-kem-512")]
 fn test_forever() {
     let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(123);
     for i in 0..u64::MAX {

diff --git a/tests/native.rs b/tests/native.rs
@@ -1,10 +1,12 @@
+#[cfg(feature = "ml-kem-512")]
 use fips203::ml_kem_512;
 use fips203::traits::{Decaps, Encaps, KeyGen, SerDes};
 use hex_literal::hex;
 use rand_core::SeedableRng;
 
 
 #[test]
+#[cfg(feature = "ml-kem-512")]
 fn wasm_match() {
     let seed = 123; // Note: this should match the value giving in the web form
     let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(seed);