From ecd9b419a26dcd0a68ebd56e68b8feb46cdc6a68 Mon Sep 17 00:00:00 2001
From: Ruediger Klaehn <rklaehn@protonmail.com>
Date: Tue, 1 Aug 2023 23:12:50 +0300
Subject: [PATCH] fix issue with single chunk sized block

also add a little brute force test
---
 src/guts.rs | 38 ++++++++++++++++++++++++++++++++++++--
 src/lib.rs  | 19 +++++--------------
 src/test.rs | 15 +++++++++------
 3 files changed, 50 insertions(+), 22 deletions(-)

diff --git a/src/guts.rs b/src/guts.rs
index 874c20ed4..6d1c73812 100644
--- a/src/guts.rs
+++ b/src/guts.rs
@@ -23,7 +23,7 @@ fn is_subtree(start_chunk: u64, len: u64) -> bool {
 /// whether the subtree is the root of the tree.
 ///
 /// Subtrees that start at a non zero chunk can not be the root.
-pub fn hash_block(start_chunk: u64, data: &[u8], is_root: bool) -> crate::Hash {
+pub fn hash_subtree(start_chunk: u64, data: &[u8], is_root: bool) -> crate::Hash {
     debug_assert!(is_subtree(start_chunk, data.len() as u64));
     debug_assert!(start_chunk == 0 || !is_root);
     let mut hasher = crate::Hasher::new_with_start_chunk(start_chunk);
@@ -133,9 +133,43 @@ mod test {
 
     #[test]
     fn test_hash_block() {
-        assert_eq!(crate::hash(b"foo"), hash_block(0, b"foo", true));
+        assert_eq!(crate::hash(b"foo"), hash_subtree(0, b"foo", true));
 
         assert_eq!(is_subtree(4, 1024 * 4 - 1), true);
         assert_eq!(is_subtree(1, 1024 * 4), false);
+
+        fn recursive_hash_block(start_chunk: u64, data: &[u8], is_root: bool) -> crate::Hash {
+            if data.len() <= CHUNK_LEN {
+                let mut hasher = ChunkState::new(start_chunk);
+                hasher.update(data);
+                hasher.finalize(is_root)
+            } else {
+                let chunks = data.len() / CHUNK_LEN + (data.len() % CHUNK_LEN != 0) as usize;
+                let chunks = chunks.next_power_of_two();
+                let mid = chunks / 2;
+                let mid_bytes = mid * CHUNK_LEN;
+                let left = recursive_hash_block(start_chunk, &data[..mid_bytes], false);
+                let right =
+                    recursive_hash_block(start_chunk + mid as u64, &data[mid_bytes..], false);
+                parent_cv(&left, &right, is_root)
+            }
+        }
+
+        let data = (0..1024 << 4).map(|i| i as u8).collect::<Vec<_>>();
+        for block_log in 0..4 {
+            let block_size = 1usize << block_log;
+            let block_size_u64 = block_size as u64;
+            for i in 0..100 {
+                let start_chunk = i * block_size_u64;
+                assert_eq!(
+                    recursive_hash_block(start_chunk, &data[..CHUNK_LEN], false),
+                    hash_subtree(start_chunk, &data[..CHUNK_LEN], false)
+                );
+                assert_eq!(
+                    recursive_hash_block(start_chunk, &data[..block_size * CHUNK_LEN], false),
+                    hash_subtree(start_chunk, &data[..block_size * CHUNK_LEN], false)
+                );
+            }
+        }
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 71562f8c8..d3b0a33a5 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -70,7 +70,6 @@
 #[cfg(feature = "zeroize")]
 extern crate zeroize_crate as zeroize; // Needed because `zeroize::Zeroize` assumes the crate is named `zeroize`.
 
-
 #[cfg(test)]
 mod test;
 
@@ -409,18 +408,6 @@ impl Output {
         Hash(platform::le_bytes_from_words_32(&cv))
     }
 
-    fn hash(&self, is_root: bool) -> Hash {
-        // debug_assert_eq!(self.counter, 0);
-        let mut cv = self.input_chaining_value;
-        let mut flags = self.flags;
-        if is_root {
-            flags |= ROOT;
-        }
-        self.platform
-            .compress_in_place(&mut cv, &self.block, self.block_len, 0, flags);
-        Hash(platform::le_bytes_from_words_32(&cv))
-    }
-
     fn root_output_block(&self) -> [u8; 2 * OUT_LEN] {
         self.platform.compress_xof(
             &self.input_chaining_value,
@@ -1327,7 +1314,11 @@ impl Hasher {
 
     fn finalize_node(&self, is_root: bool) -> Hash {
         let output = self.final_output();
-        output.hash(is_root)
+        if is_root {
+            output.root_hash()
+        } else {
+            output.chaining_value().into()
+        }
     }
 
     /// Finalize the hash state and return an [`OutputReader`], which can
diff --git a/src/test.rs b/src/test.rs
index 0d94f4465..d8a4f19a8 100644
--- a/src/test.rs
+++ b/src/test.rs
@@ -658,11 +658,13 @@ fn test_zeroize() {
     assert_eq!(hasher.chunk_state.buf_len, 0);
     assert_eq!(hasher.chunk_state.blocks_compressed, 0);
     assert_eq!(hasher.chunk_state.flags, 0);
-    assert!(matches!(hasher.chunk_state.platform, crate::Platform::Portable));
+    assert!(matches!(
+        hasher.chunk_state.platform,
+        crate::Platform::Portable
+    ));
     assert_eq!(hasher.key, [0; 8]);
     assert_eq!(&*hasher.cv_stack, &[[0u8; 32]; 0]);
 
-
     let mut output_reader = crate::OutputReader {
         inner: crate::Output {
             input_chaining_value: [42; 8],
@@ -675,14 +677,15 @@ fn test_zeroize() {
         position_within_block: 42,
     };
 
-
     output_reader.zeroize();
     assert_eq!(output_reader.inner.input_chaining_value, [0; 8]);
     assert_eq!(output_reader.inner.block, [0; 64]);
     assert_eq!(output_reader.inner.counter, 0);
     assert_eq!(output_reader.inner.block_len, 0);
     assert_eq!(output_reader.inner.flags, 0);
-    assert!(matches!(output_reader.inner.platform, crate::Platform::Portable));
+    assert!(matches!(
+        output_reader.inner.platform,
+        crate::Platform::Portable
+    ));
     assert_eq!(output_reader.position_within_block, 0);
-
-}
\ No newline at end of file
+}