feat: add zeta handling to v2 contracts

contracts/prototypes/evm/ERC20CustodyNew.sol

@@ -2,7 +2,7 @@
 pragma solidity ^0.8.0;
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
-import "./interfaces.sol";
+import "./IGatewayEVM.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@openzeppelin/contracts/security/ReentrancyGuard.sol";
 

contracts/prototypes/evm/GatewayEVM.sol

@@ -6,28 +6,43 @@
 import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
-import "./interfaces.sol";
+import "./IGatewayEVM.sol";
 
-// The GatewayEVM contract is the endpoint to call smart contracts on external chains
-// The contract doesn't hold any funds and should never have active allowances
+/**
+ * @title GatewayEVM
+ * @notice The GatewayEVM contract is the endpoint to call smart contracts on external chains.
+ * @dev The contract doesn't hold any funds and should never have active allowances.
+ */
 contract GatewayEVM is Initializable, OwnableUpgradeable, UUPSUpgradeable, IGatewayEVMErrors, IGatewayEVMEvents {
     using SafeERC20 for IERC20;
 
+    /// @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.
+    address public zetaConnector;
+
+    /// @notice The address of the Zeta token contract.
+    address public zetaToken;
+
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor() {
         _disableInitializers();
     }
 
-    function initialize(address _tssAddress) public initializer {
+    function initialize(address _tssAddress, address _zetaToken) public initializer {
         __Ownable_init();
         __UUPSUpgradeable_init();
 
-        if (_tssAddress == address(0)) revert ZeroAddress();
+        if (_tssAddress == address(0) || _zetaToken == address(0)) {
+            revert ZeroAddress();
+        }
 
         tssAddress = _tssAddress;
+        zetaToken = _zetaToken;
     }
 
     function _authorizeUpgrade(address newImplementation) internal override onlyOwner() {}
@@ -59,7 +74,7 @@
         address to,
         uint256 amount,
         bytes calldata data
-    ) public returns (bytes memory) {
+    ) public {
         if (amount == 0) revert InsufficientETHAmount();
         // Approve the target contract to spend the tokens
         if(!resetApproval(token, to)) revert ApprovalFailed();
@@ -71,15 +86,13 @@
         // Reset approval
         if(!resetApproval(token, to)) revert ApprovalFailed();
 
-        // Transfer any remaining tokens back to the custody contract
+        // Transfer any remaining tokens back to the custody/connector contract
         uint256 remainingBalance = IERC20(token).balanceOf(address(this));
         if (remainingBalance > 0) {
-            IERC20(token).safeTransfer(address(custody), remainingBalance);
+            IERC20(token).safeTransfer(getAssetHandler(token), remainingBalance);
         }
 
         emit ExecutedWithERC20(token, to, amount, data);
-
-        return result;
     }
 
     // Deposit ETH to tss
@@ -92,10 +105,11 @@
         emit Deposit(msg.sender, receiver, msg.value, address(0), "");
     }
 
-    // Deposit ERC20 tokens to custody
+    // Deposit ERC20 tokens to custody/connector
     function deposit(address receiver, uint256 amount, address asset) external {
         if (amount == 0) revert InsufficientERC20Amount();
-        IERC20(asset).safeTransferFrom(msg.sender, address(custody), amount);
+
+        IERC20(asset).safeTransferFrom(msg.sender, getAssetHandler(asset), amount);
 
         emit Deposit(msg.sender, receiver, amount, asset, "");
     }
@@ -110,10 +124,11 @@
         emit Deposit(msg.sender, receiver, msg.value, address(0), payload);
     }
 
-    // Deposit ERC20 tokens to custody and call an omnichain smart contract
+    // Deposit ERC20 tokens to custody/connector and call an omnichain smart contract
     function depositAndCall(address receiver, uint256 amount, address asset, bytes calldata payload) external {
         if (amount == 0) revert InsufficientERC20Amount();
-        IERC20(asset).safeTransferFrom(msg.sender, address(custody), amount);
+
+        IERC20(asset).safeTransferFrom(msg.sender, getAssetHandler(asset), amount);
 
         emit Deposit(msg.sender, receiver, amount, asset, payload);
     }
@@ -125,10 +140,28 @@
 
     function setCustody(address _custody) external {
         if (custody != address(0)) revert CustodyInitialized();
+        if (_custody == address(0)) revert ZeroAddress();
+
         custody = _custody;
     }
 
+     function setConnector(address _zetaConnector) external {
+        if (zetaConnector != address(0)) revert CustodyInitialized();
+        if (_zetaConnector == address(0)) revert ZeroAddress();
+
+        zetaConnector = _zetaConnector;
+    }
+
     function resetApproval(address token, address to) private returns (bool) {
         return IERC20(token).approve(to, 0);
     }
+
+    function getAssetHandler(address token) private returns (address) {
+        address assetHandler = address(custody);
+        if (token == zetaToken) {
+            assetHandler = address(zetaConnector);
+        }
+
+        return assetHandler;
+    }
 }

contracts/prototypes/evm/GatewayEVMUpgradeTest.sol

@@ -6,41 +6,32 @@
 import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
-
+import "./IGatewayEVM.sol";
 
 // NOTE: Purpose of this contract is to test upgrade process, the only difference should be name of Executed event
 // The Gateway contract is the endpoint to call smart contracts on external chains
 // The contract doesn't hold any funds and should never have active allowances
 /// @custom:oz-upgrades-from GatewayEVM
-contract GatewayEVMUpgradeTest is Initializable, OwnableUpgradeable, UUPSUpgradeable {
+contract GatewayEVMUpgradeTest is Initializable, OwnableUpgradeable, UUPSUpgradeable, IGatewayEVMErrors, IGatewayEVMEvents {
     using SafeERC20 for IERC20;
 
-    error ExecutionFailed();
-    error SendFailed();
-    error InsufficientETHAmount();
-    error ZeroAddress();
-    error ApprovalFailed();
-
     address public custody;
     address public tssAddress;
+    address public zetaConnector;
+    address public zetaToken;
 
     event ExecutedV2(address indexed destination, uint256 value, bytes data);
-    event ExecutedWithERC20(address indexed token, address indexed to, uint256 amount, bytes data);
-    event SendERC20(bytes recipient, address indexed asset, uint256 amount);
-    event Send(bytes recipient, uint256 amount);
 
-    /// @custom:oz-upgrades-unsafe-allow constructor
-    constructor() {
-        _disableInitializers();
-    }
+    constructor() {}
 
-    function initialize(address _tssAddress) public initializer {
+    function initialize(address _tssAddress, address _zetaToken) public initializer {
         __Ownable_init();
         __UUPSUpgradeable_init();
 
         if (_tssAddress == address(0)) revert ZeroAddress();
-        
+
         tssAddress = _tssAddress;
+        zetaToken = _zetaToken;
     }
 
     function _authorizeUpgrade(address newImplementation) internal override onlyOwner() {}
@@ -73,47 +64,95 @@
         address to,
         uint256 amount,
         bytes calldata data
-    ) external returns (bytes memory) {
+    ) public returns (bytes memory) {
+        if (amount == 0) revert InsufficientETHAmount();
         // Approve the target contract to spend the tokens
-        if(!IERC20(token).approve(to, 0)) revert ApprovalFailed();
+        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);
 
         // Reset approval
-        if(!IERC20(token).approve(to, 0)) revert ApprovalFailed();
+        if(!resetApproval(token, to)) revert ApprovalFailed();
 
-        // Transfer any remaining tokens back to the custody contract
+        // Transfer any remaining tokens back to the custody/connector contract
         uint256 remainingBalance = IERC20(token).balanceOf(address(this));
         if (remainingBalance > 0) {
-            IERC20(token).safeTransfer(address(custody), remainingBalance);
+             address destination = address(custody);
+            if (token == zetaToken) {
+                destination = address(zetaConnector);
+            }
+            IERC20(token).safeTransfer(address(destination), remainingBalance);
         }
 
         emit ExecutedWithERC20(token, to, amount, data);
 
         return result;
     }
 
-    // Transfer specified token amount to ERC20Custody and emits event
-    function sendERC20(bytes calldata recipient, address token, uint256 amount) external {
-        IERC20(token).safeTransferFrom(msg.sender, address(custody), amount);
+    // Deposit ETH to tss
+    function deposit(address receiver) external payable {
+        if (msg.value == 0) revert InsufficientETHAmount();
+        (bool deposited, ) = tssAddress.call{value: msg.value}("");
 
-        emit SendERC20(recipient, token, amount);
+        if (deposited == false) revert DepositFailed();
+
+        emit Deposit(msg.sender, receiver, msg.value, address(0), "");
     }
 
-    // Transfer specified ETH amount to TSS address and emits event
-    function send(bytes calldata recipient, uint256 amount) external payable {
+    // Deposit ERC20 tokens to custody
+    function deposit(address receiver, uint256 amount, address asset) external {
+        if (amount == 0) revert InsufficientERC20Amount();
+
+        address destination = address(custody);
+        if (asset == zetaToken) {
+            destination = address(zetaConnector);
+        }
+        IERC20(asset).safeTransferFrom(msg.sender, address(destination), amount);
+
+        emit Deposit(msg.sender, receiver, amount, asset, "");
+    }
+
+    // Deposit ETH to tss and call an omnichain smart contract
+    function depositAndCall(address receiver, bytes calldata payload) external payable {
         if (msg.value == 0) revert InsufficientETHAmount();
+        (bool deposited, ) = tssAddress.call{value: msg.value}("");
 
-        (bool sent, ) = tssAddress.call{value: msg.value}("");
+        if (deposited == false) revert DepositFailed();
+
+        emit Deposit(msg.sender, receiver, msg.value, address(0), payload);
+    }
+
+    // Deposit ERC20 tokens to custody and call an omnichain smart contract
+    function depositAndCall(address receiver, uint256 amount, address asset, bytes calldata payload) external {
+        if (amount == 0) revert InsufficientERC20Amount();
+
+        address destination = address(custody);
+        if (asset == zetaToken) {
+            destination = address(zetaConnector);
+        }
+        IERC20(asset).safeTransferFrom(msg.sender, address(destination), amount);
 
-        if (sent == false) revert SendFailed();
+        emit Deposit(msg.sender, receiver, amount, asset, payload);
+    }
 
-        emit Send(recipient, msg.value);
+    // Call an omnichain smart contract without asset transfer
+    function call(address receiver, bytes calldata payload) external {
+        emit Call(msg.sender, receiver, payload);
     }
 
     function setCustody(address _custody) external {
+        if (custody != address(0)) revert CustodyInitialized();
         custody = _custody;
     }
+
+     function setConnector(address _zetaConnector) external {
+        if (zetaConnector != address(0)) revert CustodyInitialized();
+        zetaConnector = _zetaConnector;
+    }
+
+    function resetApproval(address token, address to) private returns (bool) {
+        return IERC20(token).approve(to, 0);
+    }
 }

contracts/prototypes/evm/interfaces.sol → contracts/prototypes/evm/IGatewayEVM.sol

@@ -18,20 +18,13 @@ interface IGatewayEVMErrors {
     error CustodyInitialized();
 }
 
-interface IReceiverEVMEvents {
-    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);
-}
-
 interface IGatewayEVM {
     function executeWithERC20(
         address token,
         address to,
         uint256 amount,
         bytes calldata data
-    ) external returns (bytes memory);
+    ) external;
 
     function execute(address destination, bytes calldata data) external payable returns (bytes memory);
 }

contracts/prototypes/evm/IReceiverEVM.sol

@@ -0,0 +1,9 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+interface IReceiverEVMEvents {
+    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);
+}

contracts/prototypes/evm/ZetaConnectorNew.sol

@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.7;
+
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "./IGatewayEVM.sol";
+import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import "@openzeppelin/contracts/security/ReentrancyGuard.sol";
+
+contract ZetaConnectorNew is ReentrancyGuard{
+    using SafeERC20 for IERC20;
+    error ZeroAddress();
+
+    IGatewayEVM public immutable gateway;
+    IERC20 public immutable zetaToken;
+
+    event Withdraw(address indexed to, uint256 amount);
+    event WithdrawAndCall(address indexed to, uint256 amount, bytes data);
+
+    constructor(address _gateway, address _zetaToken) {
+        if (_gateway == address(0) || _zetaToken == address(0)) {
+            revert ZeroAddress();
+        }
+        gateway = IGatewayEVM(_gateway);
+        zetaToken = IERC20(_zetaToken);
+    }
+
+    // Withdraw is called by TSS address, it directly transfers zetaToken to the destination address without contract call
+    // TODO: Finalize access control
+    // https://github.com/zeta-chain/protocol-contracts/issues/204
+    function withdraw(address to, uint256 amount) external nonReentrant {
+        // TODO: mint?
+        zetaToken.safeTransfer(to, amount);
+
+        emit Withdraw(to, amount);
+    }
+
+    // WithdrawAndCall is called by TSS address, it transfers zetaToken and call a contract
+    // For this, it passes through the Gateway contract, it transfers zetaToken to the Gateway contract and then calls the contract
+    // TODO: Finalize access control
+    // https://github.com/zeta-chain/protocol-contracts/issues/204
+    function withdrawAndCall(address to, uint256 amount, bytes calldata data) public nonReentrant {
+        // TODO: mint?
+        // Transfer zetaToken to the Gateway contract
+        zetaToken.safeTransfer(address(gateway), amount);
+
+        // Forward the call to the Gateway contract
+        gateway.executeWithERC20(address(zetaToken), to, amount, data);
+
+        emit WithdrawAndCall(to, amount, data);
+    }
+}