chore: apply keccak padding to blocks rather than bytes

Author: TomAFrench

src/keccak256.nr

@@ -33,9 +33,6 @@ pub fn keccak256<let N: u32>(input: [u8; N], message_size: u32) -> [u8; 32] {
     let max_blocks = (N + BLOCK_SIZE_IN_BYTES) / BLOCK_SIZE_IN_BYTES;
     let real_max_blocks = (message_size + BLOCK_SIZE_IN_BYTES) / BLOCK_SIZE_IN_BYTES;
 
-    // Apply Keccak padding (0x01 after message, 0x80 at block end)
-    apply_keccak_padding(&mut block_bytes, message_size, real_max_blocks);
-
     // populate a vector of 64-bit limbs from our byte array
     let mut sliced_buffer =
         [0; (((N / BLOCK_SIZE_IN_BYTES) + 1) * BLOCK_SIZE_IN_BYTES) / WORD_SIZE];
@@ -59,6 +56,9 @@ pub fn keccak256<let N: u32>(input: [u8; N], message_size: u32) -> [u8; 32] {
         sliced_buffer[i] = sliced as u64;
     }
 
+    // Apply Keccak padding (0x01 after message, 0x80 at block end)
+    apply_keccak_padding( &mut sliced_buffer, message_size, real_max_blocks);
+
     //2. sponge_absorb
     let mut state: [u64; NUM_KECCAK_LANES] = [0; NUM_KECCAK_LANES];
     // `real_max_blocks` is guaranteed to at least be `1`
@@ -126,24 +126,22 @@ pub fn keccak256<let N: u32>(input: [u8; N], message_size: u32) -> [u8; 32] {
     result
 }
 
-// Apply Keccak padding to the block_bytes array
+// Apply Keccak padding to the u64 block array
 // Append 0x01 after message, then 0x80 at end of block
 // If both padding bytes collide at the same byte, combine them as 0x81
 #[inline_always]
-pub(crate) fn apply_keccak_padding<let BLOCK_BYTES: u32>(
-    block_bytes: &mut [u8; BLOCK_BYTES],
+pub(crate) fn apply_keccak_padding<let N: u32>(
+     block_array: &mut [u64; N],
     message_size: u32,
     real_max_blocks: u32,
-) {
-    let real_blocks_bytes = real_max_blocks * BLOCK_SIZE_IN_BYTES;
-
-    if message_size == real_blocks_bytes - 1 {
-        // Combine both padding bits: 0x01 | 0x80 = 0x81
-        block_bytes[message_size] = 0x81;
-    } else {
-        block_bytes[message_size] = 0x01;
-        block_bytes[real_blocks_bytes - 1] = 0x80;
-    }
+)  {
+    // Calculate limb index and byte offset within the limb (little-endian)
+    let start_limb_index = message_size / WORD_SIZE;
+    let start_byte_offset = message_size % WORD_SIZE;
+    
+    block_array[start_limb_index] += 0x01 << (8 * (start_byte_offset as u64));        
+    // The end padding byte (0x80) always goes at byte 7 of the last limb
+    block_array[real_max_blocks * LIMBS_PER_BLOCK - 1] += 0x80 << 56;
 }
 
 comptime fn unroll_loop(start: u32, end: u32, body: fn(u32) -> Quoted) -> Quoted {

src/keccak256/tests.nr

@@ -3,52 +3,48 @@ use super::{apply_keccak_padding, keccak256};
 #[test]
 fn test_padding_no_collision() {
     // Message size 134: padding bytes don't collide
-    // Expected: block_bytes[134] = 0x01, block_bytes[135] = 0x80
-    let mut block_bytes = [0; 136];
-    let real_max_blocks = (134 + 136) / 136;
-    apply_keccak_padding(&mut block_bytes, 134, real_max_blocks);
+    // Limb 16, byte offsets 6 and 7
+    // block_array[16] = 0x01 << 48 (start) + 0x80 << 56 (end)
+    let mut block_array = [0; 17]; // 1 block = 17 limbs
+    let real_max_blocks = 1;
+    apply_keccak_padding(&mut block_array, 134, real_max_blocks);
 
-    assert_eq(block_bytes[134], 0x01, "First padding byte should be 0x01");
-    assert_eq(block_bytes[135], 0x80, "Second padding byte should be 0x80");
+    assert_eq(block_array[16], 0x8001000000000000);
 }
 
 #[test]
 fn test_padding_collision_135() {
-    // Message size 135: padding bytes collide at position 135
-    // Expected: block_bytes[135] = 0x81
-    let mut block_bytes = [0; 272];
-    let real_max_blocks = (135 + 136) / 136;
-    apply_keccak_padding(&mut block_bytes, 135, real_max_blocks);
-
-    assert_eq(block_bytes[135], 0x81, "Combined padding byte should be 0x81");
-    // Ensure no stray 0x80 elsewhere
-    for i in 136..272 {
-        assert_eq(block_bytes[i], 0x00, "No padding beyond position 135");
-    }
+    // Message size 135: padding bytes collide at byte 135 (offset 7 in limb 16)
+    // block_array[16] = 0x01 << 56 + 0x80 << 56 = 0x81 << 56
+    let mut block_array = [0; 17]; // 1 block = 17 limbs
+    let real_max_blocks = 1;
+    apply_keccak_padding(&mut block_array, 135, real_max_blocks);
+
+    assert_eq(block_array[16], 0x8100000000000000);
 }
 
 #[test]
 fn test_padding_zero_length() {
     // Message size 0: padding starts immediately
-    // Expected: block_bytes[0] = 0x01, block_bytes[135] = 0x80
-    let mut block_bytes = [0; 136];
-    let real_max_blocks = (0 + 136) / 136;
-    apply_keccak_padding(&mut block_bytes, 0, real_max_blocks);
+    // block_array[0] = 0x01, block_array[16] = 0x80 << 56
+    let mut block_array = [0; 17]; // 1 block = 17 limbs
+    let real_max_blocks = 1;
+    apply_keccak_padding(&mut block_array, 0, real_max_blocks);
 
-    assert_eq(block_bytes[0], 0x01, "First padding byte at position 0");
-    assert_eq(block_bytes[135], 0x80, "Second padding byte at end of block");
+    assert_eq(block_array[0], 0x01);
+    assert_eq(block_array[16], 0x8000000000000000);
 }
 
 #[test]
 fn test_padding_full_block() {
     // Message size 136: exactly fills one block, needs another for padding
-    // Expected: block_bytes[136] = 0x01, block_bytes[271] = 0x80
-    let mut block_bytes = [0; 272];
-    let real_max_blocks = (136 + 136) / 136;
-    apply_keccak_padding(&mut block_bytes, 136, real_max_blocks);
+    // block_array[17] = 0x01, block_array[33] = 0x80 << 56
+    let mut block_array = [0; 34]; // 2 blocks = 34 limbs
+    let real_max_blocks = 2;
+    apply_keccak_padding(&mut block_array, 136, real_max_blocks);
 
-    assert_eq(block_bytes[136], 0x01, "First padding byte after full block");
-    assert_eq(block_bytes[271], 0x80, "Second padding byte at end of next block");
+    assert_eq(block_array[17], 0x01);
+    assert_eq(block_array[33], 0x8000000000000000);
 }
 
 #[test]