test that xof_many doesn't write more blocks than requested

4 files changed, 64 insertions, 10 deletions

diff --git a/c/blake3_c_rust_bindings/src/test.rs b/c/blake3_c_rust_bindings/src/test.rs
index da88367..6a8e3e3 100644
--- a/c/blake3_c_rust_bindings/src/test.rs
+++ b/c/blake3_c_rust_bindings/src/test.rs
@@ -371,6 +371,12 @@ type XofManyFunction = unsafe extern "C" 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 = 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.
@@ -380,11 +386,6 @@ pub fn test_xof_many_fn(xof_many_function: XofManyFunction) {
     for counter in initial_counters {
         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 = KEYED_HASH;
         // 31 (16 + 8 + 4 + 2 + 1) outputs
         const OUTPUT_SIZE: usize = 31 * BLOCK_LEN;
 
@@ -416,6 +417,32 @@ 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.as_ptr(),
+                block.as_ptr(),
+                block_len as u8,
+                0,
+                flags,
+                output.as_mut_ptr(),
+                block_count,
+            );
+        }
+        for i in 0..array.len() {
+            if i < output_start || output_end <= i {
+                assert_eq!(0, array[i], "index {i}");
+            }
+        }
+    }
 }
 
 // Testing the portable implementation against itself is circular, but why not.
diff --git a/c/blake3_dispatch.c b/c/blake3_dispatch.c
index a90c757..cf5bad7 100644
--- a/c/blake3_dispatch.c
+++ b/c/blake3_dispatch.c
@@ -228,6 +228,10 @@ void blake3_xof_many(const uint32_t cv[8],
                      const uint8_t block[BLAKE3_BLOCK_LEN],
                      uint8_t block_len, uint64_t counter, uint8_t flags,
                      uint8_t out[64], size_t outblocks) {
+  if (outblocks == 0) {
+    // The current assembly implementation always outputs at least 1 block.
+    return;
+  }
 #if defined(IS_X86)
   const enum cpu_feature features = get_cpu_features();
 #if defined(__unix__) && !defined(BLAKE3_NO_AVX512)
diff --git a/src/platform.rs b/src/platform.rs
index 6e4ae25..3f47718 100644
--- a/src/platform.rs
+++ b/src/platform.rs
@@ -286,6 +286,10 @@ impl Platform {
         flags: u8,
         out: &mut [u8],
     ) {
+        if out.is_empty() {
+            // The current assembly implementation always outputs at least 1 block.
+            return;
+        }
         match self {
             // Safe because detect() checked for platform support.
             #[cfg(blake3_avx512_ffi)]
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]