diff --git a/examples/token/contracts/Connected.sol b/examples/token/contracts/Connected.sol index 7842b837..e703e425 100644 --- 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 index 3920b576..681a6a57 100644 --- 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 {} }