Skip to content

Commit 94c2497

Browse files
0xClandestine0xClandestine
committed
sha
1 parent be34d3d commit 94c2497

File tree

4 files changed

+75
-52
lines changed

4 files changed

+75
-52
lines changed

‎foundry.toml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,9 @@
8080
no_match_path = "script/**/*.sol"
8181
fs_permissions = [{ access = "read-write", path = "./"}]
8282

83+
# If enabled, allows internal expectRevert calls to be used in tests.
84+
allow_internal_expect_revert = true
85+
8386
[profile.default.fmt]
8487
# Single-line vs multi-line statement blocks
8588
single_line_statement_blocks = "preserve" # Options: "single", "multi", "preserve"

‎lib/forge-std‎

‎src/test/unit/libraries/Merkle.t.sol‎

Lines changed: 66 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -6,41 +6,46 @@ import "src/contracts/libraries/Merkle.sol";
66
import "src/test/utils/Murky.sol";
77

88
abstract contract MerkleBaseTest is Test, MurkyBase {
9-
bool usingSha; // Whether to use Keccak or Sha256 for tree + proof generation.
109
bytes32[] leaves; // The contents of the merkle tree (unsorted).
1110
bytes32 root; // The root of the merkle tree.
1211
bytes[] proofs; // The proofs for each leaf in the tree.
1312

13+
function setUp() public {
14+
leaves = _genLeaves(vm.randomBool() ? 9 : 10);
15+
proofs = _genProofs(leaves);
16+
root = _genRoot(leaves);
17+
}
18+
1419
/// -----------------------------------------------------------------------
1520
/// Keccak + Sha256 Tests
1621
/// -----------------------------------------------------------------------
1722

1823
/// @notice Verifies that (Murky's) proofs are compatible with our implementation.
1924
function testFuzz_verifyInclusion_ValidProof(uint) public {
20-
checkAllProofs(true);
25+
_checkAllProofs(true);
2126
}
2227

2328
/// @notice Verifies that an empty proof(s) is invalid.
2429
function testFuzz_verifyInclusion_EmptyProofs(uint) public {
2530
proofs = new bytes[](proofs.length);
26-
checkAllProofs(false);
31+
_checkAllProofs(false);
2732
}
2833

2934
/// @notice Verifies valid proofs cannot be used to prove invalid leaves.
3035
function testFuzz_verifyInclusion_WrongProofs(uint) public {
3136
bytes memory proof0 = proofs[0];
3237
bytes memory proof1 = proofs[1];
3338
(proofs[0], proofs[1]) = (proof1, proof0);
34-
checkSingleProof(false, 0);
35-
checkSingleProof(false, 1);
39+
_checkSingleProof(false, 0);
40+
_checkSingleProof(false, 1);
3641
}
3742

3843
/// @notice Verifies that a valid proof with excess data appended is invalid.
3944
function testFuzz_verifyInclusion_ExcessProofLength(uint) public {
4045
unchecked {
4146
proofs[0] = abi.encodePacked(proofs[0], vm.randomBytes(vm.randomUint(1, 10) * 32));
4247
}
43-
checkSingleProof(false, 0);
48+
_checkSingleProof(false, 0);
4449
}
4550

4651
/// @notice Verifies that a valid proof with a manipulated word is invalid.
@@ -53,35 +58,54 @@ abstract contract MerkleBaseTest is Test, MurkyBase {
5358
mstore(m, manipulated)
5459
}
5560
proofs[0] = proof;
56-
checkSingleProof(false, 0);
61+
_checkSingleProof(false, 0);
5762
}
5863

5964
/// -----------------------------------------------------------------------
6065
/// Assertions
6166
/// -----------------------------------------------------------------------
6267

6368
/// @dev Checks that all proofs are valid for their respective leaves.
64-
function checkAllProofs(bool status) internal virtual {
69+
function _checkAllProofs(bool status) internal virtual {
6570
function (bytes memory proof, bytes32 root, bytes32 leaf, uint256 index) returns (bool) verifyInclusion =
66-
usingSha ? Merkle.verifyInclusionSha256 : Merkle.verifyInclusionKeccak;
71+
usingSha() ? Merkle.verifyInclusionSha256 : Merkle.verifyInclusionKeccak;
6772
for (uint i = 0; i < leaves.length; ++i) {
68-
assertEq(verifyInclusion(proofs[i], root, leaves[i], i), status);
73+
if (proofs[i].length == 0) {
74+
vm.expectRevert(Merkle.InvalidProofLength.selector);
75+
verifyInclusion(proofs[i], root, leaves[i], i);
76+
} else {
77+
assertEq(verifyInclusion(proofs[i], root, leaves[i], i), status);
78+
}
6979
}
7080
}
7181

7282
/// @dev Checks that a single proof is valid for its respective leaf.
73-
function checkSingleProof(bool status, uint index) internal virtual {
74-
function (bytes memory proof, bytes32 root, bytes32 leaf, uint256 index) returns (bool) verifyInclusion =
75-
usingSha ? Merkle.verifyInclusionSha256 : Merkle.verifyInclusionKeccak;
76-
assertEq(verifyInclusion(proofs[index], root, leaves[index], index), status);
83+
function _checkSingleProof(bool status, uint index) internal virtual {
84+
function (bytes memory proof, bytes32 root, bytes32 leaf, uint256 index) view returns (bool) verifyInclusion =
85+
usingSha() ? Merkle.verifyInclusionSha256 : Merkle.verifyInclusionKeccak;
86+
if (proofs[index].length == 0) {
87+
vm.expectRevert(Merkle.InvalidProofLength.selector);
88+
verifyInclusion(proofs[index], root, leaves[index], index);
89+
} else {
90+
assertEq(verifyInclusion(proofs[index], root, leaves[index], index), status);
91+
}
7792
}
7893

7994
/// -----------------------------------------------------------------------
8095
/// Helpers
8196
/// -----------------------------------------------------------------------
8297

98+
/// @dev Efficiently pads the length of leaves to the next power of 2 by appending zeros.
99+
function _padLeaves(bytes32[] memory leaves) internal view virtual returns (bytes32[] memory paddedLeaves) {
100+
uint numLeaves = _roundUpPow2(leaves.length);
101+
paddedLeaves = new bytes32[](numLeaves);
102+
for (uint i = 0; i < leaves.length; ++i) {
103+
paddedLeaves[i] = leaves[i];
104+
}
105+
}
106+
83107
/// @dev Effeciently generates a random list of leaves without iterative hashing.
84-
function getLeaves(uint numLeaves) internal view virtual returns (bytes32[] memory leaves) {
108+
function _genLeaves(uint numLeaves) internal view virtual returns (bytes32[] memory leaves) {
85109
bytes memory _leavesAsBytes = vm.randomBytes(numLeaves * 32);
86110
/// @solidity memory-safe-assembly
87111
assembly {
@@ -91,23 +115,26 @@ abstract contract MerkleBaseTest is Test, MurkyBase {
91115
}
92116

93117
/// @dev Generates proofs for each leaf in the tree.
94-
/// Intended to be overridden by the below child contracts.
95-
function getProofs(bytes32[] memory leaves) public view virtual returns (bytes[] memory proofs);
96-
}
118+
function _genProofs(bytes32[] memory leaves) internal view virtual returns (bytes[] memory proofs) {
119+
uint numLeaves = _roundUpPow2(leaves.length);
120+
bytes32[] memory paddedLeaves = _padLeaves(leaves);
121+
proofs = new bytes[](leaves.length);
122+
for (uint i = 0; i < leaves.length; ++i) {
123+
proofs[i] = abi.encodePacked(getProof(paddedLeaves, i));
124+
}
125+
}
97126

98-
contract MerkleKeccakTest is MerkleBaseTest, MerkleKeccak {
99-
function setUp() public {
100-
usingSha = false;
101-
leaves = getLeaves(vm.randomBool() ? 9 : 10);
102-
root = Merkle.merkleizeKeccak(leaves);
103-
proofs = getProofs(leaves);
127+
/// @dev Computes the merkle root using the appropriate hash function
128+
function _genRoot(bytes32[] memory leaves) internal view virtual returns (bytes32) {
129+
function (bytes32[] memory leaves) view returns (bytes32) merkleize = usingSha() ? Merkle.merkleizeSha256 : Merkle.merkleizeKeccak;
130+
if (usingSha()) leaves = _padLeaves(leaves);
131+
return merkleize(leaves);
104132
}
105133

106-
function nextPowerOf2(uint v) internal pure returns (uint) {
134+
/// @dev Rounds up to the next power of 2.
135+
/// https://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2
136+
function _roundUpPow2(uint v) internal pure returns (uint) {
107137
unchecked {
108-
// Round up to the next power of 2 using the method described here:
109-
// https://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2
110-
if (v == 0) return 0;
111138
v -= 1;
112139
v |= v >> 1;
113140
v |= v >> 2;
@@ -121,18 +148,17 @@ contract MerkleKeccakTest is MerkleBaseTest, MerkleKeccak {
121148
}
122149
}
123150

124-
function getProofs(bytes32[] memory leaves) public view virtual override returns (bytes[] memory proofs) {
125-
// Merkle.merkleizeKeccak pads to next power of 2, so we need to match that.
126-
uint numLeaves = nextPowerOf2(leaves.length);
127-
bytes32[] memory paddedLeaves = new bytes32[](numLeaves);
128-
for (uint i = 0; i < leaves.length; ++i) {
129-
// TODO: Point leaves to paddedLeaves using assembly to avoid loop.
130-
paddedLeaves[i] = leaves[i];
131-
}
151+
function usingSha() internal view virtual returns (bool);
152+
}
132153

133-
proofs = new bytes[](leaves.length);
134-
for (uint i = 0; i < leaves.length; ++i) {
135-
proofs[i] = abi.encodePacked(getProof(paddedLeaves, i));
136-
}
154+
contract MerkleKeccakTest is MerkleBaseTest, MerkleKeccak {
155+
function usingSha() internal view virtual override returns (bool) {
156+
return false;
157+
}
158+
}
159+
160+
contract MerkleShaTest is MerkleBaseTest, MerkleSha {
161+
function usingSha() internal view virtual override returns (bool) {
162+
return true;
137163
}
138164
}

‎src/test/utils/Murky.sol‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -179,17 +179,11 @@ contract MerkleSha is MurkyBase {
179179

180180
function hashLeafPairs(bytes32 left, bytes32 right) public view override returns (bytes32 _hash) {
181181
assembly {
182-
switch lt(left, right)
183-
case 0 {
184-
mstore(0x0, right)
185-
mstore(0x20, left)
186-
}
187-
default {
188-
mstore(0x0, left)
189-
mstore(0x20, right)
190-
}
191-
_hash := mload(iszero(staticcall(gas(), SHA256_PRECOMPILE, 0x0, 0x40, 0x0, 0x20)))
192-
if iszero(returndatasize()) { invalid() }
182+
mstore(0x0, left)
183+
mstore(0x20, right)
184+
let success := staticcall(gas(), SHA256_PRECOMPILE, 0x0, 0x40, 0x0, 0x20)
185+
if iszero(success) { revert(0, 0) }
186+
_hash := mload(0x0)
193187
}
194188
}
195189
}

0 commit comments

Comments
 (0)