diff --git a/.gitmodules b/.gitmodules index 888d42d..0f07815 100644 --- a/.gitmodules +++ b/.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 diff --git a/contracts/lib/LibKeccak.sol b/contracts/lib/LibKeccak.sol index 3f6df59..3669d0e 100644 --- a/contracts/lib/LibKeccak.sol +++ b/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. diff --git a/foundry.toml b/foundry.toml index e87b983..9e473bc 100644 --- a/foundry.toml +++ b/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 diff --git a/lib/solady b/lib/solady new file mode 160000 index 0000000..dcdddfd --- /dev/null +++ b/lib/solady @@ -0,0 +1 @@ +Subproject commit dcdddfd2b1b26353e96dbd75ec81db88f352384c diff --git a/test/LibKeccak.t.sol b/test/LibKeccak.t.sol index 64bf096..378d5ad 100644 --- a/test/LibKeccak.t.sol +++ b/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();