Fix padding on perfect multiples of the block size

contracts/StatefulSponge.sol

@@ -55,23 +55,36 @@ contract StatefulSponge {
         assembly {
             padded_ := mload(0x40)
 
+            // Grab the original length of `_data`
             let len := _data.length
-            let offset := _data.offset
+
+            let dataPtr := add(padded_, 0x20)
+            let endPtr := add(dataPtr, len)
 
             // Copy the data into memory.
-            calldatacopy(add(padded_, 0x20), offset, len)
+            calldatacopy(dataPtr, _data.offset, len)
+
+            let modBlockSize := mod(len, BLOCK_SIZE_BYTES)
+            switch modBlockSize
+            case false {
+                // If the input is a perfect multiple of the block size, then we add a full extra block of padding.
+                mstore8(endPtr, 0x01)
+                mstore8(sub(add(endPtr, BLOCK_SIZE_BYTES), 0x01), 0x80)
 
-            let padStart := mod(len, BLOCK_SIZE_BYTES)
-            if or(padStart, iszero(len)) {
-                let remaining := sub(BLOCK_SIZE_BYTES, padStart)
+                // Update the length of the data to include the padding.
+                mstore(padded_, add(len, BLOCK_SIZE_BYTES))
+            }
+            default {
+                // If the input is not a perfect multiple of the block size, then we add a partial block of padding.
+                // This should entail a set bit after the input, followed by as many zero bits as necessary to fill
+                // the block, followed by a single 1 bit in the lowest-order bit of the final block.
 
-                // Pad the input data
-                let dataPtr := add(padded_, 0x20)
+                let remaining := sub(BLOCK_SIZE_BYTES, modBlockSize)
                 let newLen := add(len, remaining)
 
                 // Store the padding bits.
-                mstore8(add(dataPtr, len), 0x01)
                 mstore8(add(dataPtr, sub(newLen, 0x01)), 0x80)
+                mstore8(endPtr, or(byte(0, mload(endPtr)), 0x01))
 
                 // Update the length of the data to include the padding. The length should be a multiple of the
                 // block size after this.

test/LibKeccak.t.sol

@@ -27,6 +27,29 @@ contract LibKeccak_Test is Test {
         assertEq(LibKeccak.squeeze(state), keccak256(new bytes(200)));
     }
 
+    function test_staticHashModuloBlockSize_success() public {
+        // Init
+        LibKeccak.StateMatrix memory state;
+
+        // Absorb 136 bytes into the sponge
+        bytes memory data = new bytes(136);
+        LibKeccak.absorb(state, data);
+        LibKeccak.permutation(state);
+
+        // Absorb another 136 bytes into the sponge
+        LibKeccak.absorb(state, data);
+        LibKeccak.permutation(state);
+
+        // Absorb the padding into the sponge. Because the input is a perfect multiple of the block size, the padding
+        // will be a full block.
+        data[135] = 0x80;
+        data[0] |= 0x01;
+        LibKeccak.absorb(state, data);
+        LibKeccak.permutation(state);
+
+        assertEq(LibKeccak.squeeze(state), keccak256(new bytes(136 * 2)));
+    }
+
     /// @dev Tests that the stateful sponge can absorb and squeeze an arbitrary amount of random data.
     function testFuzz_statefulSponge_success(bytes memory _data) public {
         vm.pauseGasMetering();