From 8a7b62385bbca0bffed10d8b7c8b65a980e79711 Mon Sep 17 00:00:00 2001
From: Tjaden Hess <tjade273@gmail.com>
Date: Wed, 14 Aug 2024 21:52:30 -0500
Subject: [PATCH] Add E2E tests for sign-with-context and require alloc for
 KATs

---
 slh-dsa/src/lib.rs                  | 85 +++++++++--------------------
 slh-dsa/tests/acvp_keygen.rs        |  5 +-
 slh-dsa/tests/acvp_sig.rs           |  1 +
 slh-dsa/tests/acvp_ver.rs           |  2 +-
 slh-dsa/tests/known_answer_tests.rs |  1 +
 5 files changed, 31 insertions(+), 63 deletions(-)

diff --git a/slh-dsa/src/lib.rs b/slh-dsa/src/lib.rs
index 3bb9bb11..0887bb3b 100644
--- a/slh-dsa/src/lib.rs
+++ b/slh-dsa/src/lib.rs
@@ -80,6 +80,7 @@ mod tests {
     use super::*;
     use rand::Rng;
     use signature::*;
+    use util::macros::test_parameter_sets;
 
     fn test_sign_verify<P: ParameterSet>() {
         let mut rng = rand::thread_rng();
@@ -89,66 +90,7 @@ mod tests {
         let sig = sk.try_sign(msg).unwrap();
         vk.verify(msg, &sig).unwrap();
     }
-
-    #[test]
-    fn test_sign_verify_shake_128f() {
-        test_sign_verify::<Shake128f>();
-    }
-
-    #[test]
-    fn test_sign_verify_shake_128s() {
-        test_sign_verify::<Shake128s>();
-    }
-
-    #[test]
-    fn test_sign_verify_shake_192f() {
-        test_sign_verify::<Shake192f>();
-    }
-
-    #[test]
-    fn test_sign_verify_shake_192s() {
-        test_sign_verify::<Shake192s>();
-    }
-
-    #[test]
-    fn test_sign_verify_shake_256f() {
-        test_sign_verify::<Shake256f>();
-    }
-
-    #[test]
-    fn test_sign_verify_shake_256s() {
-        test_sign_verify::<Shake256s>();
-    }
-
-    #[test]
-    fn test_sign_verify_sha2_128f() {
-        test_sign_verify::<Sha2_128f>();
-    }
-
-    #[test]
-    fn test_sign_verify_sha2_128s() {
-        test_sign_verify::<Sha2_128s>();
-    }
-
-    #[test]
-    fn test_sign_verify_sha2_192f() {
-        test_sign_verify::<Sha2_192f>();
-    }
-
-    #[test]
-    fn test_sign_verify_sha2_192s() {
-        test_sign_verify::<Sha2_192s>();
-    }
-
-    #[test]
-    fn test_sign_verify_sha2_256f() {
-        test_sign_verify::<Sha2_256f>();
-    }
-
-    #[test]
-    fn test_sign_verify_sha2_256s() {
-        test_sign_verify::<Sha2_256s>();
-    }
+    test_parameter_sets!(test_sign_verify);
 
     // Check signature fails on modified message
     #[test]
@@ -212,4 +154,27 @@ mod tests {
             "Two successive randomized signatures over the same message should not be equal"
         );
     }
+
+    #[test]
+    fn test_sign_verify_nonempty_context() {
+        let mut rng = rand::thread_rng();
+        let sk = SigningKey::<Shake128f>::new(&mut rng);
+        let vk = sk.verifying_key();
+        let msg = b"Hello, world!";
+        let ctx = b"Test context";
+        let sig = sk.try_sign_with_context(msg, ctx, None).unwrap();
+        vk.try_verify_with_context(msg, ctx, &sig).unwrap();
+    }
+
+    #[test]
+    fn test_sign_verify_wrong_context() {
+        let mut rng = rand::thread_rng();
+        let sk = SigningKey::<Shake128f>::new(&mut rng);
+        let vk = sk.verifying_key();
+        let msg = b"Hello, world!";
+        let ctx = b"Test context!";
+        let wrong_ctx = b"Wrong context";
+        let sig = sk.try_sign_with_context(msg, ctx, None).unwrap();
+        assert!(vk.try_verify_with_context(msg, wrong_ctx, &sig).is_err());
+    }
 }
diff --git a/slh-dsa/tests/acvp_keygen.rs b/slh-dsa/tests/acvp_keygen.rs
index dafc849f..6f30be72 100644
--- a/slh-dsa/tests/acvp_keygen.rs
+++ b/slh-dsa/tests/acvp_keygen.rs
@@ -1,4 +1,5 @@
 #![allow(non_snake_case)]
+#![cfg(feature = "alloc")]
 
 use serde::Deserialize;
 use signature::Keypair;
@@ -39,8 +40,8 @@ macro_rules! parameter_case {
             &$test_case.pkSeed,
         );
         let vk = sk.verifying_key();
-        assert_eq!(sk.to_bytes().to_vec(), $test_case.sk);
-        assert_eq!(vk.to_bytes().to_vec(), $test_case.pk);
+        assert_eq!(sk.to_vec(), $test_case.sk);
+        assert_eq!(vk.to_vec(), $test_case.pk);
     }};
 }
 
diff --git a/slh-dsa/tests/acvp_sig.rs b/slh-dsa/tests/acvp_sig.rs
index 89655718..6170c814 100644
--- a/slh-dsa/tests/acvp_sig.rs
+++ b/slh-dsa/tests/acvp_sig.rs
@@ -1,4 +1,5 @@
 #![allow(non_snake_case)]
+#![cfg(feature = "alloc")]
 
 use serde::Deserialize;
 use slh_dsa::*;
diff --git a/slh-dsa/tests/acvp_ver.rs b/slh-dsa/tests/acvp_ver.rs
index e58ea887..3bb99ea2 100644
--- a/slh-dsa/tests/acvp_ver.rs
+++ b/slh-dsa/tests/acvp_ver.rs
@@ -1,9 +1,9 @@
 #![allow(non_snake_case)]
+#![cfg(feature = "alloc")]
 
 use serde::Deserialize;
 use slh_dsa::*;
 
-
 const KEYGEN_KAT_JSON: &str = include_str!("acvp/SLH-DSA-sigVer-FIPS205/internalProjection.json");
 
 #[derive(Deserialize, Debug)]
diff --git a/slh-dsa/tests/known_answer_tests.rs b/slh-dsa/tests/known_answer_tests.rs
index 2c378fcc..89c4b75a 100644
--- a/slh-dsa/tests/known_answer_tests.rs
+++ b/slh-dsa/tests/known_answer_tests.rs
@@ -3,6 +3,7 @@
 //! with PQCgenKAT_sign.c modified on line 59 to reduce iterations from 100 to 10
 //!
 //! These tests call the `slh_*_internal` functions directly, bypassing context processing.
+#![cfg(feature = "alloc")]
 use std::{array::from_fn, fmt::Write};
 
 use aes::Aes256;