add withdrawAndCall logic

This commit is contained in:
Drygin 2022-06-22 19:54:24 +03:00
parent 84cf86c4d2
commit 370e3892a6
8 changed files with 310 additions and 4 deletions

View File

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

View File

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

View File

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

View File

@ -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 {

View File

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

15
src/withdrawWorker.js Normal file
View File

@ -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 }

View File

@ -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,

View File

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