test that xof_many doesn't write more blocks than requested

1 files changed, 24 insertions, 5 deletions

diff --git a/src/test.rs b/src/test.rs
index 6ce2e50..251a783 100644
--- a/src/test.rs
+++ b/src/test.rs
@@ -217,6 +217,12 @@ type XofManyFunction = unsafe fn(
 
 // A shared helper function for platform-specific tests.
 pub fn test_xof_many_fn(xof_many_function: XofManyFunction) {
+    let mut block = [0; BLOCK_LEN];
+    let block_len = 42;
+    crate::test::paint_test_input(&mut block[..block_len]);
+    let cv = [40, 41, 42, 43, 44, 45, 46, 47];
+    let flags = crate::KEYED_HASH;
+
     // Test a few different initial counter values.
     // - 0: The base case.
     // - u32::MAX: The low word of the counter overflows for all inputs except the first.
@@ -227,11 +233,6 @@ pub fn test_xof_many_fn(xof_many_function: XofManyFunction) {
         #[cfg(feature = "std")]
         dbg!(counter);
 
-        let mut block = [0; BLOCK_LEN];
-        let block_len = 42;
-        crate::test::paint_test_input(&mut block[..block_len]);
-        let cv = [40, 41, 42, 43, 44, 45, 46, 47];
-        let flags = crate::KEYED_HASH;
         // 31 (16 + 8 + 4 + 2 + 1) outputs
         const OUTPUT_SIZE: usize = 31 * BLOCK_LEN;
 
@@ -252,6 +253,24 @@ pub fn test_xof_many_fn(xof_many_function: XofManyFunction) {
 
         assert_eq!(portable_out, test_out);
     }
+
+    // Test that xof_many doesn't write more blocks than requested. Note that the current assembly
+    // implementation always outputs at least one block, so we don't test the zero case.
+    for block_count in 1..=32 {
+        let mut array = [0; BLOCK_LEN * 33];
+        let output_start = 17;
+        let output_len = block_count * BLOCK_LEN;
+        let output_end = output_start + output_len;
+        let output = &mut array[output_start..output_end];
+        unsafe {
+            xof_many_function(&cv, &block, block_len as u8, 0, flags, output);
+        }
+        for i in 0..array.len() {
+            if i < output_start || output_end <= i {
+                assert_eq!(0, array[i], "index {i}");
+            }
+        }
+    }
 }
 
 #[test]