chore: change Gateway contract to be upgradable

contracts/prototypes/ERC20CustodyNew.sol

@@ -2,19 +2,19 @@
 pragma solidity 0.8.7;
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
-import "./Gateway.sol";
+import "./interfaces.sol";
 
 // As the current version, ERC20CustodyNew hold the ERC20s deposited on ZetaChain
 // This version include a functionality allowing to call a contract
 // ERC20Custody doesn't call smart contract directly, it passes through the Gateway contract
 contract ERC20CustodyNew {
-    Gateway public gateway;
+    IGateway public gateway;
 
     event Withdraw(address indexed token, address indexed to, uint256 amount);
     event WithdrawAndCall(address indexed token, address indexed to, uint256 amount, bytes data);
 
     constructor(address _gateway) {
-        gateway = Gateway(_gateway);
+        gateway = IGateway(_gateway);
     }
 
     // Withdraw is called by TSS address, it directly transfers the tokens to the destination address without contract call

contracts/prototypes/Gateway.sol

@@ -2,18 +2,33 @@
 pragma solidity 0.8.7;
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
-import "./ERC20CustodyNew.sol";
+import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+
 
 // 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 Gateway {
+contract Gateway is Initializable, OwnableUpgradeable, UUPSUpgradeable {
     error ExecutionFailed();
 
-    ERC20CustodyNew public custody;
+    address public custody;
 
     event Executed(address indexed destination, uint256 value, bytes data);
     event ExecutedWithERC20(address indexed token, address indexed to, uint256 amount, bytes data);
 
+    /// @custom:oz-upgrades-unsafe-allow constructor
+    constructor() {
+        _disableInitializers();
+    }
+
+    function initialize() public initializer {
+        __Ownable_init();
+        __UUPSUpgradeable_init();
+    }
+
+    function _authorizeUpgrade(address newImplementation) internal override onlyOwner() {}
+
     function _execute(address destination, bytes calldata data) internal returns (bytes memory) {
         (bool success, bytes memory result) = destination.call{value: msg.value}(data);
 
@@ -66,6 +81,6 @@ contract Gateway {
     }
 
     function setCustody(address _custody) external {
-        custody = ERC20CustodyNew(_custody);
+        custody = _custody;
     }
 }

contracts/prototypes/GatewayUpgradeTest.sol

@@ -0,0 +1,87 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.7;
+
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+
+
+// NOTE: Purpose of this contract is to test upgrade process, the only difference should be event names
+// 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 GatewayUpgradeTest is Initializable, OwnableUpgradeable, UUPSUpgradeable {
+    error ExecutionFailed();
+
+    address public custody;
+
+    event ExecutedV2(address indexed destination, uint256 value, bytes data);
+    event ExecutedWithERC20V2(address indexed token, address indexed to, uint256 amount, bytes data);
+
+    /// @custom:oz-upgrades-unsafe-allow constructor
+    constructor() {
+        _disableInitializers();
+    }
+
+    function initialize() public initializer {
+        __Ownable_init();
+        __UUPSUpgradeable_init();
+    }
+
+    function _authorizeUpgrade(address newImplementation) internal override onlyOwner() {}
+
+    function _execute(address destination, bytes calldata data) internal returns (bytes memory) {
+        (bool success, bytes memory result) = destination.call{value: msg.value}(data);
+
+        if (!success) {
+            revert ExecutionFailed();
+        }
+
+        return result;
+    }
+
+    // Called by the TSS
+    // Execution without ERC20 tokens, it is payable and can be used in the case of WithdrawAndCall for Gas ZRC20
+    // It can be also used for contract call without asset movement
+    function execute(address destination, bytes calldata data) external payable returns (bytes memory) {
+        bytes memory result = _execute(destination, data);
+
+        emit ExecutedV2(destination, msg.value, data);
+
+        return result;
+    }
+
+    // Called by the ERC20Custody contract
+    // It call a function using ERC20 transfer
+    // Since the goal is to allow calling contract not designed for ZetaChain specifically, it uses ERC20 allowance system
+    // It provides allowance to destination contract and call destination contract. In the end, it remove remaining allowance and transfer remaining tokens back to the custody contract for security purposes
+    function executeWithERC20(
+        address token,
+        address to,
+        uint256 amount,
+        bytes calldata data
+    ) external returns (bytes memory) {
+        // Approve the target contract to spend the tokens
+        IERC20(token).approve(to, amount);
+
+        // Execute the call on the target contract
+        bytes memory result = _execute(to, data);
+
+        // Reset approval
+        IERC20(token).approve(to, 0);
+
+        // Transfer any remaining tokens back to the custody contract
+        uint256 remainingBalance = IERC20(token).balanceOf(address(this));
+        if (remainingBalance > 0) {
+            IERC20(token).transfer(address(custody), remainingBalance);
+        }
+
+        emit ExecutedWithERC20V2(token, to, amount, data);
+
+        return result;
+    }
+
+    function setCustody(address _custody) external {
+        custody = _custody;
+    }
+}

contracts/prototypes/interfaces.sol

@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.7;
+
+interface IGateway {
+    function executeWithERC20(
+        address token,
+        address to,
+        uint256 amount,
+        bytes calldata data
+    ) external returns (bytes memory);
+
+    function execute(address destination, bytes calldata data) external payable returns (bytes memory);
+}

hardhat.config.ts

@@ -5,6 +5,7 @@ import "tsconfig-paths/register";
 import "hardhat-abi-exporter";
 import "uniswap-v2-deploy-plugin";
 import "./tasks/addresses";
+import "@openzeppelin/hardhat-upgrades";
 
 import { getHardhatConfigNetworks } from "@zetachain/networks";
 import * as dotenv from "dotenv";

package.json

@@ -11,6 +11,8 @@
     "@nomiclabs/hardhat-ethers": "^2.0.5",
     "@nomiclabs/hardhat-waffle": "^2.0.3",
     "@openzeppelin/contracts": "^4.8.3",
+    "@openzeppelin/contracts-upgradeable": "^4.8.3",
+    "@openzeppelin/hardhat-upgrades": "1.28.0",
     "@typechain/ethers-v5": "^10.1.0",
     "@typechain/hardhat": "^6.1.2",
     "@types/chai": "^4.3.1",