|
1 | 1 | // SPDX-License-Identifier: MIT |
2 | 2 | pragma solidity ^0.8.22; |
3 | 3 |
|
4 | | -// Dummy merkle proof generation utilities for testing |
| 4 | +/** |
| 5 | + * @title MerkleTestHelper |
| 6 | + * @dev This contract builds a Merkle tree from a list of addresses, stores the root, |
| 7 | + * and provides a function to retrieve a Merkle proof for a given address. |
| 8 | + * |
| 9 | + * NOTE: Generating Merkle trees on-chain is gas-expensive, so this is typically |
| 10 | + * done only in testing scenarios or for very short lists. |
| 11 | + */ |
5 | 12 | contract MerkleTestHelper { |
6 | | - // This is a placeholder helper. In a real test, you'd generate a real merkle tree offline. |
7 | | - // Here we hardcode a single allowlisted address and its proof. |
8 | | - bytes32[] internal _proof; |
| 13 | + address[] internal _allowedAddrs; |
9 | 14 | bytes32 internal _root; |
10 | | - address internal _allowedAddr; |
11 | 15 |
|
12 | | - constructor(address allowedAddr) { |
13 | | - _allowedAddr = allowedAddr; |
14 | | - // For simplicity, root = keccak256(abi.encodePacked(_allowedAddr)) |
15 | | - // Proof is empty since this is a single-leaf tree. |
16 | | - _root = keccak256(abi.encodePacked(_allowedAddr)); |
| 16 | + /** |
| 17 | + * @dev Constructor that takes in an array of addresses, builds a Merkle tree, and stores the root. |
| 18 | + */ |
| 19 | + constructor(address[] memory allowedAddresses) { |
| 20 | + // Copy addresses to storage |
| 21 | + for (uint256 i = 0; i < allowedAddresses.length; i++) { |
| 22 | + _allowedAddrs.push(allowedAddresses[i]); |
| 23 | + } |
| 24 | + |
| 25 | + // Build leaves from the addresses |
| 26 | + bytes32[] memory leaves = _buildLeaves(_allowedAddrs); |
| 27 | + |
| 28 | + // Compute merkle root |
| 29 | + _root = _computeMerkleRoot(leaves); |
17 | 30 | } |
18 | 31 |
|
| 32 | + /** |
| 33 | + * @notice Returns the Merkle root of the addresses list. |
| 34 | + */ |
19 | 35 | function getRoot() external view returns (bytes32) { |
20 | 36 | return _root; |
21 | 37 | } |
22 | 38 |
|
| 39 | + /** |
| 40 | + * @notice Returns the Merkle proof for a given address. |
| 41 | + * @dev If the address is not found or is not part of the _allowedAddrs array, |
| 42 | + * this will return an empty array. |
| 43 | + */ |
23 | 44 | function getProofFor(address addr) external view returns (bytes32[] memory) { |
24 | | - if (addr == _allowedAddr) { |
25 | | - // Single-leaf tree: no proof necessary except empty array |
| 45 | + // Find the index of the address in our stored list |
| 46 | + (bool isInList, uint256 index) = _findAddressIndex(addr); |
| 47 | + if (!isInList) { |
| 48 | + // Return empty proof if address doesn't exist in the allowed list |
26 | 49 | return new bytes32[](0); |
27 | | - } else { |
28 | | - // No valid proof |
29 | | - bytes32[] memory emptyProof; |
30 | | - return emptyProof; |
31 | 50 | } |
| 51 | + |
| 52 | + // Build leaves in memory |
| 53 | + bytes32[] memory leaves = _buildLeaves(_allowedAddrs); |
| 54 | + |
| 55 | + // Build the proof for the leaf at the found index |
| 56 | + return _buildProof(leaves, index); |
| 57 | + } |
| 58 | + |
| 59 | + /** |
| 60 | + * @dev Creates an array of leaves by double hashing each address: |
| 61 | + * keccak256(bytes.concat(keccak256(abi.encodePacked(address)))) |
| 62 | + */ |
| 63 | + function _buildLeaves(address[] memory addrs) internal pure returns (bytes32[] memory) { |
| 64 | + bytes32[] memory leaves = new bytes32[](addrs.length); |
| 65 | + for (uint256 i = 0; i < addrs.length; i++) { |
| 66 | + leaves[i] = keccak256(bytes.concat(keccak256(abi.encode(addrs[i])))); |
| 67 | + } |
| 68 | + return leaves; |
| 69 | + } |
| 70 | + |
| 71 | + /** |
| 72 | + * @dev Computes the Merkle root from an array of leaves. |
| 73 | + * Pairs each leaf, hashing them together until only one root remains. |
| 74 | + * If there is an odd number of leaves at a given level, the last leaf is "promoted" (copied up). |
| 75 | + */ |
| 76 | + function _computeMerkleRoot(bytes32[] memory leaves) internal pure returns (bytes32) { |
| 77 | + require(leaves.length > 0, "No leaves to build a merkle root"); |
| 78 | + |
| 79 | + uint256 n = leaves.length; |
| 80 | + while (n > 1) { |
| 81 | + for (uint256 i = 0; i < n / 2; i++) { |
| 82 | + // Sort the pair before hashing |
| 83 | + (bytes32 left, bytes32 right) = leaves[2 * i] < leaves[2 * i + 1] |
| 84 | + ? (leaves[2 * i], leaves[2 * i + 1]) |
| 85 | + : (leaves[2 * i + 1], leaves[2 * i]); |
| 86 | + leaves[i] = keccak256(abi.encodePacked(left, right)); |
| 87 | + } |
| 88 | + // If odd, promote last leaf |
| 89 | + if (n % 2 == 1) { |
| 90 | + leaves[n / 2] = leaves[n - 1]; |
| 91 | + n = (n / 2) + 1; |
| 92 | + } else { |
| 93 | + n = n / 2; |
| 94 | + } |
| 95 | + } |
| 96 | + |
| 97 | + // The first element is now the root |
| 98 | + return leaves[0]; |
| 99 | + } |
| 100 | + |
| 101 | + /** |
| 102 | + * @dev Builds a Merkle proof for the leaf at the given index. |
| 103 | + * We recompute the pairing tree on the fly, capturing the "sibling" each time. |
| 104 | + */ |
| 105 | + function _buildProof(bytes32[] memory leaves, uint256 targetIndex) internal pure returns (bytes32[] memory) { |
| 106 | + bytes32[] memory proof = new bytes32[](_proofLength(leaves.length)); |
| 107 | + uint256 proofPos = 0; |
| 108 | + uint256 n = leaves.length; |
| 109 | + uint256 index = targetIndex; |
| 110 | + |
| 111 | + while (n > 1) { |
| 112 | + bool isIndexEven = (index % 2) == 0; |
| 113 | + uint256 pairIndex = isIndexEven ? index + 1 : index - 1; |
| 114 | + |
| 115 | + if (pairIndex < n) { |
| 116 | + // Add the sibling to the proof without sorting |
| 117 | + proof[proofPos] = leaves[pairIndex]; |
| 118 | + proofPos++; |
| 119 | + } |
| 120 | + |
| 121 | + // Move up to the next level |
| 122 | + for (uint256 i = 0; i < n / 2; i++) { |
| 123 | + // Sort pairs when building the next level |
| 124 | + (bytes32 left, bytes32 right) = leaves[2 * i] < leaves[2 * i + 1] |
| 125 | + ? (leaves[2 * i], leaves[2 * i + 1]) |
| 126 | + : (leaves[2 * i + 1], leaves[2 * i]); |
| 127 | + leaves[i] = keccak256(abi.encodePacked(left, right)); |
| 128 | + } |
| 129 | + |
| 130 | + // Handle odd number of leaves |
| 131 | + if (n % 2 == 1) { |
| 132 | + leaves[n / 2] = leaves[n - 1]; |
| 133 | + n = (n / 2) + 1; |
| 134 | + } else { |
| 135 | + n = n / 2; |
| 136 | + } |
| 137 | + |
| 138 | + index = index / 2; |
| 139 | + } |
| 140 | + |
| 141 | + // Trim unused proof elements |
| 142 | + uint256 trimSize = 0; |
| 143 | + for (uint256 i = proof.length; i > 0; i--) { |
| 144 | + if (proof[i - 1] != 0) { |
| 145 | + break; |
| 146 | + } |
| 147 | + trimSize++; |
| 148 | + } |
| 149 | + |
| 150 | + bytes32[] memory trimmedProof = new bytes32[](proof.length - trimSize); |
| 151 | + for (uint256 i = 0; i < trimmedProof.length; i++) { |
| 152 | + trimmedProof[i] = proof[i]; |
| 153 | + } |
| 154 | + |
| 155 | + return trimmedProof; |
32 | 156 | } |
33 | 157 |
|
34 | | - function getAllowedAddress() external view returns (address) { |
35 | | - return _allowedAddr; |
| 158 | + /** |
| 159 | + * @dev Helper to find the index of a given address in the _allowedAddrs array. |
| 160 | + */ |
| 161 | + function _findAddressIndex(address addr) internal view returns (bool, uint256) { |
| 162 | + for (uint256 i = 0; i < _allowedAddrs.length; i++) { |
| 163 | + if (_allowedAddrs[i] == addr) { |
| 164 | + return (true, i); |
| 165 | + } |
| 166 | + } |
| 167 | + return (false, 0); |
| 168 | + } |
| 169 | + |
| 170 | + /** |
| 171 | + * @dev Computes an upper bound for the proof length (worst-case). |
| 172 | + * For n leaves, the maximum proof length is ~log2(n). |
| 173 | + * Here we just do a simple upper bound for clarity. |
| 174 | + */ |
| 175 | + function _proofLength(uint256 n) internal pure returns (uint256) { |
| 176 | + // If n=1, no proof. Otherwise, each tree level can contribute 1 node in the proof path. |
| 177 | + // A simplistic approach: log2(n) <= 256 bits for typical usage, but we do this in-line: |
| 178 | + uint256 count = 0; |
| 179 | + while (n > 1) { |
| 180 | + n = (n + 1) / 2; // integer division round up |
| 181 | + count++; |
| 182 | + } |
| 183 | + return count; |
36 | 184 | } |
37 | 185 | } |
0 commit comments