feat: implement message context v2

contracts/evm/ERC20Custody.sol

@@ -2,7 +2,7 @@
 pragma solidity 0.8.26;
 
 import { IERC20Custody } from "./interfaces/IERC20Custody.sol";
-import { IGatewayEVM, MessageContext } from "./interfaces/IGatewayEVM.sol";
+import { IGatewayEVM, MessageContextV2 } from "./interfaces/IGatewayEVM.sol";
 
 import { RevertContext } from "../../contracts/Revert.sol";
 
@@ -146,7 +146,7 @@ contract ERC20Custody is
     /// @param amount Amount of tokens to withdraw.
     /// @param data Calldata to pass to the contract call.
     function withdrawAndCall(
-        MessageContext calldata messageContext,
+        MessageContextV2 calldata messageContext,
         address to,
         address token,
         uint256 amount,

contracts/evm/GatewayEVM.sol

@@ -5,7 +5,7 @@ import { INotSupportedMethods } from "../../contracts/Errors.sol";
 import { RevertContext, RevertOptions, Revertable } from "../../contracts/Revert.sol";
 import { ZetaConnectorBase } from "./ZetaConnectorBase.sol";
 import { IERC20Custody } from "./interfaces/IERC20Custody.sol";
-import { Callable, IGatewayEVM, MessageContext } from "./interfaces/IGatewayEVM.sol";
+import "./interfaces/IGatewayEVM.sol";
 
 import "@openzeppelin/contracts-upgradeable/access/AccessControlUpgradeable.sol";
 import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
@@ -129,7 +129,7 @@ contract GatewayEVM is
     /// @param data Calldata to pass to the call.
     /// @return The result of the call.
     function execute(
-        MessageContext calldata messageContext,
+        MessageContextV2 calldata messageContext,
         address destination,
         bytes calldata data
     )
@@ -165,7 +165,7 @@ contract GatewayEVM is
     /// @param amount Amount of tokens to transfer.
     /// @param data Calldata to pass to the call.
     function executeWithERC20(
-        MessageContext calldata messageContext,
+        MessageContextV2 calldata messageContext,
         address token,
         address to,
         uint256 amount,
@@ -445,14 +445,20 @@ contract GatewayEVM is
     /// @param data Calldata to pass to the call.
     /// @return The result of the call.
     function _executeAuthenticatedCall(
-        MessageContext calldata messageContext,
+        MessageContextV2 calldata messageContext,
         address destination,
         bytes calldata data
     )
         private
         returns (bytes memory)
     {
-        return Callable(destination).onCall{ value: msg.value }(messageContext, data);
+        if (messageContext.amount == 0) {
+            return Callable(destination).onCall{ value: msg.value }(
+                MessageContext({ sender: messageContext.sender }), data
+            );
+        } else {
+            return CallableV2(destination).onCall{ value: msg.value }(messageContext, data);
+        }
     }
 
     // @dev prevent spoofing onCall and onRevert functions
@@ -463,7 +469,7 @@ contract GatewayEVM is
                 functionSelector := calldataload(data.offset)
             }
 
-            if (functionSelector == Callable.onCall.selector) {
+            if (functionSelector == Callable.onCall.selector || functionSelector == CallableV2.onCall.selector) {
                 revert NotAllowedToCallOnCall();
             }
 

contracts/evm/ZetaConnectorBase.sol

@@ -10,12 +10,7 @@ import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 
 import { RevertContext } from "../../contracts/Revert.sol";
-import {
-    IGatewayEVM,
-    IGatewayEVMErrors,
-    IGatewayEVMEvents,
-    MessageContext
-} from "../../contracts/evm/interfaces/IGatewayEVM.sol";
+import { IGatewayEVM, MessageContextV2 } from "../../contracts/evm/interfaces/IGatewayEVM.sol";
 import "../../contracts/evm/interfaces/IZetaConnector.sol";
 
 /// @title ZetaConnectorBase
@@ -123,7 +118,7 @@ abstract contract ZetaConnectorBase is
     /// @param data The calldata to pass to the contract call.
     /// @param internalSendHash A hash used for internal tracking of the transaction.
     function withdrawAndCall(
-        MessageContext calldata messageContext,
+        MessageContextV2 calldata messageContext,
         address to,
         uint256 amount,
         bytes calldata data,

contracts/evm/ZetaConnectorNative.sol

@@ -54,7 +54,7 @@ contract ZetaConnectorNative is ZetaConnectorBase {
     // https://github.com/zeta-chain/protocol-contracts/issues/398)
     /// @dev This function can only be called by the TSS address.
     function withdrawAndCall(
-        MessageContext calldata messageContext,
+        MessageContextV2 calldata messageContext,
         address to,
         uint256 amount,
         bytes calldata data,

contracts/evm/ZetaConnectorNonNative.sol

@@ -68,7 +68,7 @@ contract ZetaConnectorNonNative is ZetaConnectorBase {
     /// @param internalSendHash A hash used for internal tracking of the transaction.
     /// @dev This function can only be called by the TSS address, and mints if supply is not reached.
     function withdrawAndCall(
-        MessageContext calldata messageContext,
+        MessageContextV2 calldata messageContext,
         address to,
         uint256 amount,
         bytes calldata data,

contracts/evm/interfaces/IERC20Custody.sol

@@ -3,7 +3,7 @@ pragma solidity 0.8.26;
 
 import { RevertContext } from "../../../contracts/Revert.sol";
 
-import { MessageContext } from "./IGatewayEVM.sol";
+import { MessageContextV2 } from "./IGatewayEVM.sol";
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 
 /// @title IERC20CustodyEvents
@@ -79,7 +79,7 @@ interface IERC20Custody is IERC20CustodyEvents, IERC20CustodyErrors {
     /// @param amount Amount of tokens to withdraw.
     /// @param data Calldata to pass to the contract call.
     function withdrawAndCall(
-        MessageContext calldata messageContext,
+        MessageContextV2 calldata messageContext,
         address token,
         address to,
         uint256 amount,

contracts/evm/interfaces/IGatewayEVM.sol

@@ -122,7 +122,7 @@ interface IGatewayEVM is IGatewayEVMErrors, IGatewayEVMEvents {
     /// @param amount The amount of tokens to transfer.
     /// @param data The calldata to pass to the contract call.
     function executeWithERC20(
-        MessageContext calldata messageContext,
+        MessageContextV2 calldata messageContext,
         address token,
         address to,
         uint256 amount,
@@ -150,7 +150,7 @@ interface IGatewayEVM is IGatewayEVMErrors, IGatewayEVMEvents {
     /// @param data Calldata to pass to the call.
     /// @return The result of the call.
     function execute(
-        MessageContext calldata messageContext,
+        MessageContextV2 calldata messageContext,
         address destination,
         bytes calldata data
     )
@@ -229,3 +229,24 @@ struct MessageContext {
 interface Callable {
     function onCall(MessageContext calldata context, bytes calldata message) external payable returns (bytes memory);
 }
+
+/// @notice Message context passed to execute function.
+/// @param sender Sender from omnichain contract.
+/// @param asset The address of the asset.
+/// @param amount The amount of the asset.
+struct MessageContextV2 {
+    address sender;
+    address asset;
+    uint256 amount;
+}
+
+/// @notice Interface implemented by contracts receiving authenticated calls with MessageContextV2.
+interface CallableV2 {
+    function onCall(
+        MessageContextV2 calldata context,
+        bytes calldata message
+    )
+        external
+        payable
+        returns (bytes memory);
+}

contracts/zevm/GatewayZEVM.sol

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity 0.8.26;
 
-import { CallOptions, IGatewayZEVM } from "./interfaces/IGatewayZEVM.sol";
+import { CallOptions, CallOptionsV2, IGatewayZEVM } from "./interfaces/IGatewayZEVM.sol";
 
 import { INotSupportedMethods } from "../../contracts/Errors.sol";
 import { AbortContext, Abortable, RevertContext, RevertOptions, Revertable } from "../../contracts/Revert.sol";
@@ -244,6 +244,44 @@ contract GatewayZEVM is
         );
     }
 
+    /// @notice Withdraw ZRC20 tokens and call a smart contract on an external chain.
+    /// @param receiver The receiver address on the external chain.
+    /// @param amount The amount of tokens to withdraw.
+    /// @param zrc20 The address of the ZRC20 token.
+    /// @param message The calldata to pass to the contract call.
+    /// @param callOptions Call options including gas limit, arbirtrary call flag and message context version.
+    /// @param revertOptions Revert options.
+    function withdrawAndCall(
+        bytes memory receiver,
+        uint256 amount,
+        address zrc20,
+        bytes calldata message,
+        CallOptionsV2 calldata callOptions,
+        RevertOptions calldata revertOptions
+    )
+        external
+        whenNotPaused
+    {
+        if (receiver.length == 0) revert ZeroAddress();
+        if (amount == 0) revert InsufficientZRC20Amount();
+        if (callOptions.gasLimit < MIN_GAS_LIMIT) revert InsufficientGasLimit();
+        if (message.length + revertOptions.revertMessage.length > MAX_MESSAGE_SIZE) revert MessageSizeExceeded();
+
+        uint256 gasFee = _withdrawZRC20WithGasLimit(amount, zrc20, callOptions.gasLimit);
+        emit WithdrawnAndCalledV2(
+            msg.sender,
+            0,
+            receiver,
+            zrc20,
+            amount,
+            gasFee,
+            IZRC20(zrc20).PROTOCOL_FLAT_FEE(),
+            message,
+            callOptions,
+            revertOptions
+        );
+    }
+
     /// @notice Withdraw ZETA tokens to an external chain.
     //// @param receiver The receiver address on the external chain.
     //// @param amount The amount of tokens to withdraw.

contracts/zevm/interfaces/IGatewayZEVM.sol

@@ -70,6 +70,30 @@ interface IGatewayZEVMEvents {
         CallOptions callOptions,
         RevertOptions revertOptions
     );
+
+    /// @notice Emitted when a withdraw and call is made.
+    /// @param sender The address from which the tokens are withdrawn.
+    /// @param chainId Chain id of external chain.
+    /// @param receiver The receiver address on the external chain.
+    /// @param zrc20 The address of the ZRC20 token.
+    /// @param value The amount of tokens withdrawn.
+    /// @param gasfee The gas fee for the withdrawal.
+    /// @param protocolFlatFee The protocol flat fee for the withdrawal.
+    /// @param message The calldata passed to the contract call.
+    /// @param callOptions Call options including gas limit, arbirtrary call flag and message context version.
+    /// @param revertOptions Revert options.
+    event WithdrawnAndCalledV2(
+        address indexed sender,
+        uint256 indexed chainId,
+        bytes receiver,
+        address zrc20,
+        uint256 value,
+        uint256 gasfee,
+        uint256 protocolFlatFee,
+        bytes message,
+        CallOptionsV2 callOptions,
+        RevertOptions revertOptions
+    );
 }
 
 /// @title IGatewayZEVMErrors
@@ -161,6 +185,23 @@ interface IGatewayZEVM is IGatewayZEVMErrors, IGatewayZEVMEvents {
     )
         external;
 
+    /// @notice Withdraw ZRC20 tokens and call a smart contract on an external chain.
+    /// @param receiver The receiver address on the external chain.
+    /// @param amount The amount of tokens to withdraw.
+    /// @param zrc20 The address of the ZRC20 token.
+    /// @param message The calldata to pass to the contract call.
+    /// @param callOptions Call options including gas limit, arbirtrary call flag and message context version.
+    /// @param revertOptions Revert options.
+    function withdrawAndCall(
+        bytes memory receiver,
+        uint256 amount,
+        address zrc20,
+        bytes calldata message,
+        CallOptionsV2 calldata callOptions,
+        RevertOptions calldata revertOptions
+    )
+        external;
+
     /// @notice Withdraw ZETA tokens and call a smart contract on an external chain.
     /// @param receiver The receiver address on the external chain.
     /// @param amount The amount of tokens to withdraw.
@@ -268,3 +309,13 @@ struct CallOptions {
     uint256 gasLimit;
     bool isArbitraryCall;
 }
+
+/// @notice CallOptions struct passed to withdrawAndCall function.
+/// @param gasLimit Gas limit.
+/// @param isArbitraryCall Indicates if call should be arbitrary or authenticated.
+/// @param isMsgContextV2 Indicates if call is using new message context.
+struct CallOptionsV2 {
+    uint256 gasLimit;
+    bool isArbitraryCall;
+    bool isMsgContextV2;
+}

test/ERC20Custody.t.sol

@@ -35,7 +35,8 @@ contract ERC20CustodyTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiv
     address tssAddress;
     address foo;
     RevertContext revertContext;
-    MessageContext arbitraryCallMessageContext = MessageContext({ sender: address(0) });
+    MessageContextV2 arbitraryCallMessageContext =
+        MessageContextV2({ sender: address(0), asset: address(0), amount: 0 });
 
     error EnforcedPause();
     error NotWhitelisted();
@@ -349,11 +350,17 @@ contract ERC20CustodyTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiv
         uint256 balanceBeforeCustody = token.balanceOf(address(custody));
 
         vm.expectEmit(true, true, true, true, address(receiver));
-        emit ReceivedOnCall(sender, message);
+        emit ReceivedOnCallV2(sender, address(token), amount, message);
         vm.expectEmit(true, true, true, true, address(custody));
         emit WithdrawnAndCalled(address(receiver), address(token), amount, message);
         vm.prank(tssAddress);
-        custody.withdrawAndCall(MessageContext({ sender: sender }), address(receiver), address(token), amount, message);
+        custody.withdrawAndCall(
+            MessageContextV2({ sender: sender, asset: address(token), amount: amount }),
+            address(receiver),
+            address(token),
+            amount,
+            message
+        );
 
         // Verify that the tokens were not transferred to the destination address
         uint256 balanceAfter = token.balanceOf(destination);