feat: add zeta handling to v2 contracts

contracts/prototypes/evm/GatewayEVM.sol

@@ -15,16 +15,21 @@
 
     address public custody;
     address public tssAddress;
+    address public zetaConnector;
+    address public zeta;
 
     constructor() {}
 
-    function initialize(address _tssAddress) public initializer {
+    function initialize(address _tssAddress, address _zeta) public initializer {
         __Ownable_init();
         __UUPSUpgradeable_init();
 
-        if (_tssAddress == address(0)) revert ZeroAddress();
+        if (_tssAddress == address(0) || _zeta == address(0)) {
+            revert ZeroAddress();
+        }
 
         tssAddress = _tssAddress;
+        zeta = _zeta;
     }
 
     function _authorizeUpgrade(address newImplementation) internal override onlyOwner() {}
@@ -69,10 +74,14 @@
         // 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);
+            address destination = address(custody);
+            if (token == zeta) {
+                destination = address(zetaConnector);
+            }
+            IERC20(token).safeTransfer(address(destination), remainingBalance);
         }
 
         emit ExecutedWithERC20(token, to, amount, data);
@@ -90,10 +99,15 @@
         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);
+
+        address destination = address(custody);
+        if (asset == zeta) {
+            destination = address(zetaConnector);
+        }
+        IERC20(asset).safeTransferFrom(msg.sender, address(destination), amount);
 
         emit Deposit(msg.sender, receiver, amount, asset, "");
     }
@@ -108,10 +122,15 @@
         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);
+
+        address destination = address(custody);
+        if (asset == zeta) {
+            destination = address(zetaConnector);
+        }
+        IERC20(asset).safeTransferFrom(msg.sender, address(destination), amount);
 
         emit Deposit(msg.sender, receiver, amount, asset, payload);
     }
@@ -123,9 +142,18 @@
 
     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);
     }

contracts/prototypes/evm/GatewayEVMUpgradeTest.sol

@@ -6,40 +6,31 @@
 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";
 
 // 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
-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 zetaAsset;
 
     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 _zetaAsset) public initializer {
         __Ownable_init();
         __UUPSUpgradeable_init();
 
         if (_tssAddress == address(0)) revert ZeroAddress();
-        
+
         tssAddress = _tssAddress;
+        zetaAsset = _zetaAsset;
     }
 
     function _authorizeUpgrade(address newImplementation) internal override onlyOwner() {}
@@ -72,47 +63,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 == zetaAsset) {
+                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 == zetaAsset) {
+            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 == zetaAsset) {
+            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/ZetaConnectorNew.sol

@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.7;
+
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "./interfaces.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 zeta;
+
+    event Withdraw(address indexed to, uint256 amount);
+    event WithdrawAndCall(address indexed to, uint256 amount, bytes data);
+
+    constructor(address _gateway, address _zeta) {
+        if (_gateway == address(0) || _zeta == address(0)) {
+            revert ZeroAddress();
+        }
+        gateway = IGatewayEVM(_gateway);
+        zeta = IERC20(_zeta);
+    }
+
+    // Withdraw is called by TSS address, it directly transfers zeta 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?
+        zeta.safeTransfer(to, amount);
+
+        emit Withdraw(to, amount);
+    }
+
+    // WithdrawAndCall is called by TSS address, it transfers zeta and call a contract
+    // For this, it passes through the Gateway contract, it transfers zeta 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 zeta to the Gateway contract
+        zeta.safeTransfer(address(gateway), amount);
+
+        // Forward the call to the Gateway contract
+        gateway.executeWithERC20(address(zeta), to, amount, data);
+
+        emit WithdrawAndCall(to, amount, data);
+    }
+}

contracts/prototypes/test/GatewayEVM.t.sol

@@ -7,6 +7,7 @@ import "forge-std/Vm.sol";
 import "contracts/prototypes/evm/GatewayEVM.sol";
 import "contracts/prototypes/evm/ReceiverEVM.sol";
 import "contracts/prototypes/evm/ERC20CustodyNew.sol";
+import "contracts/prototypes/evm/ZetaConnectorNew.sol";
 import "contracts/prototypes/evm/TestERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
@@ -18,7 +19,9 @@ contract GatewayEVMTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiver
     GatewayEVM gateway;
     ReceiverEVM receiver;
     ERC20CustodyNew custody;
+    ZetaConnectorNew zetaConnector;
     TestERC20 token;
+    TestERC20 zeta;
     address owner;
     address destination;
     address tssAddress;
@@ -29,11 +32,15 @@ contract GatewayEVMTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiver
         tssAddress = address(0x5678);
 
         token = new TestERC20("test", "TTK");
+        zeta = new TestERC20("zeta", "ZETA");
+
         gateway = new GatewayEVM();
         custody = new ERC20CustodyNew(address(gateway));
+        zetaConnector = new ZetaConnectorNew(address(gateway), address(zeta));
 
-        gateway.initialize(tssAddress);
+        gateway.initialize(tssAddress, address(zeta));
         gateway.setCustody(address(custody));
+        gateway.setConnector(address(zetaConnector));
 
         // Mint initial supply to the owner
         token.mint(owner, 1000000);

contracts/prototypes/test/GatewayEVMZEVM.t.sol

@@ -7,6 +7,7 @@ import "forge-std/Vm.sol";
 import "contracts/prototypes/evm/GatewayEVM.sol";
 import "contracts/prototypes/evm/ReceiverEVM.sol";
 import "contracts/prototypes/evm/ERC20CustodyNew.sol";
+import "contracts/prototypes/evm/ZetaConnectorNew.sol";
 import "contracts/prototypes/evm/TestERC20.sol";
 import "contracts/prototypes/evm/ReceiverEVM.sol";
 
@@ -25,6 +26,8 @@ contract GatewayEVMZEVMTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IGate
     using SafeERC20 for IERC20;
 
     GatewayEVM gatewayEVM;
+    ZetaConnectorNew zetaConnector;
+    TestERC20 zeta;
     ERC20CustodyNew custody;
     TestERC20 token;
     ReceiverEVM receiverEVM;
@@ -47,11 +50,17 @@ contract GatewayEVMZEVMTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IGate
         ownerZEVM = address(0x4321);
 
         token = new TestERC20("test", "TTK");
+        zeta = new TestERC20("zeta", "ZETA");
+
         gatewayEVM = new GatewayEVM();
         custody = new ERC20CustodyNew(address(gatewayEVM));
+        zetaConnector = new ZetaConnectorNew(address(gatewayEVM), address(zeta));
+
+        gatewayEVM.initialize(tssAddress, address(zeta));
+        custody = new ERC20CustodyNew(address(gatewayEVM));
 
-        gatewayEVM.initialize(tssAddress);
         gatewayEVM.setCustody(address(custody));
+        gatewayEVM.setConnector(address(zetaConnector));
 
         // Mint initial supply to the ownerEVM
         token.mint(ownerEVM, 1000000);