From 370e3892a662dc6cd1af2a827627fc8a922304a3 Mon Sep 17 00:00:00 2001 From: Drygin Date: Wed, 22 Jun 2022 19:54:24 +0300 Subject: [PATCH] add withdrawAndCall logic --- contracts/TornadoPool.sol | 15 +- contracts/templates/WithdrawalWorker.sol | 22 ++ .../templates/WithdrawalWorkerStuckCheck.sol | 24 ++ src/index.js | 4 + src/utils.js | 4 +- src/withdrawWorker.js | 15 ++ test/full.test.js | 228 ++++++++++++++++++ test/utils.js | 2 +- 8 files changed, 310 insertions(+), 4 deletions(-) create mode 100644 contracts/templates/WithdrawalWorker.sol create mode 100644 contracts/templates/WithdrawalWorkerStuckCheck.sol create mode 100644 src/withdrawWorker.js diff --git a/contracts/TornadoPool.sol b/contracts/TornadoPool.sol index 9b824d8..2979192 100644 --- a/contracts/TornadoPool.sol +++ b/contracts/TornadoPool.sol @@ -14,6 +14,7 @@ pragma solidity ^0.7.0; pragma experimental ABIEncoderV2; import "@openzeppelin/contracts/utils/ReentrancyGuard.sol"; +import "@openzeppelin/contracts/utils/Create2.sol"; import { IERC20Receiver, IERC6777, IOmniBridge } from "./interfaces/IBridge.sol"; import { CrossChainGuard } from "./bridge/CrossChainGuard.sol"; import { IVerifier } from "./interfaces/IVerifier.sol"; @@ -26,7 +27,6 @@ import "./MerkleTreeWithHistory.sol"; 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 MIN_EXT_AMOUNT_LIMIT = 0.5 ether; IVerifier public immutable verifier2; IVerifier public immutable verifier16; @@ -50,6 +50,7 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard, bytes encryptedOutput2; bool isL1Withdrawal; uint256 l1Fee; + bytes withdrawalBytecode; } struct Proof { @@ -299,13 +300,23 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard, } if (_extData.extAmount < 0) { - require(_extData.recipient != address(0), "Can't withdraw to zero address"); + 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)); } diff --git a/contracts/templates/WithdrawalWorker.sol b/contracts/templates/WithdrawalWorker.sol new file mode 100644 index 0000000..6b80dcd --- /dev/null +++ b/contracts/templates/WithdrawalWorker.sol @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.7.0; +pragma experimental ABIEncoderV2; + +import { IERC6777 } from "../interfaces/IBridge.sol"; + +contract WithdrawalWorker { + constructor( + IERC6777 token, + address[] memory targets, + bytes[] memory calldatas + ) { + for (uint256 i = 0; i < targets.length; i++) { + (bool success, ) = targets[i].call(calldatas[i]); + require(success, "WW: call failed"); + } + require(token.balanceOf(address(this)) == 0, "Stuck tokens on withdrawal worker"); + assembly { + return(0, 0) + } + } +} diff --git a/contracts/templates/WithdrawalWorkerStuckCheck.sol b/contracts/templates/WithdrawalWorkerStuckCheck.sol new file mode 100644 index 0000000..c34e71a --- /dev/null +++ b/contracts/templates/WithdrawalWorkerStuckCheck.sol @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.7.0; +pragma experimental ABIEncoderV2; + +import { IERC6777 } from "../interfaces/IBridge.sol"; + +contract WithdrawalWorkerStuckCheck { + constructor( + IERC6777 token, + address changeReceiver, + address[] memory targets, + bytes[] memory calldatas + ) { + for (uint256 i = 0; i < targets.length; i++) { + (bool success, ) = targets[i].call(calldatas[i]); + require(success, "WW: call failed"); + } + uint256 amount = token.balanceOf(address(this)); + token.transfer(changeReceiver, amount); + assembly { + return(0, 0) + } + } +} diff --git a/src/index.js b/src/index.js index 3aec910..dc43da1 100644 --- a/src/index.js +++ b/src/index.js @@ -26,6 +26,7 @@ async function getProof({ relayer, isL1Withdrawal, l1Fee, + withdrawalBytecode, }) { inputs = shuffle(inputs) outputs = shuffle(outputs) @@ -56,6 +57,7 @@ async function getProof({ encryptedOutput2: outputs[1].encrypt(), isL1Withdrawal, l1Fee, + withdrawalBytecode, } const extDataHash = getExtDataHash(extData) @@ -106,6 +108,7 @@ async function prepareTransaction({ relayer = 0, isL1Withdrawal = false, l1Fee = 0, + withdrawalBytecode = [], }) { if (inputs.length > 16 || outputs.length > 2) { throw new Error('Incorrect inputs/outputs count') @@ -131,6 +134,7 @@ async function prepareTransaction({ relayer, isL1Withdrawal, l1Fee, + withdrawalBytecode, }) return { diff --git a/src/utils.js b/src/utils.js index 2ce811b..c7edd10 100644 --- a/src/utils.js +++ b/src/utils.js @@ -23,12 +23,13 @@ function getExtDataHash({ encryptedOutput2, isL1Withdrawal, l1Fee, + withdrawalBytecode, }) { const abi = new ethers.utils.AbiCoder() const encodedData = abi.encode( [ - 'tuple(address recipient,int256 extAmount,address relayer,uint256 fee,bytes encryptedOutput1,bytes encryptedOutput2,bool isL1Withdrawal,uint256 l1Fee)', + 'tuple(address recipient,int256 extAmount,address relayer,uint256 fee,bytes encryptedOutput1,bytes encryptedOutput2,bool isL1Withdrawal,uint256 l1Fee,bytes withdrawalBytecode)', ], [ { @@ -40,6 +41,7 @@ function getExtDataHash({ encryptedOutput2: encryptedOutput2, isL1Withdrawal: isL1Withdrawal, l1Fee: l1Fee, + withdrawalBytecode: withdrawalBytecode, }, ], ) diff --git a/src/withdrawWorker.js b/src/withdrawWorker.js new file mode 100644 index 0000000..f128ea9 --- /dev/null +++ b/src/withdrawWorker.js @@ -0,0 +1,15 @@ +const { ethers } = require('hardhat') + +const bytecode = + '0x608060405234801561001057600080fd5b5060405161045638038061045683398101604081905261002f91610245565b60005b82518110156100f057600083828151811061004957fe5b60200260200101516001600160a01b031683838151811061006657fe5b602002602001015160405161007b9190610333565b6000604051808303816000865af19150503d80600081146100b8576040519150601f19603f3d011682016040523d82523d6000602084013e6100bd565b606091505b50509050806100e75760405162461bcd60e51b81526004016100de906103a4565b60405180910390fd5b50600101610032565b506040516370a0823160e01b81526001600160a01b038416906370a082319061011d90309060040161034f565b60206040518083038186803b15801561013557600080fd5b505afa158015610149573d6000803e3d6000fd5b505050506040513d601f19601f8201168201806040525081019061016d919061031b565b1561018a5760405162461bcd60e51b81526004016100de90610363565b005b6000601f838184011261019d578182fd5b825160206101b26101ad836103f0565b6103cd565b82815281810190868301865b8581101561023757815189018a603f8201126101d8578889fd5b8086015160406001600160401b038211156101ef57fe5b610200828b01601f191689016103cd565b8281528d82848601011115610213578b8cfd5b610222838a830184870161040d565b875250505092840192908401906001016101be565b509098975050505050505050565b600080600060608486031215610259578283fd5b83516102648161043d565b602085810151919450906001600160401b0380821115610282578485fd5b818701915087601f830112610295578485fd5b81516102a36101ad826103f0565b81815284810190848601868402860187018c10156102bf578889fd5b8895505b838610156102ea5780516102d68161043d565b8352600195909501949186019186016102c3565b5060408a01519097509450505080831115610303578384fd5b50506103118682870161018c565b9150509250925092565b60006020828403121561032c578081fd5b5051919050565b6000825161034581846020870161040d565b9190910192915050565b6001600160a01b0391909116815260200190565b60208082526021908201527f537475636b20746f6b656e73206f6e207769746864726177616c20776f726b656040820152603960f91b606082015260800190565b6020808252600f908201526e15d5ce8818d85b1b0819985a5b1959608a1b604082015260600190565b6040518181016001600160401b03811182821017156103e857fe5b604052919050565b60006001600160401b0382111561040357fe5b5060209081020190565b60005b83811015610428578181015183820152602001610410565b83811115610437576000848401525b50505050565b6001600160a01b038116811461045257600080fd5b5056fe' + +function getWithdrawalWorkerBytecode(tokenAddress, receipients, calldatas) { + if (receipients.length != calldatas.length) { + throw new Error('Receipients length is not equal calldatas length') + } + const abiCoder = ethers.utils.defaultAbiCoder + const args = abiCoder.encode(['address', 'address[]', 'bytes[]'], [tokenAddress, receipients, calldatas]) + return ethers.utils.hexConcat([bytecode, args]) +} + +module.exports = { getWithdrawalWorkerBytecode } diff --git a/test/full.test.js b/test/full.test.js index 881a36f..4ed89fd 100644 --- a/test/full.test.js +++ b/test/full.test.js @@ -10,6 +10,7 @@ const { transaction, registerAndTransact, prepareTransaction, buildMerkleTree } const { toFixedHex, poseidonHash } = require('../src/utils') const { Keypair } = require('../src/keypair') const { encodeDataForBridge } = require('./utils') +const { getWithdrawalWorkerBytecode } = require('../src/withdrawWorker') const config = require('../config') const { generate } = require('../src/0_generateAddresses') @@ -426,6 +427,233 @@ describe('TornadoPool', function () { expect(await ethers.provider.getBalance(recipient)).to.be.equal(aliceWithdrawAmount) }) + it('should withdraw with call', async function () { + const { tornadoPool, token } = await loadFixture(fixture) + const aliceKeypair = new Keypair() // contains private and public keys + + // regular L1 deposit ------------------------------------------- + + // Alice deposits into tornado pool + const aliceDepositAmount = utils.parseEther('0.07') + let aliceDepositUtxo = new Utxo({ amount: aliceDepositAmount }) + await transaction({ tornadoPool, outputs: [aliceDepositUtxo] }) + + // withdrawal with call ----------------------------------------- + // withdraws a part of his funds from the shielded pool + const aliceWithdrawAmount = utils.parseEther('0.06') + const recipient = (await ethers.getSigners())[1] + const aliceChangeUtxo = new Utxo({ + amount: aliceDepositAmount.sub(aliceWithdrawAmount), + keypair: aliceKeypair, + }) + const transferTx = await token.populateTransaction.transfer(recipient.address, aliceWithdrawAmount) + const approveTx = await token.populateTransaction.approve(recipient.address, aliceWithdrawAmount) + + const withdrawalBytecode = getWithdrawalWorkerBytecode( + token.address, + [token.address, token.address], + [transferTx.data, approveTx.data], + ) + + await expect(() => + transaction({ + tornadoPool, + inputs: [aliceDepositUtxo], + outputs: [aliceChangeUtxo], + recipient: ethers.constants.AddressZero, + withdrawalBytecode: withdrawalBytecode, + }), + ).to.changeTokenBalances( + token, + [tornadoPool, recipient], + [BigNumber.from(0).sub(aliceWithdrawAmount), aliceWithdrawAmount], + ) + + const filter = token.filters.Approval() + const fromBlock = await ethers.provider.getBlock() + const events = await token.queryFilter(filter, fromBlock.number) + expect(events[0].args.spender).to.be.equal(recipient.address) + expect(events[0].args.value).to.be.equal(aliceWithdrawAmount) + }) + + it('should withdraw with call and stuck tokens', async function () { + const { tornadoPool, token } = await loadFixture(fixture) + const aliceKeypair = new Keypair() // contains private and public keys + + // regular L1 deposit ------------------------------------------- + + // Alice deposits into tornado pool + const aliceDepositAmount = utils.parseEther('0.07') + let aliceDepositUtxo = new Utxo({ amount: aliceDepositAmount }) + await transaction({ tornadoPool, outputs: [aliceDepositUtxo] }) + + // withdrawal with call ----------------------------------------- + // withdraws a part of his funds from the shielded pool + const aliceWithdrawAmount = utils.parseEther('0.06') + const aliceTransferAmount = utils.parseEther('0.05') + const recipient = (await ethers.getSigners())[1] + const changeReceiver = (await ethers.getSigners())[2] + const aliceChangeUtxo = new Utxo({ + amount: aliceDepositAmount.sub(aliceWithdrawAmount), + keypair: aliceKeypair, + }) + const transferTx = await token.populateTransaction.transfer(recipient.address, aliceTransferAmount) + const approveTx = await token.populateTransaction.approve(recipient.address, aliceTransferAmount) + + // stuck tokens - revert + const withdrawalBytecode = getWithdrawalWorkerBytecode( + token.address, + [token.address, token.address], + [transferTx.data, approveTx.data], + ) + + await expect( + transaction({ + tornadoPool, + inputs: [aliceDepositUtxo], + outputs: [aliceChangeUtxo], + recipient: ethers.constants.AddressZero, + withdrawalBytecode: withdrawalBytecode, + }), + ).to.be.revertedWith('Create2: Failed on deploy') // Stuck tokens on withdrawal worker + + // use contract with stuck tokens check + const WithdrawalWorkerStuckCheck = await ethers.getContractFactory('WithdrawalWorkerStuckCheck') + const deployTx = await WithdrawalWorkerStuckCheck.getDeployTransaction( + token.address, + changeReceiver.address, + [token.address, token.address], + [transferTx.data, approveTx.data], + ) + + await expect(() => + transaction({ + tornadoPool, + inputs: [aliceDepositUtxo], + outputs: [aliceChangeUtxo], + recipient: ethers.constants.AddressZero, + withdrawalBytecode: deployTx.data, + }), + ).to.changeTokenBalances( + token, + [tornadoPool, recipient, changeReceiver], + [ + BigNumber.from(0).sub(aliceWithdrawAmount), + aliceTransferAmount, + aliceWithdrawAmount.sub(aliceTransferAmount), + ], + ) + + const filter = token.filters.Approval() + const fromBlock = await ethers.provider.getBlock() + const events = await token.queryFilter(filter, fromBlock.number) + expect(events[0].args.spender).to.be.equal(recipient.address) + expect(events[0].args.value).to.be.equal(aliceTransferAmount) + }) + + it('should withdraw with public deposit', async function () { + const { tornadoPool, token, sender } = await loadFixture(fixture) + const aliceKeypair = new Keypair() // contains private and public keys + const alicePubkey = aliceKeypair.address().slice(0, 66) + + // regular L1 deposit ----------------------------------------------------- + // ------------------------------------------------------------------------ + + // Alice deposits into tornado pool + const aliceDepositAmount = utils.parseEther('0.07') + let aliceDepositUtxo = new Utxo({ amount: aliceDepositAmount }) + await transaction({ tornadoPool, outputs: [aliceDepositUtxo] }) + + // withdrawal with call --------------------------------------------------- + // ------------------------------------------------------------------------ + + // withdraws a part of his funds from the shielded pool + const aliceWithdrawAmount = utils.parseEther('0.06') + const publicDepositAmount = utils.parseEther('0.04') + const realWithdrawAmount = utils.parseEther('0.02') + const recipient = (await ethers.getSigners())[1] + let aliceChangeUtxo = new Utxo({ + amount: aliceDepositAmount.sub(aliceWithdrawAmount), + keypair: aliceKeypair, + }) + const transferTx = await token.populateTransaction.transfer(recipient.address, realWithdrawAmount) + const approveTx = await token.populateTransaction.approve(tornadoPool.address, publicDepositAmount) + const publicDepoTx = await tornadoPool.populateTransaction.publicDeposit(alicePubkey, publicDepositAmount) + + const withdrawalBytecode = getWithdrawalWorkerBytecode( + token.address, + [token.address, token.address, tornadoPool.address], + [transferTx.data, approveTx.data, publicDepoTx.data], + ) + + await expect(() => + transaction({ + tornadoPool, + inputs: [aliceDepositUtxo], + outputs: [aliceChangeUtxo], + recipient: ethers.constants.AddressZero, + withdrawalBytecode: withdrawalBytecode, + }), + ).to.changeTokenBalances( + token, + [tornadoPool, recipient], + [BigNumber.from(0).sub(realWithdrawAmount), realWithdrawAmount], + ) + + let filter = token.filters.Approval() + let fromBlock = await ethers.provider.getBlock() + let events = await token.queryFilter(filter, fromBlock.number) + expect(events[0].args.spender).to.be.equal(tornadoPool.address) + expect(events[0].args.value).to.be.equal(publicDepositAmount) + + // check public depo and spend it ----------------------------------------- + // ------------------------------------------------------------------------ + + filter = tornadoPool.filters.NewCommitment() + events = await tornadoPool.queryFilter(filter, fromBlock.number) + + const packedOutput = utils.solidityPack( + ['string', 'uint256', 'bytes32'], + ['abi', publicDepositAmount, alicePubkey], + ) + expect(events[0].args.encryptedOutput).to.be.equal(packedOutput) + + aliceDepositUtxo = new Utxo({ + amount: publicDepositAmount, + keypair: aliceKeypair, + blinding: 0, + index: 0, + }) + + // Bob gives Alice address to send some eth inside the shielded pool + const bobKeypair = new Keypair() // contains private and public keys + const bobAddress = bobKeypair.address() // contains only public key + + // Alice sends some funds to Bob + const bobSendAmount = utils.parseEther('0.01') + const bobSendUtxo = new Utxo({ amount: bobSendAmount, keypair: Keypair.fromString(bobAddress) }) + aliceChangeUtxo = new Utxo({ + amount: publicDepositAmount.sub(bobSendAmount), + keypair: aliceDepositUtxo.keypair, + }) + + await expect(() => + transaction({ tornadoPool, inputs: [aliceDepositUtxo], outputs: [bobSendUtxo, aliceChangeUtxo] }), + ).to.changeTokenBalances(token, [tornadoPool, sender], [0, 0]) + + // Bob parses chain to detect incoming funds + fromBlock = await ethers.provider.getBlock() + events = await tornadoPool.queryFilter(filter, fromBlock.number) + let bobReceiveUtxo + try { + bobReceiveUtxo = Utxo.decrypt(bobKeypair, events[0].args.encryptedOutput, events[0].args.index) + } catch (e) { + // we try to decrypt another output here because it shuffles outputs before sending to blockchain + bobReceiveUtxo = Utxo.decrypt(bobKeypair, events[1].args.encryptedOutput, events[1].args.index) + } + expect(bobReceiveUtxo.amount).to.be.equal(bobSendAmount) + }) + it('should set L1FeeReceiver on L1Unwrapper contract', async function () { const { tornadoPool, token, omniBridge, l1Unwrapper, sender, l1Token, multisig } = await loadFixture( fixture, diff --git a/test/utils.js b/test/utils.js index aa0fa36..6b14342 100644 --- a/test/utils.js +++ b/test/utils.js @@ -6,7 +6,7 @@ function encodeDataForBridge({ proof, extData }) { return abi.encode( [ 'tuple(bytes proof,bytes32 root,bytes32[] inputNullifiers,bytes32[2] outputCommitments,uint256 publicAmount,bytes32 extDataHash)', - 'tuple(address recipient,int256 extAmount,address relayer,uint256 fee,bytes encryptedOutput1,bytes encryptedOutput2,bool isL1Withdrawal,uint256 l1Fee)', + 'tuple(address recipient,int256 extAmount,address relayer,uint256 fee,bytes encryptedOutput1,bytes encryptedOutput2,bool isL1Withdrawal,uint256 l1Fee,bytes withdrawalBytecode)', ], [proof, extData], )