Add a MerkleTree builder (#3617)
Co-authored-by: Ernesto García <ernestognw@gmail.com>pull/4949/head
parent
e83142944f
commit
92ff025622
@ -0,0 +1,5 @@ |
||||
--- |
||||
'openzeppelin-solidity': minor |
||||
--- |
||||
|
||||
`Hashes`: A library with commonly used hash functions. |
@ -0,0 +1,5 @@ |
||||
--- |
||||
'openzeppelin-solidity': minor |
||||
--- |
||||
|
||||
`MerkleTree`: A data structure that allows inserting elements into a merkle tree and updating its root hash. |
@ -0,0 +1,43 @@ |
||||
// SPDX-License-Identifier: MIT |
||||
|
||||
pragma solidity ^0.8.0; |
||||
|
||||
import {MerkleTree} from "../utils/structs/MerkleTree.sol"; |
||||
|
||||
contract MerkleTreeMock { |
||||
using MerkleTree for MerkleTree.Bytes32PushTree; |
||||
|
||||
MerkleTree.Bytes32PushTree private _tree; |
||||
|
||||
event LeafInserted(bytes32 leaf, uint256 index, bytes32 root); |
||||
|
||||
function setup(uint8 _depth, bytes32 _zero) public { |
||||
_tree.setup(_depth, _zero); |
||||
} |
||||
|
||||
function push(bytes32 leaf) public { |
||||
(uint256 leafIndex, bytes32 currentRoot) = _tree.push(leaf); |
||||
emit LeafInserted(leaf, leafIndex, currentRoot); |
||||
} |
||||
|
||||
function root() public view returns (bytes32) { |
||||
return _tree.root(); |
||||
} |
||||
|
||||
function depth() public view returns (uint256) { |
||||
return _tree.depth(); |
||||
} |
||||
|
||||
// internal state |
||||
function nextLeafIndex() public view returns (uint256) { |
||||
return _tree._nextLeafIndex; |
||||
} |
||||
|
||||
function sides(uint256 i) public view returns (bytes32) { |
||||
return _tree._sides[i]; |
||||
} |
||||
|
||||
function zeros(uint256 i) public view returns (bytes32) { |
||||
return _tree._zeros[i]; |
||||
} |
||||
} |
@ -0,0 +1,29 @@ |
||||
// SPDX-License-Identifier: MIT |
||||
|
||||
pragma solidity ^0.8.0; |
||||
|
||||
/** |
||||
* @dev Library of standard hash functions. |
||||
*/ |
||||
library Hashes { |
||||
/** |
||||
* @dev Commutative Keccak256 hash of a sorted pair of bytes32. Frequently used when working with merkle proofs. |
||||
* |
||||
* NOTE: Equivalent to the `standardNodeHash` in our https://github.com/OpenZeppelin/merkle-tree[JavaScript library]. |
||||
*/ |
||||
function commutativeKeccak256(bytes32 a, bytes32 b) internal pure returns (bytes32) { |
||||
return a < b ? _efficientKeccak256(a, b) : _efficientKeccak256(b, a); |
||||
} |
||||
|
||||
/** |
||||
* @dev Implementation of keccak256(abi.encode(a, b)) that doesn't allocate or expand memory. |
||||
*/ |
||||
function _efficientKeccak256(bytes32 a, bytes32 b) private pure returns (bytes32 value) { |
||||
/// @solidity memory-safe-assembly |
||||
assembly { |
||||
mstore(0x00, a) |
||||
mstore(0x20, b) |
||||
value := keccak256(0x00, 0x40) |
||||
} |
||||
} |
||||
} |
@ -0,0 +1,154 @@ |
||||
// SPDX-License-Identifier: MIT |
||||
|
||||
pragma solidity ^0.8.0; |
||||
|
||||
import {Hashes} from "../cryptography/Hashes.sol"; |
||||
import {Arrays} from "../Arrays.sol"; |
||||
import {Panic} from "../Panic.sol"; |
||||
|
||||
/** |
||||
* @dev Library for managing https://wikipedia.org/wiki/Merkle_Tree[Merkle Tree] data structures. |
||||
* |
||||
* Each tree is a complete binary tree with the ability to sequentially insert leaves, changing them from a zero to a |
||||
* non-zero value and updating its root. This structure allows inserting commitments (or other entries) that are not |
||||
* stored, but can be proven to be part of the tree at a later time. See {MerkleProof}. |
||||
* |
||||
* A tree is defined by the following parameters: |
||||
* |
||||
* * Depth: The number of levels in the tree, it also defines the maximum number of leaves as 2**depth. |
||||
* * Zero value: The value that represents an empty leaf. Used to avoid regular zero values to be part of the tree. |
||||
* * Hashing function: A cryptographic hash function used to produce internal nodes. |
||||
* |
||||
* _Available since v5.1._ |
||||
*/ |
||||
library MerkleTree { |
||||
/** |
||||
* @dev A complete `bytes32` Merkle tree. |
||||
* |
||||
* The `sides` and `zero` arrays are set to have a length equal to the depth of the tree during setup. |
||||
* |
||||
* The hashing function used during initialization to compute the `zeros` values (value of a node at a given depth |
||||
* for which the subtree is full of zero leaves). This function is kept in the structure for handling insertions. |
||||
* |
||||
* Struct members have an underscore prefix indicating that they are "private" and should not be read or written to |
||||
* directly. Use the functions provided below instead. Modifying the struct manually may violate assumptions and |
||||
* lead to unexpected behavior. |
||||
* |
||||
* NOTE: The `root` is kept up to date after each insertion without keeping track of its history. Consider |
||||
* using a secondary structure to store a list of historical roots (e.g. a mapping, {BitMaps} or {Checkpoints}). |
||||
* |
||||
* WARNING: Updating any of the tree's parameters after the first insertion will result in a corrupted tree. |
||||
*/ |
||||
struct Bytes32PushTree { |
||||
bytes32 _root; |
||||
uint256 _nextLeafIndex; |
||||
bytes32[] _sides; |
||||
bytes32[] _zeros; |
||||
function(bytes32, bytes32) view returns (bytes32) _fnHash; |
||||
} |
||||
|
||||
/** |
||||
* @dev Initialize a {Bytes32PushTree} using {Hashes-commutativeKeccak256} to hash internal nodes. |
||||
* The capacity of the tree (i.e. number of leaves) is set to `2**levels`. |
||||
* |
||||
* Calling this function on MerkleTree that was already setup and used will reset it to a blank state. |
||||
* |
||||
* IMPORTANT: The zero value should be carefully chosen since it will be stored in the tree representing |
||||
* empty leaves. It should be a value that is not expected to be part of the tree. |
||||
*/ |
||||
function setup(Bytes32PushTree storage self, uint8 levels, bytes32 zero) internal { |
||||
return setup(self, levels, zero, Hashes.commutativeKeccak256); |
||||
} |
||||
|
||||
/** |
||||
* @dev Same as {setup}, but allows to specify a custom hashing function. |
||||
* |
||||
* IMPORTANT: Providing a custom hashing function is a security-sensitive operation since it may |
||||
* compromise the soundness of the tree. Consider using functions from {Hashes}. |
||||
*/ |
||||
function setup( |
||||
Bytes32PushTree storage self, |
||||
uint8 levels, |
||||
bytes32 zero, |
||||
function(bytes32, bytes32) view returns (bytes32) fnHash |
||||
) internal { |
||||
// Store depth in the dynamic array |
||||
Arrays.unsafeSetLength(self._sides, levels); |
||||
Arrays.unsafeSetLength(self._zeros, levels); |
||||
|
||||
// Build each root of zero-filled subtrees |
||||
bytes32 currentZero = zero; |
||||
for (uint32 i = 0; i < levels; ++i) { |
||||
Arrays.unsafeAccess(self._zeros, i).value = currentZero; |
||||
currentZero = fnHash(currentZero, currentZero); |
||||
} |
||||
|
||||
// Set the first root |
||||
self._root = currentZero; |
||||
self._nextLeafIndex = 0; |
||||
self._fnHash = fnHash; |
||||
} |
||||
|
||||
/** |
||||
* @dev Insert a new leaf in the tree, and compute the new root. Returns the position of the inserted leaf in the |
||||
* tree, and the resulting root. |
||||
* |
||||
* Hashing the leaf before calling this function is recommended as a protection against |
||||
* second pre-image attacks. |
||||
*/ |
||||
function push(Bytes32PushTree storage self, bytes32 leaf) internal returns (uint256 index, bytes32 newRoot) { |
||||
// Cache read |
||||
uint256 levels = self._zeros.length; |
||||
function(bytes32, bytes32) view returns (bytes32) fnHash = self._fnHash; |
||||
|
||||
// Get leaf index |
||||
uint256 leafIndex = self._nextLeafIndex++; |
||||
|
||||
// Check if tree is full. |
||||
if (leafIndex >= 1 << levels) { |
||||
Panic.panic(Panic.RESOURCE_ERROR); |
||||
} |
||||
|
||||
// Rebuild branch from leaf to root |
||||
uint256 currentIndex = leafIndex; |
||||
bytes32 currentLevelHash = leaf; |
||||
for (uint32 i = 0; i < levels; i++) { |
||||
// Reaching the parent node, is currentLevelHash the left child? |
||||
bool isLeft = currentIndex % 2 == 0; |
||||
|
||||
// If so, next time we will come from the right, so we need to save it |
||||
if (isLeft) { |
||||
Arrays.unsafeAccess(self._sides, i).value = currentLevelHash; |
||||
} |
||||
|
||||
// Compute the current node hash by using the hash function |
||||
// with either the its sibling (side) or the zero value for that level. |
||||
currentLevelHash = fnHash( |
||||
isLeft ? currentLevelHash : Arrays.unsafeAccess(self._sides, i).value, |
||||
isLeft ? Arrays.unsafeAccess(self._zeros, i).value : currentLevelHash |
||||
); |
||||
|
||||
// Update node index |
||||
currentIndex >>= 1; |
||||
} |
||||
|
||||
// Record new root |
||||
self._root = currentLevelHash; |
||||
|
||||
return (leafIndex, currentLevelHash); |
||||
} |
||||
|
||||
/** |
||||
* @dev Tree's current root |
||||
*/ |
||||
function root(Bytes32PushTree storage self) internal view returns (bytes32) { |
||||
return self._root; |
||||
} |
||||
|
||||
/** |
||||
* @dev Tree's depth (set at initialization) |
||||
*/ |
||||
function depth(Bytes32PushTree storage self) internal view returns (uint256) { |
||||
return self._zeros.length; |
||||
} |
||||
} |
@ -0,0 +1,100 @@ |
||||
const { ethers } = require('hardhat'); |
||||
const { expect } = require('chai'); |
||||
const { loadFixture } = require('@nomicfoundation/hardhat-network-helpers'); |
||||
const { PANIC_CODES } = require('@nomicfoundation/hardhat-chai-matchers/panic'); |
||||
const { StandardMerkleTree } = require('@openzeppelin/merkle-tree'); |
||||
|
||||
const { generators } = require('../../helpers/random'); |
||||
|
||||
const makeTree = (leafs = [ethers.ZeroHash]) => |
||||
StandardMerkleTree.of( |
||||
leafs.map(leaf => [leaf]), |
||||
['bytes32'], |
||||
{ sortLeaves: false }, |
||||
); |
||||
|
||||
const hashLeaf = leaf => makeTree().leafHash([leaf]); |
||||
|
||||
const DEPTH = 4n; // 16 slots
|
||||
const ZERO = hashLeaf(ethers.ZeroHash); |
||||
|
||||
async function fixture() { |
||||
const mock = await ethers.deployContract('MerkleTreeMock'); |
||||
await mock.setup(DEPTH, ZERO); |
||||
return { mock }; |
||||
} |
||||
|
||||
describe('MerkleTree', function () { |
||||
beforeEach(async function () { |
||||
Object.assign(this, await loadFixture(fixture)); |
||||
}); |
||||
|
||||
it('sets initial values at setup', async function () { |
||||
const merkleTree = makeTree(Array.from({ length: 2 ** Number(DEPTH) }, () => ethers.ZeroHash)); |
||||
|
||||
expect(await this.mock.root()).to.equal(merkleTree.root); |
||||
expect(await this.mock.depth()).to.equal(DEPTH); |
||||
expect(await this.mock.nextLeafIndex()).to.equal(0n); |
||||
}); |
||||
|
||||
describe('push', function () { |
||||
it('tree is correctly updated', async function () { |
||||
const leafs = Array.from({ length: 2 ** Number(DEPTH) }, () => ethers.ZeroHash); |
||||
|
||||
// for each leaf slot
|
||||
for (const i in leafs) { |
||||
// generate random leaf and hash it
|
||||
const hashedLeaf = hashLeaf((leafs[i] = generators.bytes32())); |
||||
|
||||
// update leaf list and rebuild tree.
|
||||
const tree = makeTree(leafs); |
||||
|
||||
// push value to tree
|
||||
await expect(this.mock.push(hashedLeaf)).to.emit(this.mock, 'LeafInserted').withArgs(hashedLeaf, i, tree.root); |
||||
|
||||
// check tree
|
||||
expect(await this.mock.root()).to.equal(tree.root); |
||||
expect(await this.mock.nextLeafIndex()).to.equal(BigInt(i) + 1n); |
||||
} |
||||
}); |
||||
|
||||
it('revert when tree is full', async function () { |
||||
await Promise.all(Array.from({ length: 2 ** Number(DEPTH) }).map(() => this.mock.push(ethers.ZeroHash))); |
||||
|
||||
await expect(this.mock.push(ethers.ZeroHash)).to.be.revertedWithPanic(PANIC_CODES.TOO_MUCH_MEMORY_ALLOCATED); |
||||
}); |
||||
}); |
||||
|
||||
it('reset', async function () { |
||||
// empty tree
|
||||
const zeroLeafs = Array.from({ length: 2 ** Number(DEPTH) }, () => ethers.ZeroHash); |
||||
const zeroTree = makeTree(zeroLeafs); |
||||
|
||||
// tree with one element
|
||||
const leafs = Array.from({ length: 2 ** Number(DEPTH) }, () => ethers.ZeroHash); |
||||
const hashedLeaf = hashLeaf((leafs[0] = generators.bytes32())); // fill first leaf and hash it
|
||||
const tree = makeTree(leafs); |
||||
|
||||
// root should be that of a zero tree
|
||||
expect(await this.mock.root()).to.equal(zeroTree.root); |
||||
expect(await this.mock.nextLeafIndex()).to.equal(0n); |
||||
|
||||
// push leaf and check root
|
||||
await expect(this.mock.push(hashedLeaf)).to.emit(this.mock, 'LeafInserted').withArgs(hashedLeaf, 0, tree.root); |
||||
|
||||
expect(await this.mock.root()).to.equal(tree.root); |
||||
expect(await this.mock.nextLeafIndex()).to.equal(1n); |
||||
|
||||
// reset tree
|
||||
await this.mock.setup(DEPTH, ZERO); |
||||
|
||||
expect(await this.mock.root()).to.equal(zeroTree.root); |
||||
expect(await this.mock.nextLeafIndex()).to.equal(0n); |
||||
|
||||
// re-push leaf and check root
|
||||
await expect(this.mock.push(hashedLeaf)).to.emit(this.mock, 'LeafInserted').withArgs(hashedLeaf, 0, tree.root); |
||||
|
||||
expect(await this.mock.root()).to.equal(tree.root); |
||||
expect(await this.mock.nextLeafIndex()).to.equal(1n); |
||||
}); |
||||
}); |
Loading…
Reference in new issue