test: connector and revert unit tests (#244)

contracts/prototypes/evm/GatewayEVM.sol

@@ -20,7 +20,6 @@ contract GatewayEVM is Initializable, OwnableUpgradeable, UUPSUpgradeable, IGate
 
     /// @notice The address of the custody contract.
     address public custody;
-
     /// @notice The address of the TSS (Threshold Signature Scheme) contract.
     address public tssAddress;
     /// @notice The address of the ZetaConnector contract.
@@ -86,11 +85,10 @@ contract GatewayEVM is Initializable, OwnableUpgradeable, UUPSUpgradeable, IGate
         uint256 amount,
         bytes calldata data
     ) public nonReentrant {
-        if (amount == 0) revert InsufficientETHAmount();
+        if (amount == 0) revert InsufficientERC20Amount();
         // Approve the target contract to spend the tokens
         if(!resetApproval(token, to)) revert ApprovalFailed();
         if(!IERC20(token).approve(to, amount)) revert ApprovalFailed();
-
         // Execute the call on the target contract
         bytes memory result = _execute(to, data);
 
@@ -100,7 +98,7 @@ contract GatewayEVM is Initializable, OwnableUpgradeable, UUPSUpgradeable, IGate
         // Transfer any remaining tokens back to the custody/connector contract
         uint256 remainingBalance = IERC20(token).balanceOf(address(this));
         if (remainingBalance > 0) {
-            transferToAssetHandler(token, amount);
+            transferToAssetHandler(token, remainingBalance);
         }
 
         emit ExecutedWithERC20(token, to, amount, data);

contracts/prototypes/evm/IReceiverEVM.sol

@@ -6,4 +6,5 @@ interface IReceiverEVMEvents {
     event ReceivedNonPayable(address sender, string[] strs, uint256[] nums, bool flag);
     event ReceivedERC20(address sender, uint256 amount, address token, address destination);
     event ReceivedNoParams(address sender);
+    event ReceivedRevert(address sender, bytes data);
 }

contracts/prototypes/evm/ReceiverEVM.sol

@@ -3,14 +3,13 @@ pragma solidity ^0.8.0;
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import "@openzeppelin/contracts/security/ReentrancyGuard.sol";
+import "./IReceiverEVM.sol";
 
-contract ReceiverEVM {
+// @notice This contract is used just for testing
+contract ReceiverEVM is IReceiverEVMEvents, ReentrancyGuard {
     using SafeERC20 for IERC20;
-
-    event ReceivedPayable(address sender, uint256 value, string str, uint256 num, bool flag);
-    event ReceivedNonPayable(address sender, string[] strs, uint256[] nums, bool flag);
-    event ReceivedERC20(address sender, uint256 amount, address token, address destination);
-    event ReceivedNoParams(address sender);
+    error ZeroAmount();
 
     // Payable function
     function receivePayable(string memory str, uint256 num, bool flag) external payable {
@@ -23,15 +22,33 @@ contract ReceiverEVM {
     }
 
     // Function using IERC20
-    function receiveERC20(uint256 amount, address token, address destination) external {
+    function receiveERC20(uint256 amount, address token, address destination) external nonReentrant {
         // Transfer tokens from the Gateway contract to the destination address
         IERC20(token).safeTransferFrom(msg.sender, destination, amount);
 
         emit ReceivedERC20(msg.sender, amount, token, destination);
     }
 
+    // Function using IERC20 to partially transfer tokens
+    function receiveERC20Partial(uint256 amount, address token, address destination) external nonReentrant {
+        uint256 amountToSend = amount / 2;
+        if (amountToSend == 0) revert ZeroAmount();
+
+        IERC20(token).safeTransferFrom(msg.sender, destination, amountToSend);
+
+        emit ReceivedERC20(msg.sender, amountToSend, token, destination);
+    }
+
     // Function without parameters
     function receiveNoParams() external {
         emit ReceivedNoParams(msg.sender);
     }
+
+    // onRevertCallback
+    function onRevert(bytes calldata data) external {
+        emit ReceivedRevert(msg.sender, data);
+    }
+
+    receive() external payable {}
+    fallback() external payable {}
 }

contracts/prototypes/evm/TestERC20.sol

@@ -3,6 +3,7 @@ pragma solidity ^0.8.0;
 
 import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 
+// @notice This contract is used just for testing
 contract TestERC20 is ERC20 {
     constructor(string memory name, string memory symbol) ERC20(name, symbol) {}
 

contracts/prototypes/evm/ZetaConnectorNative.sol

@@ -12,13 +12,13 @@ contract ZetaConnectorNative is ZetaConnectorNewBase {
         ZetaConnectorNewBase(_gateway, _zetaToken)
     {}
 
-    // Withdraw is called by TSS address, it directly transfers zetaToken to the destination address without contract call
+    // @dev withdraw is called by TSS address, it directly transfers zetaToken to the destination address without contract call
     function withdraw(address to, uint256 amount, bytes32 internalSendHash) external override nonReentrant {
         IERC20(zetaToken).safeTransfer(to, amount);
         emit Withdraw(to, amount);
     }
 
-    // WithdrawAndCall is called by TSS address, it transfers zetaToken to the gateway and calls a contract
+    // @dev withdrawAndCall is called by TSS address, it transfers zetaToken to the gateway and calls a contract
     function withdrawAndCall(address to, uint256 amount, bytes calldata data, bytes32 internalSendHash) external override nonReentrant {
         // Transfer zetaToken to the Gateway contract
         IERC20(zetaToken).safeTransfer(address(gateway), amount);
@@ -29,9 +29,19 @@ contract ZetaConnectorNative is ZetaConnectorNewBase {
         emit WithdrawAndCall(to, amount, data);
     }
 
-    // Function to handle token transfer
-    function receiveTokens(uint256 amount) external override nonReentrant {
-        // Transfer tokens from the sender to this contract
+    // @dev withdrawAndRevert is called by TSS address, it transfers zetaToken to the gateway and calls onRevert on a contract
+    function withdrawAndRevert(address to, uint256 amount, bytes calldata data, bytes32 internalSendHash) external override nonReentrant {
+        // Transfer zetaToken to the Gateway contract
+        IERC20(zetaToken).safeTransfer(address(gateway), amount);
+
+        // Forward the call to the Gateway contract
+        gateway.revertWithERC20(address(zetaToken), to, amount, data);
+
+        emit WithdrawAndRevert(to, amount, data);
+    }
+
+    // @dev receiveTokens handles token transfer and burn them
+    function receiveTokens(uint256 amount) external override {
         IERC20(zetaToken).safeTransferFrom(msg.sender, address(this), amount);
     }
 }

contracts/prototypes/evm/ZetaConnectorNewBase.sol

@@ -16,6 +16,7 @@ abstract contract ZetaConnectorNewBase is ReentrancyGuard {
 
     event Withdraw(address indexed to, uint256 amount);
     event WithdrawAndCall(address indexed to, uint256 amount, bytes data);
+    event WithdrawAndRevert(address indexed to, uint256 amount, bytes data);
 
     constructor(address _gateway, address _zetaToken) {
         if (_gateway == address(0) || _zetaToken == address(0)) {
@@ -29,5 +30,7 @@ abstract contract ZetaConnectorNewBase is ReentrancyGuard {
 
     function withdrawAndCall(address to, uint256 amount, bytes calldata data, bytes32 internalSendHash) external virtual;
 
+    function withdrawAndRevert(address to, uint256 amount, bytes calldata data, bytes32 internalSendHash) external virtual;
+
     function receiveTokens(uint256 amount) external virtual;
 }

contracts/prototypes/evm/ZetaConnectorNonNative.sol

@@ -10,13 +10,13 @@ contract ZetaConnectorNonNative is ZetaConnectorNewBase {
         ZetaConnectorNewBase(_gateway, _zetaToken)
     {}
 
-    // Withdraw is called by TSS address, it mints zetaToken to the destination address
+    // @dev withdraw is called by TSS address, it mints zetaToken to the destination address
     function withdraw(address to, uint256 amount, bytes32 internalSendHash) external override nonReentrant {
         IZetaNonEthNew(zetaToken).mint(to, amount, internalSendHash);
         emit Withdraw(to, amount);
     }
 
-    // WithdrawAndCall is called by TSS address, it mints zetaToken and calls a contract
+    // @dev withdrawAndCall is called by TSS address, it mints zetaToken and calls a contract
     function withdrawAndCall(address to, uint256 amount, bytes calldata data, bytes32 internalSendHash) external override nonReentrant {
         // Mint zetaToken to the Gateway contract
         IZetaNonEthNew(zetaToken).mint(address(gateway), amount, internalSendHash);
@@ -27,9 +27,19 @@ contract ZetaConnectorNonNative is ZetaConnectorNewBase {
         emit WithdrawAndCall(to, amount, data);
     }
 
-    // Function to handle token transfer and burn them
+    // @dev withdrawAndRevert is called by TSS address, it mints zetaToken to the gateway and calls onRevert on a contract
+    function withdrawAndRevert(address to, uint256 amount, bytes calldata data, bytes32 internalSendHash) external override nonReentrant {
+        // Mint zetaToken to the Gateway contract
+        IZetaNonEthNew(zetaToken).mint(address(gateway), amount, internalSendHash);
+
+        // Forward the call to the Gateway contract
+        gateway.revertWithERC20(address(zetaToken), to, amount, data);
+
+        emit WithdrawAndRevert(to, amount, data);
+    }
+
+    // @dev receiveTokens handles token transfer and burn them
     function receiveTokens(uint256 amount) external override {
-        // Burn the tokens
         IZetaNonEthNew(zetaToken).burnFrom(msg.sender, amount);
     }
 }

contracts/prototypes/zevm/GatewayZEVM.sol

@@ -11,7 +11,6 @@ import "../../zevm/interfaces/zContract.sol";
 import "./IGatewayZEVM.sol";
 import "../../zevm/interfaces/IWZETA.sol";
 
-
 // The GatewayZEVM contract is the endpoint to call smart contracts on omnichain
 // The contract doesn't hold any funds and should never have active allowances
 contract GatewayZEVM is IGatewayZEVMEvents, IGatewayZEVMErrors, Initializable, OwnableUpgradeable, UUPSUpgradeable, ReentrancyGuardUpgradeable {
@@ -38,6 +37,11 @@ contract GatewayZEVM is IGatewayZEVMEvents, IGatewayZEVMErrors, Initializable, O
 
     function _authorizeUpgrade(address newImplementation) internal override onlyOwner() {}
 
+    /// @dev Receive function to receive ZETA from WETH9.withdraw().
+    receive() external payable {
+        if (msg.sender != zetaToken && msg.sender != FUNGIBLE_MODULE_ADDRESS) revert OnlyWZETAOrFungible();
+    }
+
     function _withdrawZRC20(uint256 amount, address zrc20) internal returns (uint256) {
         (address gasZRC20, uint256 gasFee) = IZRC20(zrc20).withdrawGasFee();
         if (!IZRC20(gasZRC20).transferFrom(msg.sender, FUNGIBLE_MODULE_ADDRESS, gasFee)) {
@@ -75,13 +79,13 @@ contract GatewayZEVM is IGatewayZEVMEvents, IGatewayZEVMErrors, Initializable, O
     // Withdraw ZETA to external chain
     function withdraw(uint256 amount) external nonReentrant {
         _transferZETA(amount, FUNGIBLE_MODULE_ADDRESS);
-        emit Withdrawal(msg.sender, address(0), abi.encodePacked(FUNGIBLE_MODULE_ADDRESS), amount, 0, 0, "");
+        emit Withdrawal(msg.sender, address(zetaToken), abi.encodePacked(FUNGIBLE_MODULE_ADDRESS), amount, 0, 0, "");
     }
 
     // Withdraw ZETA and call smart contract on external chain
     function withdrawAndCall(uint256 amount, bytes calldata message) external nonReentrant {
         _transferZETA(amount, FUNGIBLE_MODULE_ADDRESS);
-        emit Withdrawal(msg.sender, address(0), abi.encodePacked(FUNGIBLE_MODULE_ADDRESS), amount, 0, 0, message);
+        emit Withdrawal(msg.sender, address(zetaToken), abi.encodePacked(FUNGIBLE_MODULE_ADDRESS), amount, 0, 0, message);
     }
 
     // Call smart contract on external chain without asset transfer

contracts/prototypes/zevm/IGatewayZEVM.sol

@@ -47,4 +47,5 @@ interface IGatewayZEVMErrors {
     error CallerIsNotFungibleModule();
     error InvalidTarget();
     error FailedZetaSent();
+    error OnlyWZETAOrFungible();
 }

contracts/prototypes/zevm/SenderZEVM.sol

@@ -5,6 +5,7 @@ import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "./IGatewayZEVM.sol";
 import "../../zevm/interfaces/IZRC20.sol";
 
+// @notice This contract is used just for testing
 contract SenderZEVM {
     address public gateway;
     error ApprovalFailed();

contracts/prototypes/zevm/TestZContract.sol

@@ -4,8 +4,10 @@ pragma solidity 0.8.7;
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "../../zevm/interfaces/zContract.sol";
 
+// @notice This contract is used just for testing
 contract TestZContract is UniversalContract {
     event ContextData(bytes origin, address sender, uint256 chainID, address msgSender, string message);
+    event ContextDataRevert(bytes origin, address sender, uint256 chainID, address msgSender, string message);
 
     function onCrossChainCall(
         zContext calldata context,
@@ -30,6 +32,9 @@ contract TestZContract is UniversalContract {
         if (message.length > 0) {
             decodedMessage = abi.decode(message, (string));
         }
-        emit ContextData(context.origin, context.sender, context.chainID, msg.sender, decodedMessage);
+        emit ContextDataRevert(context.origin, context.sender, context.chainID, msg.sender, decodedMessage);
     }
+
+    receive() external payable {}
+    fallback() external payable {}
 }

contracts/zevm/WZETA.sol

@@ -1,61 +1,68 @@
-pragma solidity ^0.4.18;
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
 
 contract WETH9 {
-    string public name = "Wrapped Zeta";
-    string public symbol = "WZETA";
+    string public name = "Wrapped Ether";
+    string public symbol = "WETH";
     uint8 public decimals = 18;
 
-    event Approval(address indexed src, address indexed guy, uint wad);
-    event Transfer(address indexed src, address indexed dst, uint wad);
-    event Deposit(address indexed dst, uint wad);
-    event Withdrawal(address indexed src, uint wad);
+    event Approval(address indexed src, address indexed guy, uint256 wad);
+    event Transfer(address indexed src, address indexed dst, uint256 wad);
+    event Deposit(address indexed dst, uint256 wad);
+    event Withdrawal(address indexed src, uint256 wad);
 
-    mapping(address => uint) public balanceOf;
-    mapping(address => mapping(address => uint)) public allowance;
+    mapping(address => uint256) public balanceOf;
+    mapping(address => mapping(address => uint256)) public allowance;
 
-    function() public payable {
+    receive() external payable {
         deposit();
     }
 
     function deposit() public payable {
         balanceOf[msg.sender] += msg.value;
-        Deposit(msg.sender, msg.value);
+        emit Deposit(msg.sender, msg.value);
     }
 
-    function withdraw(uint wad) public {
-        require(balanceOf[msg.sender] >= wad);
+    function withdraw(uint256 wad) public {
+        require(balanceOf[msg.sender] >= wad, "");
         balanceOf[msg.sender] -= wad;
-        msg.sender.transfer(wad);
-        Withdrawal(msg.sender, wad);
+        payable(msg.sender).transfer(wad);
+        emit Withdrawal(msg.sender, wad);
     }
 
-    function totalSupply() public view returns (uint) {
-        return this.balance;
+    function totalSupply() public view returns (uint256) {
+        return address(this).balance;
     }
 
-    function approve(address guy, uint wad) public returns (bool) {
+    function approve(address guy, uint256 wad) public returns (bool) {
         allowance[msg.sender][guy] = wad;
-        Approval(msg.sender, guy, wad);
+        emit Approval(msg.sender, guy, wad);
         return true;
     }
 
-    function transfer(address dst, uint wad) public returns (bool) {
+    function transfer(address dst, uint256 wad) public returns (bool) {
         return transferFrom(msg.sender, dst, wad);
     }
 
-    function transferFrom(address src, address dst, uint wad) public returns (bool) {
-        require(balanceOf[src] >= wad);
+    function transferFrom(
+        address src,
+        address dst,
+        uint256 wad
+    ) public returns (bool) {
+        require(balanceOf[src] >= wad, "");
 
-        if (src != msg.sender && allowance[src][msg.sender] != uint(-1)) {
-            require(allowance[src][msg.sender] >= wad);
+        if (
+            src != msg.sender && allowance[src][msg.sender] != type(uint256).max
+        ) {
+            require(allowance[src][msg.sender] >= wad, "");
             allowance[src][msg.sender] -= wad;
         }
 
         balanceOf[src] -= wad;
         balanceOf[dst] += wad;
 
-        Transfer(src, dst, wad);
+        emit Transfer(src, dst, wad);
 
         return true;
     }
-}
+}