From c01c9a3b07cedd58131612c9d578cc66c1785e1b Mon Sep 17 00:00:00 2001
From: MaxMustermann2 <82761650+MaxMustermann2@users.noreply.github.com>
Date: Tue, 30 Jul 2024 09:09:04 +0000
Subject: [PATCH] fix: check client chain before setPeer

---
 src/core/ExocoreGateway.sol            | 20 ++++++++++++++++++++
 src/interfaces/precompiles/IAssets.sol |  6 ++++++
 src/libraries/Errors.sol               |  6 ++++++
 test/foundry/unit/ExocoreGateway.t.sol |  6 ++++++
 test/mocks/AssetsMock.sol              |  4 ++++
 test/mocks/ExocoreGatewayMock.sol      | 12 ++++++++++++
 6 files changed, 54 insertions(+)

diff --git a/src/core/ExocoreGateway.sol b/src/core/ExocoreGateway.sol
index b71edb7d..711a75f8 100644
--- a/src/core/ExocoreGateway.sol
+++ b/src/core/ExocoreGateway.sol
@@ -164,6 +164,10 @@ contract ExocoreGateway is
         onlyOwner
         whenNotPaused
     {
+        // The registration of the client chain is done here and nowhere else.
+        // Elsewhere, the precompile is responsible for the checks. The precompile
+        // is not called here at all, and hence, such a check must be made manually.
+        _validateClientChainIdRegistered(clientChainId);
         super.setPeer(clientChainId, clientChainGateway);
     }
 
@@ -176,6 +180,7 @@ contract ExocoreGateway is
         string[] calldata names,
         string[] calldata metaData
     ) external payable onlyOwner whenNotPaused nonReentrant {
+        // The registration of the client chain is left for the precompile to validate.
         _validateWhitelistTokensInput(tokens, decimals, tvlLimits, names, metaData);
 
         bool success;
@@ -274,6 +279,21 @@ contract ExocoreGateway is
         }
     }
 
+    /// @dev Validates that the client chain id is registered.
+    /// @dev This is designed to be called only in the cases wherein the precompile isn't used.
+    /// @dev In all other situations, it is the responsibility of the precompile to perform such
+    ///      checks.
+    /// @param clientChainId The client chain id.
+    function _validateClientChainIdRegistered(uint32 clientChainId) internal view {
+        (bool success, bool isRegistered) = ASSETS_CONTRACT.isRegisteredClientChain(clientChainId);
+        if (!success) {
+            revert Errors.ExocoreGatewayFailedToCheckClientChainId();
+        }
+        if (!isRegistered) {
+            revert Errors.ExocoreGatewayNotRegisteredClientChainId();
+        }
+    }
+
     /// @dev The internal version of registerOrUpdateClientChain.
     /// @param clientChainId The client chain id.
     /// @param addressLength The length of the address type on the client chain.
diff --git a/src/interfaces/precompiles/IAssets.sol b/src/interfaces/precompiles/IAssets.sol
index ba0e2775..65580341 100644
--- a/src/interfaces/precompiles/IAssets.sol
+++ b/src/interfaces/precompiles/IAssets.sol
@@ -80,4 +80,10 @@ interface IAssets {
     /// @dev Returns the chain indices of the client chains.
     function getClientChains() external view returns (bool, uint32[] memory);
 
+    /// @dev Checks if the client chain is registered, given the chain ID.
+    /// @param clientChainID is the layerZero chainID if it is supported.
+    /// @return success true if the query is successful
+    /// @return isRegistered true if the client chain is registered
+    function isRegisteredClientChain(uint32 clientChainID) external view returns (bool success, bool isRegistered);
+
 }
diff --git a/src/libraries/Errors.sol b/src/libraries/Errors.sol
index 903dd64a..933d6983 100644
--- a/src/libraries/Errors.sol
+++ b/src/libraries/Errors.sol
@@ -163,6 +163,12 @@ library Errors {
     /// @dev ExocoreGateway: failed to decode client chain ids
     error ExocoreGatewayFailedToDecodeClientChainIds();
 
+    /// @dev ExocoreGateway: client chain should be registered before.
+    error ExocoreGatewayNotRegisteredClientChainId();
+
+    /// @dev ExocoreGateway: failed to check if the client id is registered
+    error ExocoreGatewayFailedToCheckClientChainId();
+
     /// @dev ExocoreGateway: thrown when associateOperatorWithEVMStaker failed
     error AssociateOperatorFailed(uint32 clientChainId, address staker, string operator);
 
diff --git a/test/foundry/unit/ExocoreGateway.t.sol b/test/foundry/unit/ExocoreGateway.t.sol
index 55168576..cc0774b7 100644
--- a/test/foundry/unit/ExocoreGateway.t.sol
+++ b/test/foundry/unit/ExocoreGateway.t.sol
@@ -501,6 +501,12 @@ contract SetPeer is SetUp {
         exocoreGateway.setPeer(anotherClientChain, newPeer);
     }
 
+    function test_RevertWhen_ClientChainNotRegistered() public {
+        vm.startPrank(exocoreValidatorSet.addr);
+        vm.expectRevert(Errors.ExocoreGatewayNotRegisteredClientChainId.selector);
+        exocoreGateway.setPeer(anotherClientChain, newPeer);
+    }
+
 }
 
 contract AddWhitelistTokens is SetUp {
diff --git a/test/mocks/AssetsMock.sol b/test/mocks/AssetsMock.sol
index e0688f8b..29283505 100644
--- a/test/mocks/AssetsMock.sol
+++ b/test/mocks/AssetsMock.sol
@@ -96,4 +96,8 @@ contract AssetsMock is IAssets {
         return abi.encodePacked(bytes32(bytes20(addr)));
     }
 
+    function isRegisteredClientChain(uint32 clientChainID) external view returns (bool, bool) {
+        return (true, isRegisteredChain[clientChainID]);
+    }
+
 }
diff --git a/test/mocks/ExocoreGatewayMock.sol b/test/mocks/ExocoreGatewayMock.sol
index 3e2a523c..45b795b7 100644
--- a/test/mocks/ExocoreGatewayMock.sol
+++ b/test/mocks/ExocoreGatewayMock.sol
@@ -22,6 +22,7 @@ import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/Own
 import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import {PausableUpgradeable} from "@openzeppelin/contracts-upgradeable/security/PausableUpgradeable.sol";
 import {ReentrancyGuardUpgradeable} from "@openzeppelin/contracts-upgradeable/security/ReentrancyGuardUpgradeable.sol";
+import {Errors} from "src/libraries/Errors.sol";
 import {OAppCoreUpgradeable} from "src/lzApp/OAppCoreUpgradeable.sol";
 
 contract ExocoreGatewayMock is
@@ -170,6 +171,7 @@ contract ExocoreGatewayMock is
         onlyOwner
         whenNotPaused
     {
+        _validateClientChainIdRegistered(clientChainId);
         super.setPeer(clientChainId, clientChainGateway);
     }
 
@@ -243,6 +245,16 @@ contract ExocoreGatewayMock is
         }
     }
 
+    function _validateClientChainIdRegistered(uint32 clientChainId) internal view {
+        (bool success, bool isRegistered) = ASSETS_CONTRACT.isRegisteredClientChain(clientChainId);
+        if (!success) {
+            revert Errors.ExocoreGatewayFailedToCheckClientChainId();
+        }
+        if (!isRegistered) {
+            revert Errors.ExocoreGatewayNotRegisteredClientChainId();
+        }
+    }
+
     function _registerOrUpdateClientChain(
         uint32 clientChainId,
         uint8 addressLength,