fix: partial_sha256_var_end constrained/unconstrained divergence on non-block-aligned messages

src/sha224.nr

@@ -59,7 +59,7 @@ unconstrained fn __sha224_var<let N: u32>(
  * @return finalized sha24 hash
  */
 pub fn partial_sha224_var_end<let N: u32>(
-    mut h: [u32; 8],
+    h: [u32; 8],
     msg: [u8; N],
     message_size: u32,
     real_message_size: u32,

src/sha224/tests.nr

@@ -237,6 +237,14 @@ fn same_msg_len_variable_padding() {
     assert_eq(var_length_hash_576, var_full_length_hash);
 }
 
+// Regression test: partial_sha224_var_end delegates to partial_sha256_var_end, so
+// the same out-of-capacity guard must reject a tail that extends beyond `msg`.
+#[test(should_fail_with = "does not fit in the input array")]
+fn regression_partial_end_tail_beyond_capacity_sha224() {
+    let msg = [7; 64];
+    let _ = partial_sha224_var_end(sha224_constants::INITIAL_STATE_SHA224, msg, 64, 100);
+}
+
 #[test]
 fn fuzz_test_partial_hash_sha224(data: [u8; 192]) {
     let mut data0 = [0; 128];
@@ -271,6 +279,105 @@ fn fuzz_test_partial_hash_interstitial_sha224(data: [u8; 192]) {
     assert_eq(hash, correct_hash);
 }
 
+// Dual-mode coverage for the SHA-224 partial-hash API, mirroring the SHA-256
+// tests. The defect in partial_sha256_var_end is reachable verbatim through
+// partial_sha224_var_end, so these exercise the same non-block-aligned totals.
+
+unconstrained fn unconstrained_partial_sha224_var_end<let N: u32>(
+    h: [u32; 8],
+    msg: [u8; N],
+    message_size: u32,
+    real_message_size: u32,
+) -> sha224_constants::HASH_SHA224 {
+    partial_sha224_var_end(h, msg, message_size, real_message_size)
+}
+
+fn assert_partial_end_matches_reference_sha224(real_message_size: u32) {
+    let mut msg = [0; 192];
+    for i in 0..msg.len() {
+        msg[i] = ((i * 7 + 1) % 256) as u8;
+    }
+    let message_size = (real_message_size / 64) * 64;
+    let hash = partial_sha224_var_end(
+        sha224_constants::INITIAL_STATE_SHA224,
+        msg,
+        message_size,
+        real_message_size,
+    );
+    let reference = sha224_var(msg, real_message_size);
+    assert_eq(hash, reference);
+}
+
+#[test]
+fn test_partial_end_sha224_size_1() {
+    assert_partial_end_matches_reference_sha224(1);
+}
+
+#[test]
+fn test_partial_end_sha224_size_63() {
+    assert_partial_end_matches_reference_sha224(63);
+}
+
+#[test]
+fn test_partial_end_sha224_size_65() {
+    assert_partial_end_matches_reference_sha224(65);
+}
+
+#[test]
+fn test_partial_end_sha224_size_100() {
+    assert_partial_end_matches_reference_sha224(100);
+}
+
+#[test]
+fn test_partial_end_sha224_size_127() {
+    assert_partial_end_matches_reference_sha224(127);
+}
+
+#[test]
+fn test_partial_end_sha224_size_191() {
+    assert_partial_end_matches_reference_sha224(191);
+}
+
+#[test]
+fn test_partial_end_sha224_size_64() {
+    assert_partial_end_matches_reference_sha224(64);
+}
+
+#[test]
+fn fuzz_partial_end_sha224_implementations_agree(msg: [u8; 192], real_message_size: u32) {
+    let real = real_message_size % 193; // 0..=192
+    let message_size = (real / 64) * 64;
+    // Safety: test function
+    let oracle = unsafe {
+        unconstrained_partial_sha224_var_end(
+            sha224_constants::INITIAL_STATE_SHA224,
+            msg,
+            message_size,
+            real,
+        )
+    };
+    let circuit = partial_sha224_var_end(
+        sha224_constants::INITIAL_STATE_SHA224,
+        msg,
+        message_size,
+        real,
+    );
+    assert_eq(circuit, oracle);
+}
+
+#[test]
+fn fuzz_partial_end_sha224_matches_reference(msg: [u8; 192], real_message_size: u32) {
+    let real = real_message_size % 193; // 0..=192
+    let message_size = (real / 64) * 64;
+    let hash = partial_sha224_var_end(
+        sha224_constants::INITIAL_STATE_SHA224,
+        msg,
+        message_size,
+        real,
+    );
+    assert_eq(hash, sha224_var(msg, real));
+}
+
 #[test]
 fn fuzz_test_partial_hash_end_sha224(data: [u8; 192]) {
     let mut data0 = [0; 64];

src/sha256.nr

@@ -266,11 +266,7 @@ fn add_padding_byte_and_compress_if_needed(
     }
 }
 
-pub(crate) fn finalize_sha256_blocks(
-    message_size: u32,
-    mut h: STATE,
-    mut msg_block: MSG_BLOCK,
-) -> HASH {
+pub(crate) fn finalize_sha256_blocks(message_size: u32, h: STATE, msg_block: MSG_BLOCK) -> HASH {
     let msg_byte_ptr = message_size % BLOCK_SIZE;
 
     let (h, mut msg_block) = add_padding_byte_and_compress_if_needed(msg_block, msg_byte_ptr, h);
@@ -291,11 +287,12 @@ pub(crate) fn finalize_sha256_blocks(
  * @return the intermediate hash state after compressing in msg to h
  */
 pub fn partial_sha256_var_interstitial<let N: u32>(
-    mut h: [u32; 8],
+    h: [u32; 8],
     msg: [u8; N],
     message_size: u32,
 ) -> [u32; 8] {
     assert(message_size % BLOCK_SIZE == 0, "Message size must be a multiple of the block size");
+    assert(message_size <= N, "Message does not fit in the input array");
     if std::runtime::is_unconstrained() {
         // Safety: running as an unconstrained function
         unsafe {
@@ -326,6 +323,15 @@ pub fn partial_sha256_var_end<let N: u32>(
     real_message_size: u32,
 ) -> [u8; 32] {
     assert(message_size % BLOCK_SIZE == 0, "Message size must be a multiple of the block size");
+    assert(
+        real_message_size >= message_size,
+        "Real message size must be at least the message size",
+    );
+    // The block-aligned prefix plus the trailing partial block must lie within `msg`.
+    assert(
+        message_size + real_message_size % BLOCK_SIZE <= N,
+        "Message does not fit in the input array",
+    );
     if std::runtime::is_unconstrained() {
         // Safety: running as an unconstrained function
         unsafe {
@@ -335,10 +341,12 @@ pub fn partial_sha256_var_end<let N: u32>(
             // This case is only hit if the msg is less than the block size,
             // or our message cannot be evenly split into blocks.
 
-            finalize_last_sha256_block(h, real_message_size, msg)
+            finalize_last_sha256_block(h, message_size, real_message_size, msg)
         }
     } else {
-        let (h, msg_block) = process_full_blocks(msg, message_size, h);
+        // The final chunk is the block-aligned `message_size` plus the trailing partial block.
+        let chunk_size = message_size + (real_message_size % BLOCK_SIZE);
+        let (h, msg_block) = process_full_blocks(msg, chunk_size, h);
         finalize_sha256_blocks(real_message_size, h, msg_block)
     }
 }
@@ -358,19 +366,26 @@ unconstrained fn __sha_partial_var_interstitial<let N: u32>(
     h
 }
 
-// Helper function to finalize the message block with padding and length
+// Helper function to finalize the message block with padding and length.
+//
+// The trailing partial block is read from `msg` starting at `partial_message_size` (the
+// block-aligned bytes already compressed into `h`), while `real_message_size` supplies the
+// length encoded into the padding.
 unconstrained fn finalize_last_sha256_block<let N: u32>(
-    mut h: STATE,
-    message_size: u32,
+    h: STATE,
+    partial_message_size: u32,
+    real_message_size: u32,
     msg: [u8; N],
 ) -> HASH {
-    let msg_byte_ptr = message_size % BLOCK_SIZE;
+    let msg_byte_ptr = real_message_size % BLOCK_SIZE;
 
-    // We now build the final un-filled block.
+    // We now build the final un-filled block from the trailing bytes.
     let msg_block: MSG_BLOCK = if msg_byte_ptr != 0 {
-        let num_full_blocks = message_size / BLOCK_SIZE;
-        let msg_start = BLOCK_SIZE * num_full_blocks;
-        build_msg_block(msg, message_size, msg_start)
+        build_msg_block(
+            msg,
+            partial_message_size + msg_byte_ptr,
+            partial_message_size,
+        )
     } else {
         // If the message size is a multiple of the block size (i.e. `msg_byte_ptr == 0`) then this block will be empty,
         // so we short-circuit in this case.
@@ -379,7 +394,7 @@ unconstrained fn finalize_last_sha256_block<let N: u32>(
 
     // Once built, we need to add the necessary padding bytes and encoded length
     let (h, mut msg_block) = add_padding_byte_and_compress_if_needed(msg_block, msg_byte_ptr, h);
-    msg_block = attach_len_to_msg_block(msg_block, message_size);
+    msg_block = attach_len_to_msg_block(msg_block, real_message_size);
 
     hash_final_block(msg_block, h)
 }