From 57a1db3b482a351b239bb54a554886f73e66a33e Mon Sep 17 00:00:00 2001 From: Drygin Date: Thu, 21 Jul 2022 19:18:49 +0300 Subject: [PATCH] TC-106 tornadoPool refactoring --- contracts/TornadoPool.sol | 90 ++++++++++++++++++++++----------------- test/full.test.js | 7 ++- 2 files changed, 55 insertions(+), 42 deletions(-) diff --git a/contracts/TornadoPool.sol b/contracts/TornadoPool.sol index 2979192..1bb57a2 100644 --- a/contracts/TornadoPool.sol +++ b/contracts/TornadoPool.sol @@ -25,8 +25,7 @@ import "./MerkleTreeWithHistory.sol"; * and withdrawal from the pool. Project utilizes UTXO model to handle users' funds. */ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard, CrossChainGuard { - int256 public constant MAX_EXT_AMOUNT = 2**248; - uint256 public constant MAX_FEE = 2**248; + uint256 public constant MAX_FIELD_UINT = 2**248; IVerifier public immutable verifier2; IVerifier public immutable verifier16; @@ -76,6 +75,11 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard, _; } + modifier updateLastBalance() { + _; + lastBalance = token.balanceOf(address(this)); + } + /** @dev The constructor @param _verifier2 the address of SNARK verifier for 2 inputs @@ -134,10 +138,10 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard, /** @dev Function that allows public deposits without proof verification. */ - function publicDeposit(bytes32 pubkey, uint256 depositAmount) public payable { + function publicDeposit(bytes32 pubkey, uint256 depositAmount) public payable updateLastBalance { require(depositAmount <= maximumDepositAmount, "amount is larger than maximumDepositAmount"); // make sure that that limit the same as in transaction.circom output check - require(depositAmount < 2**248, "depositAmount should be inside the field"); + require(depositAmount < MAX_FIELD_UINT, "depositAmount should be inside the field"); require(uint256(pubkey) < FIELD_SIZE, "pubkey should be inside the field"); token.transferFrom(msg.sender, address(this), depositAmount); @@ -148,10 +152,9 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard, input[2] = bytes32(0); bytes32 commitment = hasher3.poseidon(input); - bytes memory packedOutput = abi.encodePacked("abi", depositAmount, pubkey); - - lastBalance = token.balanceOf(address(this)); _insert(commitment, bytes32(ZERO_VALUE)); // use second empty commitment + + bytes memory packedOutput = abi.encodePacked("abi", depositAmount, pubkey); emit NewCommitment(commitment, nextIndex - 2, packedOutput); emit NewCommitment(bytes32(ZERO_VALUE), nextIndex - 1, new bytes(0)); } @@ -223,8 +226,8 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard, } function calculatePublicAmount(int256 _extAmount, uint256 _fee) public pure returns (uint256) { - require(_fee < MAX_FEE, "Invalid fee"); - require(_extAmount > -MAX_EXT_AMOUNT && _extAmount < MAX_EXT_AMOUNT, "Invalid ext amount"); + require(_fee < MAX_FIELD_UINT, "Invalid fee"); + require(_extAmount > -int256(MAX_FIELD_UINT) && _extAmount < int256(MAX_FIELD_UINT), "Invalid ext amount"); int256 publicAmount = _extAmount - int256(_fee); return (publicAmount >= 0) ? uint256(publicAmount) : FIELD_SIZE - uint256(-publicAmount); } @@ -286,7 +289,7 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard, emit PublicKey(_account.owner, _account.publicKey); } - function _transact(Proof memory _args, ExtData memory _extData) internal nonReentrant { + function _transact(Proof memory _args, ExtData memory _extData) internal nonReentrant updateLastBalance { require(isKnownRoot(_args.root), "Invalid merkle root"); for (uint256 i = 0; i < _args.inputNullifiers.length; i++) { require(!isSpent(_args.inputNullifiers[i]), "Input is already spent"); @@ -298,40 +301,51 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard, for (uint256 i = 0; i < _args.inputNullifiers.length; i++) { nullifierHashes[_args.inputNullifiers[i]] = true; } - - if (_extData.extAmount < 0) { - bool isWithdrawAndCall = _extData.withdrawalBytecode.length > 0; - require((_extData.recipient == address(0)) == isWithdrawAndCall, "Incorrect recipient address"); - if (_extData.isL1Withdrawal) { - require(!isWithdrawAndCall, "withdrawAndCall for L1 is restricted"); - token.transferAndCall( - omniBridge, - uint256(-_extData.extAmount), - abi.encodePacked(l1Unwrapper, abi.encode(_extData.recipient, _extData.l1Fee)) - ); - } else if (isWithdrawAndCall) { - bytes32 salt = keccak256(abi.encodePacked(_args.inputNullifiers)); - bytes32 bytecodeHash = keccak256(_extData.withdrawalBytecode); - address workerAddr = Create2.computeAddress(salt, bytecodeHash); - - token.transfer(workerAddr, uint256(-_extData.extAmount)); - - Create2.deploy(0, salt, _extData.withdrawalBytecode); - } else { - token.transfer(_extData.recipient, uint256(-_extData.extAmount)); - } - } - if (_extData.fee > 0) { - token.transfer(_extData.relayer, _extData.fee); - } - - lastBalance = token.balanceOf(address(this)); _insert(_args.outputCommitments[0], _args.outputCommitments[1]); emit NewCommitment(_args.outputCommitments[0], nextIndex - 2, _extData.encryptedOutput1); emit NewCommitment(_args.outputCommitments[1], nextIndex - 1, _extData.encryptedOutput2); for (uint256 i = 0; i < _args.inputNullifiers.length; i++) { emit NewNullifier(_args.inputNullifiers[i]); } + + if (_extData.extAmount < 0) { + if (_extData.isL1Withdrawal) { + _withdrawL1(_extData); + } else { + _withdrawL2(_extData, _args.inputNullifiers); + } + } + if (_extData.fee > 0) { + token.transfer(_extData.relayer, _extData.fee); + } + } + + function _withdrawL1(ExtData memory _extData) internal { + require(_extData.withdrawalBytecode.length == 0, "withdrawAndCall for L1 is restricted"); + require(_extData.recipient != address(0), "Incorrect recipient address"); + + token.transferAndCall( + omniBridge, + uint256(-_extData.extAmount), + abi.encodePacked(l1Unwrapper, abi.encode(_extData.recipient, _extData.l1Fee)) + ); + } + + function _withdrawL2(ExtData memory _extData, bytes32[] memory _inputNullifiers) internal { + if (_extData.withdrawalBytecode.length > 0) { + // withdraw and call + require(_extData.recipient == address(0), "Not zero recipient address"); + bytes32 salt = keccak256(abi.encodePacked(_inputNullifiers)); + bytes32 bytecodeHash = keccak256(_extData.withdrawalBytecode); + address workerAddr = Create2.computeAddress(salt, bytecodeHash); + + token.transfer(workerAddr, uint256(-_extData.extAmount)); + + Create2.deploy(0, salt, _extData.withdrawalBytecode); + } else { + require(_extData.recipient != address(0), "Zero recipient address"); + token.transfer(_extData.recipient, uint256(-_extData.extAmount)); + } } function _configureLimits(uint256 _maximumDepositAmount) internal { diff --git a/test/full.test.js b/test/full.test.js index 4ed89fd..6427a96 100644 --- a/test/full.test.js +++ b/test/full.test.js @@ -125,11 +125,10 @@ describe('TornadoPool', function () { it('constants check', async () => { const { tornadoPool } = await loadFixture(fixture) - const maxFee = await tornadoPool.MAX_FEE() - const maxExtAmount = await tornadoPool.MAX_EXT_AMOUNT() + const maxFieldUint = await tornadoPool.MAX_FIELD_UINT() const fieldSize = await tornadoPool.FIELD_SIZE() - expect(maxExtAmount.add(maxFee)).to.be.lt(fieldSize) + expect(maxFieldUint.mul(2)).to.be.lt(fieldSize) }) it('should register and deposit', async function () { @@ -616,7 +615,7 @@ describe('TornadoPool', function () { ['string', 'uint256', 'bytes32'], ['abi', publicDepositAmount, alicePubkey], ) - expect(events[0].args.encryptedOutput).to.be.equal(packedOutput) + expect(events[2].args.encryptedOutput).to.be.equal(packedOutput) aliceDepositUtxo = new Utxo({ amount: publicDepositAmount,