correct blocks

Author: jialinli98

src/sha224.nr

@@ -21,7 +21,7 @@ pub fn sha224_var<let N: u32>(msg: [u8; N], message_size: u64) -> sha224_constan
         let (mut h, mut msg_block) =
             process_full_blocks(msg, message_size, sha224_constants::INITIAL_STATE_SHA224);
 
-        let hash: HASH = finalize_sha256_blocks(msg, message_size, N, h, msg_block);
+        let hash: HASH = finalize_sha256_blocks::<N>(message_size, h, msg_block);
         // Convert 32-byte hash to 28-byte hash for SHA-224
         let mut hash_sha224: sha224_constants::HASH_SHA224 = [0; 28];
         for i in 0..28 {

src/sha256.nr

@@ -39,7 +39,7 @@ pub fn sha256_var<let N: u32>(msg: [u8; N], message_size: u64) -> HASH {
     } else {
         let (mut h, mut msg_block) = process_full_blocks(msg, message_size, INITIAL_STATE);
 
-        finalize_sha256_blocks(msg, message_size, N, h, msg_block)
+        finalize_sha256_blocks::<N>(message_size, h, msg_block)
     }
 }
 
@@ -119,14 +119,14 @@ pub(crate) fn process_full_blocks<let N: u32>(
     // Consider a message with an unknown number of bytes, `msg_size. It can be seen that this will have `msg_size / BLOCK_SIZE` full blocks.
     // - `states[i]` should then be the state after processing the first `i` blocks.
     // - `blocks[i]` should then be the next message block after processing the first `i` blocks.
+    // blocks[first_partially_filled_block_index] is the last block that is partially filled or all 0 if the message is a multiple of the block size.
     //
     // In other words:
     //
-    // blocks = [block 1, block 2, ..., block N / BLOCK_SIZE, empty block]
+    // blocks = [block 1, block 2, ..., block N / BLOCK_SIZE, block N / BLOCK_SIZE + 1]
     // states = [INITIAL_STATE, state after block 1, state after block 2, ..., state after block N / BLOCK_SIZE]
     //
     // We place the initial state in `states[0]` as in the case where the `message_size < BLOCK_SIZE` then there are no full blocks to process and no compressions should occur.
-    // Similarly, we place an empty block in `blocks[N/BLOCK_SIZE]` as in the case where the `message_size == N` then the padding bits will be written to an empty block.
     let mut blocks: [MSG_BLOCK; N / BLOCK_SIZE + 1] = std::mem::zeroed();
     let mut states: [STATE; N / BLOCK_SIZE + 1] = [h; N / BLOCK_SIZE + 1];
 
@@ -153,6 +153,19 @@ pub(crate) fn process_full_blocks<let N: u32>(
         states[i + 1] = sha256_compression(new_msg_block, states[i]);
     }
 
+    // If message_size/BLOCK_SIZE == N/BLOCK_SIZE, and message_size/BLOCK_SIZE has a remainder, we need to process the last block.
+    if message_size % BLOCK_SIZE != 0 {
+        let new_msg_block =
+            // Safety: separate verification function
+            unsafe { build_msg_block(msg, message_size, BLOCK_SIZE * num_blocks) };
+
+        // Verify the block we are compressing was appropriately constructed
+        verify_msg_block(msg, message_size, new_msg_block, BLOCK_SIZE * num_blocks);
+        verify_msg_block_zeros(new_msg_block, message_size % BLOCK_SIZE, INT_BLOCK_SIZE);
+
+        blocks[num_blocks] = new_msg_block;
+    }
+
     (states[first_partially_filled_block_index], blocks[first_partially_filled_block_index])
 }
 
@@ -241,13 +254,6 @@ fn verify_msg_block<let N: u32>(
     }
 }
 
-// Verify the block we are compressing was appropriately padded with zeros by `build_msg_block`.
-// This is only relevant for the last, potentially partially filled block.
-fn verify_msg_block_padding(msg_block: MSG_BLOCK, msg_byte_ptr: BLOCK_BYTE_PTR) {
-    // Check all the way to the end of the block.
-    verify_msg_block_zeros(msg_block, msg_byte_ptr, INT_BLOCK_SIZE);
-}
-
 // Verify that a region of ints in the message block are (partially) zeroed,
 // up to an (exclusive) maximum which can either be the end of the block
 // or just where the size is to be written.
@@ -512,33 +518,11 @@ fn hash_final_block(msg_block: MSG_BLOCK, mut state: STATE) -> HASH {
 }
 
 pub(crate) fn finalize_sha256_blocks<let N: u32>(
-    msg: [u8; N],
     message_size: u32,
-    total_len: u32,
     mut h: STATE,
     mut msg_block: MSG_BLOCK,
 ) -> HASH {
     let mut msg_byte_ptr = message_size % BLOCK_SIZE;
-    let modulo = total_len % BLOCK_SIZE;
-    // Handle setup of the final msg block.
-    // This case is only hit if the msg is less than the block size,
-    // or our message cannot be evenly split into blocks.
-    if modulo != 0 {
-        let num_blocks = total_len / BLOCK_SIZE;
-        let msg_start = BLOCK_SIZE * num_blocks;
-        let new_msg_block =
-            // Safety: separate verification function
-            unsafe { build_msg_block(msg, message_size, msg_start) };
-
-        if msg_start < message_size {
-            msg_block = new_msg_block;
-        }
-
-        verify_msg_block(msg, message_size, msg_block, msg_start);
-        if msg_start < message_size {
-            verify_msg_block_padding(msg_block, msg_byte_ptr);
-        }
-    }
 
     // If we had modulo == 0 then it means the last block was full,
     // and we can reset the pointer to zero to overwrite it.
@@ -630,7 +614,7 @@ pub fn partial_sha256_var_end<let N: u32>(
         }
     } else {
         let (mut h, mut msg_block) = process_full_blocks(msg, message_size, h);
-        finalize_sha256_blocks(msg, real_message_size, N, h, msg_block)
+        finalize_sha256_blocks::<N>(real_message_size, h, msg_block)
     }
 }