more tests

contracts/Mocks/MockAMB.sol

@@ -1,5 +1,6 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.7.0;
+pragma abicoder v2;
 
 import { IAMB } from "../interfaces/IBridge.sol";
 
@@ -7,6 +8,11 @@ contract MockAMB is IAMB {
   address public xDomainMessageSender;
   bytes32 public xDomainMessageChainId;
 
+  struct Call {
+    address who;
+    bytes callData;
+  }
+
   constructor(address _xDomainMessageSender, uint256 _xDomainMessageChainId) {
     xDomainMessageSender = _xDomainMessageSender;
     xDomainMessageChainId = bytes32(uint256(_xDomainMessageChainId));
@@ -24,8 +30,10 @@ contract MockAMB is IAMB {
     return xDomainMessageChainId;
   }
 
-  function execute(address _who, bytes calldata _calldata) external returns (bool success, bytes memory result) {
-    (success, result) = _who.call(_calldata);
+  function execute(Call[] calldata _calls) external returns (bool success, bytes memory result) {
+    for (uint256 i = 0; i < _calls.length; i++) {
+      (success, result) = _calls[i].who.call(_calls[i].callData);
       require(success, string(result));
     }
+  }
 }

contracts/Mocks/MockOmniBridge.sol

@@ -1,11 +1,17 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.7.0;
+pragma abicoder v2;
 
 import { IAMB, IOmniBridge } from "../interfaces/IBridge.sol";
 
 contract MockOmniBridge is IOmniBridge {
   IAMB public AMB;
 
+  struct Call {
+    address who;
+    bytes callData;
+  }
+
   constructor(IAMB _AMB) {
     AMB = _AMB;
   }
@@ -14,9 +20,10 @@ contract MockOmniBridge is IOmniBridge {
     return AMB;
   }
 
-  function execute(address _who, bytes calldata _calldata) external returns (bool success, bytes memory result) {
-    (success, result) = _who.call(_calldata);
-    require(success, string(result));
+  function execute(Call[] calldata _calls) external returns (bool success, bytes memory result) {
+    for (uint256 i = 0; i < _calls.length; i++) {
+      (success, result) = _calls[i].who.call(_calls[i].callData);
+    }
   }
 
   event OnTokenTransfer(address contr, address from, address receiver, uint256 value, bytes data);

contracts/TornadoPool.sol

@@ -155,12 +155,15 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard,
     require(token.balanceOf(address(this)) >= uint256(_extData.extAmount) + lastBalance, "bridge did not send enough tokens");
     require(uint256(_extData.extAmount) <= maximumDepositAmount, "amount is larger than maximumDepositAmount");
     uint256 sentAmount = token.balanceOf(address(this)) - lastBalance;
-    try TornadoPool(address(this)).bridgeTransact(_args, _extData) {} catch (bytes memory) {
+    try TornadoPool(address(this)).onTransact(_args, _extData) {} catch (bytes memory) {
       token.transfer(multisig, sentAmount);
     }
   }
 
-  function bridgeTransact(Proof memory _args, ExtData memory _extData) external {
+  /**
+   * @dev Wrapper for the internal func _transact to call it using try-catch from onTokenBridged
+   */
+  function onTransact(Proof memory _args, ExtData memory _extData) external {
     require(msg.sender == address(this), "can be called only from onTokenBridged");
     _transact(_args, _extData);
   }

test/full.test.js

@@ -68,14 +68,14 @@ describe('TornadoPool', function () {
 
     await token.approve(tornadoPool.address, utils.parseEther('10000'))
 
-    return { tornadoPool, token, proxy, omniBridge, amb, gov }
+    return { tornadoPool, token, proxy, omniBridge, amb, gov, multisig }
   }
 
   describe('Upgradeability tests', () => {
     it('admin should be gov', async () => {
       const { proxy, amb, gov } = await loadFixture(fixture)
       const { data } = await proxy.populateTransaction.admin()
-      const { result } = await amb.callStatic.execute(proxy.address, data)
+      const { result } = await amb.callStatic.execute([{ who: proxy.address, callData: data }])
       expect('0x' + result.slice(26)).to.be.equal(gov.address.toLowerCase())
     })
 
@@ -96,7 +96,7 @@ describe('TornadoPool', function () {
         newDepositLimit,
       )
 
-      await amb.execute(tornadoPool.address, data)
+      await amb.execute([{ who: tornadoPool.address, callData: data }])
 
       expect(await tornadoPool.maximumDepositAmount()).to.be.equal(newDepositLimit)
       expect(await tornadoPool.minimalWithdrawalAmount()).to.be.equal(newWithdrawalLimit)
@@ -240,9 +240,14 @@ describe('TornadoPool', function () {
       aliceDepositUtxo.amount,
       onTokenBridgedData,
     )
-    // emulating bridge. first it sends tokens and then calls onTokenBridged method
-    await token.transfer(tornadoPool.address, aliceDepositAmount)
-    await omniBridge.execute(tornadoPool.address, onTokenBridgedTx.data)
+    // emulating bridge. first it sends tokens to omnibridge mock then it sends to the pool
+    await token.transfer(omniBridge.address, aliceDepositAmount)
+    const transferTx = await token.populateTransaction.transfer(tornadoPool.address, aliceDepositAmount)
+
+    await omniBridge.execute([
+      { who: token.address, callData: transferTx.data }, // send tokens to pool
+      { who: tornadoPool.address, callData: onTokenBridgedTx.data }, // call onTokenBridgedTx
+    ])
 
     // withdraws a part of his funds from the shielded pool
     const aliceWithdrawAmount = utils.parseEther('0.06')
@@ -265,6 +270,62 @@ describe('TornadoPool', function () {
     expect(omniBridgeBalance).to.be.equal(aliceWithdrawAmount)
   })
 
+  it('should transfer funds to multisig in case of L1 deposit fail', async function () {
+    const { tornadoPool, token, omniBridge, multisig } = await loadFixture(fixture)
+    const aliceKeypair = new Keypair() // contains private and public keys
+
+    // Alice deposits into tornado pool
+    const aliceDepositAmount = utils.parseEther('0.07')
+    const aliceDepositUtxo = new Utxo({ amount: aliceDepositAmount, keypair: aliceKeypair })
+    const { args, extData } = await prepareTransaction({
+      tornadoPool,
+      outputs: [aliceDepositUtxo],
+    })
+
+    args.proof = args.proof.slice(0, -2)
+
+    const onTokenBridgedData = encodeDataForBridge({
+      proof: args,
+      extData,
+    })
+
+    const onTokenBridgedTx = await tornadoPool.populateTransaction.onTokenBridged(
+      token.address,
+      aliceDepositUtxo.amount,
+      onTokenBridgedData,
+    )
+    // emulating bridge. first it sends tokens to omnibridge mock then it sends to the pool
+    await token.transfer(omniBridge.address, aliceDepositAmount)
+    const transferTx = await token.populateTransaction.transfer(tornadoPool.address, aliceDepositAmount)
+
+    const lastRoot = await tornadoPool.getLastRoot()
+    await omniBridge.execute([
+      { who: token.address, callData: transferTx.data }, // send tokens to pool
+      { who: tornadoPool.address, callData: onTokenBridgedTx.data }, // call onTokenBridgedTx
+    ])
+
+    const multisigBalance = await token.balanceOf(multisig.address)
+    expect(multisigBalance).to.be.equal(aliceDepositAmount)
+    expect(await tornadoPool.getLastRoot()).to.be.equal(lastRoot)
+  })
+
+  it('should revert if onTransact called directly', async () => {
+    const { tornadoPool } = await loadFixture(fixture)
+    const aliceKeypair = new Keypair() // contains private and public keys
+
+    // Alice deposits into tornado pool
+    const aliceDepositAmount = utils.parseEther('0.07')
+    const aliceDepositUtxo = new Utxo({ amount: aliceDepositAmount, keypair: aliceKeypair })
+    const { args, extData } = await prepareTransaction({
+      tornadoPool,
+      outputs: [aliceDepositUtxo],
+    })
+
+    await expect(tornadoPool.onTransact(args, extData)).to.be.revertedWith(
+      'can be called only from onTokenBridged',
+    )
+  })
+
   it('should work with 16 inputs', async function () {
     const { tornadoPool } = await loadFixture(fixture)
     const aliceDepositAmount = utils.parseEther('0.07')