onlyGateway

zeta-chain · Nov 4, 2024 · c01dac3 · c01dac3
1 parent e959412
commit c01dac3
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/examples/token/contracts/Connected.sol b/examples/token/contracts/Connected.sol
@@ -10,6 +10,11 @@ contract Connected is ERC20, Ownable {
     GatewayEVM public immutable gateway;
     address public counterparty;
 
+    modifier onlyGateway() {
+        require(msg.sender == address(gateway), "Caller is not the gateway");
+        _;
+    }
+
     function setCounterparty(address contractAddress) external onlyOwner {
         counterparty = contractAddress;
     }
@@ -55,7 +60,7 @@ contract Connected is ERC20, Ownable {
     function onCall(
         MessageContext calldata messageContext,
         bytes calldata message
-    ) external payable returns (bytes4) {
+    ) external payable onlyGateway returns (bytes4) {
         if (messageContext.sender != counterparty) revert("Unauthorized");
         (address receiver, uint256 amount) = abi.decode(
             message,
@@ -64,7 +69,7 @@ contract Connected is ERC20, Ownable {
         _mint(receiver, amount);
     }
 
-    function onRevert(RevertContext calldata context) external {}
+    function onRevert(RevertContext calldata context) external onlyGateway {}
 
     receive() external payable {}
 

diff --git a/examples/token/contracts/Universal.sol b/examples/token/contracts/Universal.sol
@@ -24,6 +24,11 @@ contract Universal is ERC20, Ownable, UniversalContract {
 
     event CounterpartySet(address indexed zrc20, bytes indexed contractAddress);
 
+    modifier onlyGateway() {
+        require(msg.sender == address(gateway), "Caller is not the gateway");
+        _;
+    }
+
     constructor(
         address payable gatewayAddress,
         address initialOwner
@@ -81,7 +86,7 @@ contract Universal is ERC20, Ownable, UniversalContract {
         address zrc20,
         uint256 amount,
         bytes calldata message
-    ) external override {
+    ) external override onlyGateway {
         if (keccak256(messageContext.origin) != keccak256(counterparty[zrc20]))
             revert("Unauthorized");
         (address receiver, uint256 tokenAmount, address destination) = abi
@@ -110,5 +115,5 @@ contract Universal is ERC20, Ownable, UniversalContract {
         }
     }
 
-    function onRevert(RevertContext calldata context) external {}
+    function onRevert(RevertContext calldata context) external onlyGateway {}
 }