Cross Chain Guard test

contracts/CrossChainUpgradeableProxy.sol

@@ -1,16 +1,13 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.7.0;
 
-import { IAMB } from "./interfaces/IBridge.sol";
+import { CrossChainGuard } from "./bridge/CrossChainGuard.sol";
 import "@openzeppelin/contracts/proxy/TransparentUpgradeableProxy.sol";
 
 /**
  * @dev TransparentUpgradeableProxy where admin acts from a different chain.
  */
-contract CrossChainUpgradeableProxy is TransparentUpgradeableProxy {
-  IAMB public immutable ambBridge;
-  bytes32 public immutable adminChainId;
-
+contract CrossChainUpgradeableProxy is TransparentUpgradeableProxy, CrossChainGuard {
   /**
    * @dev Initializes an upgradeable proxy backed by the implementation at `_logic`.
    */
@@ -18,25 +15,23 @@ contract CrossChainUpgradeableProxy is TransparentUpgradeableProxy {
     address _logic,
     address _admin,
     bytes memory _data,
-    IAMB _ambBridge,
+    address _ambBridge,
     uint256 _adminChainId
-  ) TransparentUpgradeableProxy(_logic, _admin, _data) {
-    ambBridge = _ambBridge;
-    adminChainId = bytes32(uint256(_adminChainId));
-  }
+  ) TransparentUpgradeableProxy(_logic, _admin, _data) CrossChainGuard(_ambBridge, _adminChainId, _admin) {}
 
   /**
    * @dev Modifier used internally that will delegate the call to the implementation unless the sender is the cross chain admin.
    */
   modifier ifAdmin() override {
-    if (
-      msg.sender == address(ambBridge) &&
-      ambBridge.messageSourceChainId() == adminChainId &&
-      ambBridge.messageSender() == _admin()
-    ) {
+    if (isCalledByOwner()) {
       _;
     } else {
       _fallback();
     }
   }
+
+  /**
+   * @dev Override to allow admin access the fallback function.
+   */
+  function _beforeFallback() internal override {}
 }

contracts/Mocks/MockAMB.sol

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.7.0;
 
-import { IAMB } from "../CrossChainUpgradeableProxy.sol";
+import { IAMB } from "../interfaces/IBridge.sol";
 
 contract MockAMB is IAMB {
   address public xDomainMessageSender;

contracts/TornadoPool.sol

@@ -14,14 +14,14 @@ pragma solidity ^0.7.0;
 pragma experimental ABIEncoderV2;
 
 import "@openzeppelin/contracts/utils/ReentrancyGuard.sol";
-import { IERC20Receiver, IERC6777 } from "./interfaces/IBridge.sol";
+import { IERC20Receiver, IERC6777, IOmniBridge } from "./interfaces/IBridge.sol";
+import { CrossChainGuard } from "./bridge/CrossChainGuard.sol";
 import { IVerifier } from "./interfaces/IVerifier.sol";
 import "./MerkleTreeWithHistory.sol";
 
-contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard {
+contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard, CrossChainGuard {
   int256 public constant MAX_EXT_AMOUNT = 2**248;
   uint256 public constant MAX_FEE = 2**248;
-  address public immutable governance;
 
   IVerifier public immutable verifier2;
   IVerifier public immutable verifier16;
@@ -63,7 +63,7 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard {
   event PublicKey(address indexed owner, bytes key);
 
   modifier onlyGovernance() {
-    require(msg.sender == governance, "only governance");
+    require(isCalledByOwner(), "only governance");
     _;
   }
 
@@ -80,14 +80,17 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard {
     IERC6777 _token,
     address _omniBridge,
     address _l1Unwrapper,
-    address _governance
-  ) MerkleTreeWithHistory(_levels, _hasher) {
+    address _governance,
+    uint256 _l1ChainId
+  )
+    MerkleTreeWithHistory(_levels, _hasher)
+    CrossChainGuard(address(IOmniBridge(_omniBridge).bridgeContract()), _l1ChainId, _governance)
+  {
     verifier2 = _verifier2;
     verifier16 = _verifier16;
     token = _token;
     omniBridge = _omniBridge;
     l1Unwrapper = _l1Unwrapper;
-    governance = _governance;
   }
 
   function initialize(uint256 _minimalWithdrawalAmount, uint256 _maximumDepositAmount) external initializer {

contracts/bridge/CrossChainGuard.sol

@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.7.0;
+
+import { IAMB } from "../interfaces/IBridge.sol";
+
+contract CrossChainGuard {
+  IAMB public immutable ambBridge;
+  bytes32 public immutable ownerChainId;
+  address public immutable owner;
+
+  constructor(
+    address _ambBridge,
+    uint256 _ownerChainId,
+    address _owner
+  ) {
+    ambBridge = IAMB(_ambBridge);
+    owner = _owner;
+    ownerChainId = bytes32(uint256(_ownerChainId));
+  }
+
+  function isCalledByOwner() public virtual returns (bool) {
+    return
+      msg.sender == address(ambBridge) && ambBridge.messageSourceChainId() == ownerChainId && ambBridge.messageSender() == owner;
+  }
+}

scripts/deployTornado.js

@@ -11,7 +11,7 @@ async function main() {
   const omniBridge = '0x59447362798334d3485c64D1e4870Fde2DDC0d75'
   const amb = '0x162e898bd0aacb578c8d5f8d6ca588c13d2a383f'
   const token = '0xCa8d20f3e0144a72C6B5d576e9Bd3Fd8557E2B04' // WBNB
-  const l1Unwrapper = '0xcf35E84bbA3506BB97cf6fAEFe6cc1A9bd843Fc2' // WBNB -> BNB
+  const l1Unwrapper = '0x2353Dcda746fa1AAD17C5650Ddf2A20112862197' // WBNB -> BNB
   const l1ChainId = 56
 
   const Verifier2 = await ethers.getContractFactory('Verifier2')
@@ -40,6 +40,7 @@ async function main() {
       omniBridge,
       l1Unwrapper,
       govAddress,
+      l1ChainId,
     ]).slice(1, -1)}\n`,
   )
   const tornadoImpl = prompt('Deploy tornado pool implementation and provide address here:\n')

test/full.test.js

@@ -47,6 +47,7 @@ describe('TornadoPool', function () {
       omniBridge.address,
       l1Unwrapper.address,
       gov.address,
+      l1ChainId,
     )
 
     const proxy = await deploy(
@@ -81,6 +82,22 @@ describe('TornadoPool', function () {
         "Transaction reverted: function selector was not recognized and there's no fallback function",
       )
     })
+
+    it('should configure', async () => {
+      const { tornadoPool, amb } = await loadFixture(fixture)
+      const newWithdrawalLimit = utils.parseEther('0.01337')
+      const newDepositLimit = utils.parseEther('1337')
+
+      const { data } = await tornadoPool.populateTransaction.configureLimits(
+        newWithdrawalLimit,
+        newDepositLimit,
+      )
+
+      await amb.execute(tornadoPool.address, data)
+
+      expect(await tornadoPool.maximumDepositAmount()).to.be.equal(newDepositLimit)
+      expect(await tornadoPool.minimalWithdrawalAmount()).to.be.equal(newWithdrawalLimit)
+    })
   })
 
   it('encrypt -> decrypt should work', () => {