add withdrawAndCall logic

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));
       }

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)
+    }
+  }
+}

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)
+    }
+  }
+}

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 {

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,
       },
     ],
   )

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 }

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,

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],
   )