From 42d3a25a4e2d0e63cc83f7ed7959b705eecb0075 Mon Sep 17 00:00:00 2001
From: Daniel Reisinger <daniel@reisinger.xyz>
Date: Wed, 16 Aug 2023 20:58:50 +0200
Subject: [PATCH] k12: fix incorrect hash calculation (#499)

If the length of the merged input string is a multiple of the chunk size,
then the previous implementation appends a chaining value for an empty
chunk.
---
 k12/src/lib.rs   | 62 +++++++++++++++++++++++++++++-------------------
 k12/tests/mod.rs | 34 ++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 24 deletions(-)

diff --git a/k12/src/lib.rs b/k12/src/lib.rs
index a5f593c4a..7f1602ec8 100644
--- a/k12/src/lib.rs
+++ b/k12/src/lib.rs
@@ -52,6 +52,32 @@ impl<'cs> KangarooTwelveCore<'cs> {
             chain_length: 0usize,
         }
     }
+
+    fn process_chunk(&mut self) {
+        debug_assert!(self.bufpos == CHUNK_SIZE);
+        if self.chain_length == 0 {
+            self.final_tshk.update(&self.buffer);
+        } else {
+            self.process_chaining_chunk();
+        }
+
+        self.chain_length += 1;
+        self.buffer = [0u8; CHUNK_SIZE];
+        self.bufpos = 0;
+    }
+
+    fn process_chaining_chunk(&mut self) {
+        debug_assert!(self.bufpos != 0);
+        if self.chain_length == 1 {
+            self.final_tshk
+                .update(&[0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
+        }
+
+        let mut result = [0u8; CHAINING_VALUE_SIZE];
+        self.chain_tshk.update(&self.buffer[..self.bufpos]);
+        self.chain_tshk.finalize_xof_reset_into(&mut result);
+        self.final_tshk.update(&result);
+    }
 }
 
 impl HashMarker for KangarooTwelveCore<'_> {}
@@ -68,27 +94,12 @@ impl UpdateCore for KangarooTwelveCore<'_> {
     #[inline]
     fn update_blocks(&mut self, blocks: &[Block<Self>]) {
         for block in blocks {
-            self.buffer[self.bufpos..self.bufpos + 128].clone_from_slice(block);
-            self.bufpos += 128;
-
-            if self.bufpos != CHUNK_SIZE {
-                continue;
-            }
-
-            if self.chain_length == 0 {
-                self.final_tshk.update(&self.buffer);
-                self.final_tshk
-                    .update(&[0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
-            } else {
-                let mut result = [0u8; CHAINING_VALUE_SIZE];
-                self.chain_tshk.update(&self.buffer);
-                self.chain_tshk.finalize_xof_reset_into(&mut result);
-                self.final_tshk.update(&result);
+            if self.bufpos == CHUNK_SIZE {
+                self.process_chunk();
             }
 
-            self.chain_length += 1;
-            self.buffer = [0u8; CHUNK_SIZE];
-            self.bufpos = 0;
+            self.buffer[self.bufpos..self.bufpos + 128].clone_from_slice(block);
+            self.bufpos += 128;
         }
     }
 }
@@ -107,6 +118,10 @@ impl ExtendableOutputCore for KangarooTwelveCore<'_> {
             |block| self.update_blocks(block),
         );
 
+        if self.bufpos == CHUNK_SIZE && buffer.get_pos() != 0 {
+            self.process_chunk();
+        }
+
         // Read leftover data from buffer
         self.buffer[self.bufpos..(self.bufpos + buffer.get_pos())]
             .copy_from_slice(buffer.get_data());
@@ -114,17 +129,16 @@ impl ExtendableOutputCore for KangarooTwelveCore<'_> {
 
         // Calculate final node
         if self.chain_length == 0 {
-            // Input didnot exceed a single chaining value
+            // Input did not exceed a single chaining value
             let tshk = TurboShake128::from_core(<TurboShake128Core>::new(0x07))
                 .chain(&self.buffer[..self.bufpos])
                 .finalize_xof_reset();
             return KangarooTwelveReaderCore { tshk };
         }
+
         // Calculate last chaining value
-        let mut result = [0u8; CHAINING_VALUE_SIZE];
-        self.chain_tshk.update(&self.buffer[..self.bufpos]);
-        self.chain_tshk.finalize_xof_reset_into(&mut result);
-        self.final_tshk.update(&result);
+        self.process_chaining_chunk();
+
         // Pad final node calculation
         self.final_tshk
             .update(length_encode(self.chain_length, &mut lenbuf));
diff --git a/k12/tests/mod.rs b/k12/tests/mod.rs
index 92456e046..36b5b29b8 100644
--- a/k12/tests/mod.rs
+++ b/k12/tests/mod.rs
@@ -73,3 +73,37 @@ fn pat_c() {
         assert_eq!(result[..], expected[i as usize][..]);
     }
 }
+
+#[test]
+fn input_multiple_of_chunk_size_minus_one() {
+    // generated with reference python implementation
+    let expected = [
+        hex!("1b577636f723643e990cc7d6a659837436fd6a103626600eb8301cd1dbe553d6"),
+        hex!("e3ded52118ea64eaf04c7531c6ccb95e32924b7c2b87b2ce68ff2f2ee46e84ef"),
+        hex!("daacf62e434bdd126fbe9e61fae38d1429e9dddfaf8f999095585c3cbf366a4a"),
+        hex!("eac3722b4b7db10af973ed7ca60e113a19fab895b46476a9aac51ead099e6ba4"),
+    ];
+    for i in 0..expected.len() {
+        let len = 8192 * (i + 1) - 1;
+        let m: Vec<u8> = (0..len).map(|j| (j % 251) as u8).collect();
+        let result = digest_and_box(&m, 32);
+        assert_eq!(result[..], expected[i as usize][..]);
+    }
+}
+
+#[test]
+fn input_multiple_of_chunk_size() {
+    // generated with reference python implementation
+    let expected = [
+        hex!("48f256f6772f9edfb6a8b661ec92dc93b95ebd05a08a17b39ae3490870c926c3"),
+        hex!("82778f7f7234c83352e76837b721fbdbb5270b88010d84fa5ab0b61ec8ce0956"),
+        hex!("f4082a8fe7d1635aa042cd1da63bf235f91c231886c29896f9fe3818c60cd360"),
+        hex!("d14f8dc243c206004ca8a996997e5ae16a8bdda288f6c90d20d7c43c1a408618"),
+    ];
+    for i in 0..expected.len() {
+        let len = 8192 * (i + 1);
+        let m: Vec<u8> = (0..len).map(|j| (j % 251) as u8).collect();
+        let result = digest_and_box(&m, 32);
+        assert_eq!(result[..], expected[i as usize][..]);
+    }
+}