From a93caef5fadbc76af370c9d11690a0547808ef8a Mon Sep 17 00:00:00 2001
From: Alexey <peretsmobil@gmail.com>
Date: Tue, 5 Oct 2021 12:08:04 +0300
Subject: [PATCH] fix crosschainproxy

---
 contracts/CrossChainUpgradeableProxy.sol | 25 +++++++-----------------
 contracts/Mocks/MockAMB.sol              |  4 ++++
 contracts/Mocks/MockOmniBridge.sol       |  2 +-
 contracts/interfaces/Bridge.sol          | 13 ++++++++++++
 test/full.test.js                        | 10 +++++-----
 5 files changed, 30 insertions(+), 24 deletions(-)
 create mode 100644 contracts/interfaces/Bridge.sol

diff --git a/contracts/CrossChainUpgradeableProxy.sol b/contracts/CrossChainUpgradeableProxy.sol
index 167207a..8399e88 100644
--- a/contracts/CrossChainUpgradeableProxy.sol
+++ b/contracts/CrossChainUpgradeableProxy.sol
@@ -1,25 +1,14 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.7.0;
 
+import { IAMB } from "./interfaces/Bridge.sol";
 import "@openzeppelin/contracts/contracts/proxy/TransparentUpgradeableProxy.sol";
 
-// https://docs.tokenbridge.net/amb-bridge/development-of-a-cross-chain-application/how-to-develop-xchain-apps-by-amb#call-a-method-in-another-chain-using-the-amb-bridge
-
-interface IAMB {
-  function messageSender() external view returns (address);
-
-  function messageSourceChainId() external view returns (bytes32);
-}
-
-interface IOmniBridge {
-  function bridgeContract() external view returns (IAMB);
-}
-
 /**
  * @dev TransparentUpgradeableProxy where admin acts from a different chain.
  */
 contract CrossChainUpgradeableProxy is TransparentUpgradeableProxy {
-  IOmniBridge public immutable omniBridge;
+  IAMB public immutable ambBridge;
   bytes32 public immutable adminChainId;
 
   /**
@@ -29,10 +18,10 @@ contract CrossChainUpgradeableProxy is TransparentUpgradeableProxy {
     address _logic,
     address _admin,
     bytes memory _data,
-    IOmniBridge _omniBridge,
+    IAMB _ambBridge,
     uint256 _adminChainId
   ) TransparentUpgradeableProxy(_logic, _admin, _data) {
-    omniBridge = _omniBridge;
+    ambBridge = _ambBridge;
     adminChainId = bytes32(uint256(_adminChainId));
   }
 
@@ -41,9 +30,9 @@ contract CrossChainUpgradeableProxy is TransparentUpgradeableProxy {
    */
   modifier ifAdmin() override {
     if (
-      msg.sender == address(omniBridge) &&
-      omniBridge.bridgeContract().messageSourceChainId() == adminChainId &&
-      omniBridge.bridgeContract().messageSender() == _admin()
+      msg.sender == address(ambBridge) &&
+      ambBridge.messageSourceChainId() == adminChainId &&
+      ambBridge.messageSender() == _admin()
     ) {
       _;
     } else {
diff --git a/contracts/Mocks/MockAMB.sol b/contracts/Mocks/MockAMB.sol
index 65d32a7..d6b430e 100644
--- a/contracts/Mocks/MockAMB.sol
+++ b/contracts/Mocks/MockAMB.sol
@@ -23,4 +23,8 @@ contract MockAMB is IAMB {
   function messageSourceChainId() external view override returns (bytes32) {
     return xDomainMessageChainId;
   }
+
+  function execute(address _who, bytes calldata _calldata) external returns (bool success, bytes memory result) {
+    (success, result) = _who.call(_calldata);
+  }
 }
diff --git a/contracts/Mocks/MockOmniBridge.sol b/contracts/Mocks/MockOmniBridge.sol
index bb2ef3a..d90817e 100644
--- a/contracts/Mocks/MockOmniBridge.sol
+++ b/contracts/Mocks/MockOmniBridge.sol
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.7.0;
 
-import { IAMB, IOmniBridge } from "../CrossChainUpgradeableProxy.sol";
+import { IAMB, IOmniBridge } from "../interfaces/Bridge.sol";
 
 contract MockOmniBridge is IOmniBridge {
   IAMB public AMB;
diff --git a/contracts/interfaces/Bridge.sol b/contracts/interfaces/Bridge.sol
new file mode 100644
index 0000000..0ed88d2
--- /dev/null
+++ b/contracts/interfaces/Bridge.sol
@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.7.0;
+
+// https://docs.tokenbridge.net/amb-bridge/development-of-a-cross-chain-application/how-to-develop-xchain-apps-by-amb#call-a-method-in-another-chain-using-the-amb-bridge
+interface IAMB {
+  function messageSender() external view returns (address);
+
+  function messageSourceChainId() external view returns (bytes32);
+}
+
+interface IOmniBridge {
+  function bridgeContract() external view returns (IAMB);
+}
diff --git a/test/full.test.js b/test/full.test.js
index 46cb21a..aa78457 100644
--- a/test/full.test.js
+++ b/test/full.test.js
@@ -50,14 +50,14 @@ describe('TornadoPool', function () {
   }
 
   async function fixtureUpgradeable() {
-    const { tornadoPool, omniBridge } = await loadFixture(fixture)
+    const { tornadoPool, omniBridge, amb } = await loadFixture(fixture)
     const [, gov] = await ethers.getSigners()
     const proxy = await deploy(
       'CrossChainUpgradeableProxy',
       tornadoPool.address,
       gov.address,
       [],
-      omniBridge.address,
+      amb.address,
       1,
     )
 
@@ -66,14 +66,14 @@ describe('TornadoPool', function () {
     const tornadoPoolProxied = TornadoPool.attach(proxy.address)
     await tornadoPoolProxied.initialize()
 
-    return { tornadoPool: tornadoPoolProxied, proxy, gov, omniBridge }
+    return { tornadoPool: tornadoPoolProxied, proxy, gov, omniBridge, amb }
   }
 
   describe('Upgradeability tests', () => {
     it('admin should be gov', async () => {
-      const { proxy, omniBridge, gov } = await loadFixture(fixtureUpgradeable)
+      const { proxy, amb, gov } = await loadFixture(fixtureUpgradeable)
       const { data } = await proxy.populateTransaction.admin()
-      const { result } = await omniBridge.callStatic.execute(proxy.address, data)
+      const { result } = await amb.callStatic.execute(proxy.address, data)
       expect('0x' + result.slice(26)).to.be.equal(gov.address.toLowerCase())
     })