diff --git a/contracts/TornadoPool.sol b/contracts/TornadoPool.sol index fa31873..a36fec3 100644 --- a/contracts/TornadoPool.sol +++ b/contracts/TornadoPool.sol @@ -70,10 +70,14 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver { bytes32 extDataHash; } + struct Account { + address owner; + bytes publicKey; + } + event NewCommitment(bytes32 commitment, uint256 index, bytes encryptedOutput); event NewNullifier(bytes32 nullifier); event PublicKey(address indexed owner, bytes key); - event EncryptedAccount(address indexed owner, bytes account); /** @dev The constructor @@ -200,16 +204,21 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver { } } - function register(bytes memory _publicKey) public { - emit PublicKey(msg.sender, _publicKey); + function register(Account memory _account) public { + require(_account.owner == msg.sender, "only owner can be registered"); + _register(_account); + } + + function _register(Account memory _account) internal { + emit PublicKey(_account.owner, _account.publicKey); } function registerAndTransact( - bytes memory _publicKey, + Account memory _account, Proof memory _proofArgs, ExtData memory _extData ) public { - register(_publicKey); + register(_account); transact(_proofArgs, _extData); } @@ -218,7 +227,7 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver { uint256 _amount, bytes calldata _data ) external override { - (bytes memory _publicKey, Proof memory _args, ExtData memory _extData) = abi.decode(_data, (bytes, Proof, ExtData)); + (Account memory _account, Proof memory _args, ExtData memory _extData) = abi.decode(_data, (Account, Proof, ExtData)); require(_token == token, "provided token is not supported"); require(msg.sender == omniBridge, "only omni bridge"); require(_amount == uint256(_extData.extAmount), "amount from bridge is incorrect"); @@ -226,8 +235,8 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver { totalDeposited += uint256(_extData.extAmount); - if (_publicKey.length != 0) { - register(_publicKey); + if (_account.owner != address(0) && _account.publicKey.length > 0) { + _register(_account); } _transact(_args, _extData); } diff --git a/src/index.js b/src/index.js index 6d1f4cc..be836e1 100644 --- a/src/index.js +++ b/src/index.js @@ -149,13 +149,13 @@ async function transaction({ tornadoPool, ...rest }) { return await receipt.wait() } -async function registerAndTransact({ tornadoPool, poolAddress, ...rest }) { +async function registerAndTransact({ tornadoPool, account, ...rest }) { const { args, extData } = await prepareTransaction({ tornadoPool, ...rest, }) - const receipt = await tornadoPool.registerAndTransact(poolAddress, args, extData, { + const receipt = await tornadoPool.registerAndTransact(account, args, extData, { gasLimit: 2e6, }) await receipt.wait() diff --git a/test/full.test.js b/test/full.test.js index 508aec3..46cb21a 100644 --- a/test/full.test.js +++ b/test/full.test.js @@ -115,7 +115,10 @@ describe('TornadoPool', function () { await registerAndTransact({ tornadoPool, outputs: [aliceDepositUtxo], - poolAddress: aliceDepositUtxo.keypair.address(), + account: { + owner: sender.address, + publicKey: aliceDepositUtxo.keypair.address(), + }, }) const filter = tornadoPool.filters.NewCommitment() @@ -208,7 +211,14 @@ describe('TornadoPool', function () { tornadoPool, outputs: [aliceDepositUtxo], }) - const transactTx = await tornadoPool.populateTransaction.registerAndTransact([], args, extData) + const transactTx = await tornadoPool.populateTransaction.registerAndTransact( + { + owner: '0x0000000000000000000000000000000000000000', + publicKey: [], + }, + args, + extData, + ) const onTokenBridgedData = '0x' + transactTx.data.slice(10) const onTokenBridgedTx = await tornadoPool.populateTransaction.onTokenBridged( token.address,