diff --git a/c/blake3_c_rust_bindings/src/test.rs b/c/blake3_c_rust_bindings/src/test.rs index da883677..6a8e3e3e 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 a90c757c..cf5bad7e 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 6e4ae259..3f47718e 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 6ce2e509..251a7834 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]