Fix memory edge case in pad & padMemory

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "lib/forge-std"]
 	path = lib/forge-std
 	url = https://github.com/foundry-rs/forge-std
+[submodule "lib/solady"]
+	path = lib/solady
+	url = https://github.com/vectorized/solady

contracts/lib/LibKeccak.sol

@@ -280,6 +280,14 @@ library LibKeccak {
             let modBlockSize := mod(len, BLOCK_SIZE_BYTES)
             switch modBlockSize
             case false {
+                // Clean the full padding block. It is possible that this memory is dirty, since solidity sometimes does
+                // not update the free memory pointer when allocating memory, for example with external calls.
+                mstore(endPtr, 0x00)
+                mstore(add(endPtr, 0x20), 0x00)
+                mstore(add(endPtr, 0x40), 0x00)
+                mstore(add(endPtr, 0x60), 0x00)
+                mstore(add(endPtr, 0x68), 0x00)
+
                 // 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)
@@ -294,10 +302,20 @@ library LibKeccak {
 
                 let remaining := sub(BLOCK_SIZE_BYTES, modBlockSize)
                 let newLen := add(len, remaining)
+                let paddedEndPtr := add(dataPtr, sub(newLen, 0x01))
+
+                // Clean the remainder to ensure that the intermediate data between the padding bits is 0. It is
+                // possible that this memory is dirty, since solidity sometimes does not update the free memory pointer
+                // when allocating memory, for example with external calls.
+                let ptr := endPtr
+                for { } lt(ptr, sub(add(paddedEndPtr, 0x01), 0x20)) { ptr := add(ptr, 0x20) } { mstore(ptr, 0x00) }
+                let partialRemainder := add(sub(paddedEndPtr, ptr), 0x01)
+                let final := sub(add(paddedEndPtr, 0x01), partialRemainder)
+                mstore(final, and(mload(final), not(shl(sub(256, shl(0x03, partialRemainder)), not(0x00)))))
 
                 // Store the padding bits.
                 mstore8(add(dataPtr, sub(newLen, 0x01)), 0x80)
-                mstore8(endPtr, or(byte(0, mload(endPtr)), 0x01))
+                mstore8(endPtr, or(byte(0x00, 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.
@@ -322,16 +340,24 @@ library LibKeccak {
 
             // Copy the data.
             let originalDataPtr := add(_data, 0x20)
-            for { let i := 0 } lt(i, len) { i := add(i, 0x20) } {
+            for { let i := 0x00 } lt(i, len) { i := add(i, 0x20) } {
                 mstore(add(dataPtr, i), mload(add(originalDataPtr, i)))
             }
 
             let modBlockSize := mod(len, BLOCK_SIZE_BYTES)
             switch modBlockSize
             case false {
+                // Clean the full padding block. It is possible that this memory is dirty, since solidity sometimes does
+                // not update the free memory pointer when allocating memory, for example with external calls.
+                mstore(endPtr, 0x00)
+                mstore(add(endPtr, 0x20), 0x00)
+                mstore(add(endPtr, 0x40), 0x00)
+                mstore(add(endPtr, 0x60), 0x00)
+                mstore(add(endPtr, 0x68), 0x00)
+
                 // 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)
+                mstore8(endPtr, 0x01)
 
                 // Update the length of the data to include the padding.
                 mstore(padded_, add(len, BLOCK_SIZE_BYTES))
@@ -343,10 +369,20 @@ library LibKeccak {
 
                 let remaining := sub(BLOCK_SIZE_BYTES, modBlockSize)
                 let newLen := add(len, remaining)
+                let paddedEndPtr := add(dataPtr, sub(newLen, 0x01))
+
+                // Clean the remainder to ensure that the intermediate data between the padding bits is 0. It is
+                // possible that this memory is dirty, since solidity sometimes does not update the free memory pointer
+                // when allocating memory, for example with external calls.
+                let ptr := endPtr
+                for { } lt(ptr, sub(add(paddedEndPtr, 0x01), 0x20)) { ptr := add(ptr, 0x20) } { mstore(ptr, 0x00) }
+                let partialRemainder := add(sub(paddedEndPtr, ptr), 0x01)
+                let final := sub(add(paddedEndPtr, 0x01), partialRemainder)
+                mstore(final, and(mload(final), not(shl(sub(256, shl(0x03, partialRemainder)), not(0x00)))))
 
                 // Store the padding bits.
-                mstore8(add(dataPtr, sub(newLen, 0x01)), 0x80)
-                mstore8(endPtr, or(byte(0, mload(endPtr)), 0x01))
+                mstore8(paddedEndPtr, 0x80)
+                mstore8(endPtr, or(byte(0x00, 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.

foundry.toml

@@ -4,9 +4,13 @@ out = "out"
 libs = ["lib"]
 optimizer_runs = 10_000_000
 evm_version = "shanghai"
+remappings = [
+  "@solady-test/=lib/solady/test/",
+  "@solady/=lib/solady/src/"
+]
 
 [fmt]
 bracket_spacing = true
 
 [fuzz]
-runs = 512
+runs = 256

test/LibKeccak.t.sol

@@ -2,12 +2,14 @@
 pragma solidity 0.8.15;
 
 import { Test, console2 as console } from "forge-std/Test.sol";
+import { TestPlus } from "@solady-test/utils/TestPlus.sol";
+import { LibString } from "@solady/utils/LibString.sol";
 
 import { LibKeccak } from "contracts/lib/LibKeccak.sol";
 import { StatefulSponge } from "contracts/StatefulSponge.sol";
 
-contract LibKeccak_Test is Test {
-    function test_staticHash_success() public {
+contract LibKeccak_Test is Test, TestPlus {
+    function test_staticHash_success() public brutalizeMemory {
         // Init
         LibKeccak.StateMatrix memory state;
 
@@ -27,7 +29,7 @@ contract LibKeccak_Test is Test {
         assertEq(LibKeccak.squeeze(state), keccak256(new bytes(200)));
     }
 
-    function test_staticHashModuloBlockSize_success() public {
+    function test_staticHashModuloBlockSize_success() public brutalizeMemory {
         // Init
         LibKeccak.StateMatrix memory state;
 
@@ -50,6 +52,64 @@ contract LibKeccak_Test is Test {
         assertEq(LibKeccak.squeeze(state), keccak256(new bytes(136 * 2)));
     }
 
+    /// @notice Tests the permutation end-to-end with brutalized memory. This ensures that the permutation does not have
+    ///         reliance on clean memory to function properly.
+    function testFuzz_hash_success(uint256 _numBytes) public brutalizeMemory {
+        _numBytes = bound(_numBytes, 0, 2 ** 10);
+
+        // Generate a pseudo-random preimage.
+        bytes memory data = new bytes(_numBytes);
+        for (uint256 i = 0; i < data.length; i++) {
+            data[i] = bytes1(uint8(_random()));
+        }
+
+        // Pad the data.
+        bytes memory paddedData = LibKeccak.padMemory(data);
+
+        // Hash the preimage.
+        LibKeccak.StateMatrix memory state;
+        for (uint256 i = 0; i < paddedData.length; i += LibKeccak.BLOCK_SIZE_BYTES) {
+            bytes memory kBlock = bytes(LibString.slice(string(paddedData), i, i + LibKeccak.BLOCK_SIZE_BYTES));
+            LibKeccak.absorb(state, kBlock);
+            LibKeccak.permutation(state);
+        }
+
+        // Assert that the hash is correct.
+        assertEq(LibKeccak.squeeze(state), keccak256(data));
+    }
+
+    /// @notice Tests that the `padCalldata` function does not write outside of the bounds of the input.
+    function testFuzz_padCalldata_memorySafety_succeeds(bytes calldata _in) public {
+        uint256 len = _in.length;
+        uint256 paddedLen = len % LibKeccak.BLOCK_SIZE_BYTES == 0
+            ? len + LibKeccak.BLOCK_SIZE_BYTES
+            : len + (LibKeccak.BLOCK_SIZE_BYTES - (len % LibKeccak.BLOCK_SIZE_BYTES));
+        uint64 freePtr;
+        assembly {
+            freePtr := mload(0x40)
+        }
+
+        // Pad memory should only write to memory in the range of [freePtr, freePtr + paddedLen + 0x20 (length word)]
+        vm.expectSafeMemory(freePtr, freePtr + uint64(paddedLen) + 0x20);
+        LibKeccak.pad(_in);
+    }
+
+    /// @notice Tests that the `padMemory` function does not write outside of the bounds of the input.
+    function testFuzz_padMemory_memorySafety_succeeds(bytes memory _in) public {
+        uint256 len = _in.length;
+        uint256 paddedLen = len % LibKeccak.BLOCK_SIZE_BYTES == 0
+            ? len + LibKeccak.BLOCK_SIZE_BYTES
+            : len + (LibKeccak.BLOCK_SIZE_BYTES - (len % LibKeccak.BLOCK_SIZE_BYTES));
+        uint64 freePtr;
+        assembly {
+            freePtr := mload(0x40)
+        }
+
+        // Pad memory should only write to memory in the range of [freePtr, freePtr + paddedLen + 0x20 (length word)]
+        vm.expectSafeMemory(freePtr, freePtr + uint64(paddedLen) + 0x20);
+        LibKeccak.padMemory(_in);
+    }
+
     /// @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();