From eadcd95bdaaa38f18ec21e8687fc76da7d35d25e Mon Sep 17 00:00:00 2001
From: mriise <mark.riise26@gmail.com>
Date: Wed, 18 May 2022 15:43:57 -0700
Subject: [PATCH] Use blak3's XOF function for output sizes other than 32 fixes
 #213

---
 src/hasher_impl.rs | 23 ++++++++++++++++++-----
 tests/lib.rs       | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/src/hasher_impl.rs b/src/hasher_impl.rs
index 925e47c6..260376df 100644
--- a/src/hasher_impl.rs
+++ b/src/hasher_impl.rs
@@ -102,6 +102,14 @@ pub mod blake3 {
         digest: [u8; S],
     }
 
+    impl<const S: usize> Blake3Hasher<S> {
+        /// blake3's XOF function, fills the given slice with hash information
+        pub fn finalize_xof_fill(&mut self, digest_out: &mut [u8]) {
+            let mut digest = self.hasher.finalize_xof();
+            digest.fill(digest_out)
+        }
+    }
+
     impl<const S: usize> Default for Blake3Hasher<S> {
         fn default() -> Self {
             let hasher = ::blake3::Hasher::new();
@@ -119,11 +127,16 @@ pub mod blake3 {
         }
 
         fn finalize(&mut self) -> &[u8] {
-            let digest = self.hasher.finalize(); //default is 32 bytes anyway
-            let digest_bytes = digest.as_bytes();
-            let digest_out = &mut self.digest[..digest_bytes.len().max(S)];
-            digest_out.copy_from_slice(digest_bytes);
-            digest_out
+            // let digest = self.hasher.finalize(); //default is 32 bytes anyway
+            // let digest_bytes = digest.as_bytes();
+            // let size = digest_bytes.len().min(S);
+            // let digest_out = &mut self.digest[..digest_bytes.len().max(S)];
+            // digest_out.copy_from_slice(digest_bytes);
+            let Blake3Hasher { hasher, digest } = self;
+            let mut output = hasher.finalize_xof();
+            output.fill(digest);
+
+            digest
         }
 
         fn reset(&mut self) {
diff --git a/tests/lib.rs b/tests/lib.rs
index 3fdcba3b..f8416932 100644
--- a/tests/lib.rs
+++ b/tests/lib.rs
@@ -335,3 +335,40 @@ fn multihash_errors() {
         "Should error on wrong hash length"
     );
 }
+
+#[test]
+fn smaller_than_default_digest() {
+    use multihash::Blake3Hasher;
+    const DIGEST_SIZE: usize = 16;
+    pub struct ContentHasher(Blake3Hasher<DIGEST_SIZE>);
+
+    pub struct ContentHash([u8; DIGEST_SIZE]);
+
+    impl ContentHasher {
+        fn new() -> ContentHasher {
+            ContentHasher(Blake3Hasher::default())
+        }
+
+        fn write(&mut self, input: &[u8]) {
+            self.0.update(input);
+        }
+
+        fn finish(&mut self) -> ContentHash {
+            let hash = multihash::Code::Blake3_256.wrap(self.0.finalize()).unwrap();
+            let resized_hash = hash.resize::<DIGEST_SIZE>().unwrap();
+
+            let mut content = ContentHash([0u8; DIGEST_SIZE]);
+            content.0.copy_from_slice(resized_hash.digest());
+            content
+        }
+
+        fn reset(&mut self) {
+            self.0.reset();
+        }
+    }
+
+    let mut hasher = ContentHasher::new();
+    hasher.write("foobar".as_bytes());
+    let _content_hash = hasher.finish();
+    hasher.reset();
+}