diff --git a/contracts/MerkleTreeWithHistory.sol b/contracts/MerkleTreeWithHistory.sol index 28aedbe..5c055c3 100644 --- a/contracts/MerkleTreeWithHistory.sol +++ b/contracts/MerkleTreeWithHistory.sol @@ -5,7 +5,7 @@ library MiMC { } contract MerkleTreeWithHistory { - uint8 levels; + uint256 public levels; uint8 constant ROOT_HISTORY_SIZE = 100; uint256[] private _roots; @@ -18,7 +18,7 @@ contract MerkleTreeWithHistory { event LeafAdded(uint256 leaf, uint32 leaf_index); - constructor(uint8 tree_levels, uint256 zero_value) public { + constructor(uint256 tree_levels, uint256 zero_value) public { levels = tree_levels; _zeros.push(zero_value); @@ -50,13 +50,13 @@ contract MerkleTreeWithHistory { function _insert(uint256 leaf) internal { uint32 leaf_index = next_index; uint32 current_index = next_index; + require(current_index != 2**(levels - 1), "Merkle tree is full"); next_index += 1; - uint256 current_level_hash = leaf; uint256 left; uint256 right; - for (uint8 i = 0; i < levels; i++) { + for (uint256 i = 0; i < levels; i++) { if (current_index % 2 == 0) { left = current_level_hash; right = _zeros[i]; @@ -78,7 +78,7 @@ contract MerkleTreeWithHistory { emit LeafAdded(leaf, leaf_index); } - function isKnownRoot(uint root) public view returns(bool) { + function isKnownRoot(uint256 root) public view returns(bool) { if (root == 0) { return false; } diff --git a/test/MerkleTreeWithHistory.test.js b/test/MerkleTreeWithHistory.test.js index 5f2f633..452a38a 100644 --- a/test/MerkleTreeWithHistory.test.js +++ b/test/MerkleTreeWithHistory.test.js @@ -12,6 +12,8 @@ const MiMC = artifacts.require('./MiMC.sol') const MerkleTree = require('../lib/MerkleTree') const MimcHasher = require('../lib/MiMC') +const { AMOUNT, MERKLE_TREE_HEIGHT, EMPTY_ELEMENT } = process.env + function BNArrayToStringArray(array) { const arrayToPrint = [] array.forEach(item => { @@ -25,8 +27,9 @@ contract('MerkleTreeWithHistory', async accounts => { let miMC const sender = accounts[0] const emptyAddress = '0x0000000000000000000000000000000000000000' - const levels = 16 - const zeroValue = 1337 + const levels = MERKLE_TREE_HEIGHT || 16 + const zeroValue = EMPTY_ELEMENT || 1337 + const value = AMOUNT || '1000000000000000000' let snapshotId let prefix = 'test' let tree @@ -145,7 +148,6 @@ contract('MerkleTreeWithHistory', async accounts => { describe('#insert', async () => { it('should insert', async () => { - let filled_subtrees let rootFromContract for (i = 1; i < 11; i++) { @@ -161,6 +163,21 @@ contract('MerkleTreeWithHistory', async accounts => { // console.log('rootFromCon', root.toString()) } }) + + it.skip('should reject if tree is full', async () => { + // consider change the merkle tree height to run the test + // it should be changed in .env and ./circuits/withdraw.circom files (default is 16) + + for (i = 0; i < 2**(MERKLE_TREE_HEIGHT - 1); i++) { + await merkleTreeWithHistory.insert(i+42).should.be.fulfilled + } + + let error = await merkleTreeWithHistory.insert(1337).should.be.rejected + error.reason.should.be.equal('Merkle tree is full') + + error = await merkleTreeWithHistory.insert(1).should.be.rejected + error.reason.should.be.equal('Merkle tree is full') + }) }) describe('#MIMC', async () => {