From 9cf7dc384d1970d34f6fe63e7e8ad189710ec3c4 Mon Sep 17 00:00:00 2001
From: Alexey <peretsmobil@gmail.com>
Date: Fri, 29 Oct 2021 13:20:25 +0300
Subject: [PATCH] try catch for bridge call

---
 contracts/TornadoPool.sol | 42 ++++++++++++++++++++++++++++++++++++++-
 test/full.test.js         |  3 ++-
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/contracts/TornadoPool.sol b/contracts/TornadoPool.sol
index 51cd01b..c0b2409 100644
--- a/contracts/TornadoPool.sol
+++ b/contracts/TornadoPool.sol
@@ -31,6 +31,7 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard,
   IERC6777 public immutable token;
   address public immutable omniBridge;
   address public immutable l1Unwrapper;
+  address public immutable multisig;
 
   uint256 public lastBalance;
   uint256 public minimalWithdrawalAmount;
@@ -70,6 +71,11 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard,
     _;
   }
 
+  modifier onlyMultisig() {
+    require(msg.sender == multisig, "only governance");
+    _;
+  }
+
   /**
     @dev The constructor
     @param _verifier2 the address of SNARK verifier for 2 inputs
@@ -81,6 +87,7 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard,
     @param _l1Unwrapper address of the L1Helper
     @param _governance owner address
     @param _l1ChainId chain id of L1
+    @param _multisig multisig on L2
   */
   constructor(
     IVerifier _verifier2,
@@ -91,7 +98,8 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard,
     address _omniBridge,
     address _l1Unwrapper,
     address _governance,
-    uint256 _l1ChainId
+    uint256 _l1ChainId,
+    address _multisig
   )
     MerkleTreeWithHistory(_levels, _hasher)
     CrossChainGuard(address(IOmniBridge(_omniBridge).bridgeContract()), _l1ChainId, _governance)
@@ -101,6 +109,7 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard,
     token = _token;
     omniBridge = _omniBridge;
     l1Unwrapper = _l1Unwrapper;
+    multisig = _multisig;
   }
 
   function initialize(uint256 _minimalWithdrawalAmount, uint256 _maximumDepositAmount) external initializer {
@@ -145,9 +154,40 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard,
     require(_amount >= uint256(_extData.extAmount), "amount from bridge is incorrect");
     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) {
+      token.transfer(multisig, sentAmount);
+    }
+  }
+
+  function bridgeTransact(Proof memory _args, ExtData memory _extData) external {
+    require(msg.sender == address(this), "can be called only from onTokenBridged");
     _transact(_args, _extData);
   }
 
+  /// @dev Method to claim junk and accidentally sent tokens
+  function rescueTokens(
+    IERC6777 _token,
+    address payable _to,
+    uint256 _balance
+  ) external onlyMultisig {
+    require(_to != address(0), "TORN: can not send to zero address");
+    require(_token != token, "can not rescue pool asset");
+
+    if (_token == IERC6777(0)) {
+      // for Ether
+      uint256 totalBalance = address(this).balance;
+      uint256 balance = _balance == 0 ? totalBalance : _balance;
+      _to.transfer(balance);
+    } else {
+      // any other erc20
+      uint256 totalBalance = _token.balanceOf(address(this));
+      uint256 balance = _balance == 0 ? totalBalance : _balance;
+      require(balance > 0, "TORN: trying to send 0 balance");
+      _token.transfer(_to, balance);
+    }
+  }
+
   function configureLimits(uint256 _minimalWithdrawalAmount, uint256 _maximumDepositAmount) public onlyGovernance {
     _configureLimits(_minimalWithdrawalAmount, _maximumDepositAmount);
   }
diff --git a/test/full.test.js b/test/full.test.js
index c157145..94b62ec 100644
--- a/test/full.test.js
+++ b/test/full.test.js
@@ -25,7 +25,7 @@ describe('TornadoPool', function () {
 
   async function fixture() {
     require('../scripts/compileHasher')
-    const [sender, gov, l1Unwrapper] = await ethers.getSigners()
+    const [sender, gov, l1Unwrapper, multisig] = await ethers.getSigners()
     const verifier2 = await deploy('Verifier2')
     const verifier16 = await deploy('Verifier16')
     const hasher = await deploy('Hasher')
@@ -48,6 +48,7 @@ describe('TornadoPool', function () {
       l1Unwrapper.address,
       gov.address,
       l1ChainId,
+      multisig.address,
     )
 
     const { data } = await tornadoPoolImpl.populateTransaction.initialize(