erc20 upgradable

v2/contracts/evm/ERC20Custody.sol

@@ -6,20 +6,30 @@ import { IGatewayEVM } from "./interfaces/IGatewayEVM.sol";
 
 import { RevertContext } from "../../contracts/Revert.sol";
 
-import "@openzeppelin/contracts/access/AccessControl.sol";
+import "@openzeppelin/contracts-upgradeable/access/AccessControlUpgradeable.sol";
+
+import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/utils/ReentrancyGuardUpgradeable.sol";
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
-import "@openzeppelin/contracts/utils/Pausable.sol";
-import "@openzeppelin/contracts/utils/ReentrancyGuard.sol";
 
 /// @title ERC20Custody
 /// @notice Holds the ERC20 tokens deposited on ZetaChain and includes functionality to call a contract.
 /// @dev This contract does not call smart contracts directly, it passes through the Gateway contract.
-contract ERC20Custody is IERC20Custody, ReentrancyGuard, AccessControl, Pausable {
+contract ERC20Custody is
+    Initializable,
+    UUPSUpgradeable,
+    IERC20Custody,
+    ReentrancyGuardUpgradeable,
+    AccessControlUpgradeable,
+    PausableUpgradeable
+{
     using SafeERC20 for IERC20;
 
     /// @notice Gateway contract.
-    IGatewayEVM public immutable gateway;
+    IGatewayEVM public gateway;
     /// @notice Mapping of whitelisted tokens => true/false.
     mapping(address => bool) public whitelisted;
     /// @notice The address of the TSS (Threshold Signature Scheme) contract.
@@ -33,12 +43,18 @@ contract ERC20Custody is IERC20Custody, ReentrancyGuard, AccessControl, Pausable
     /// @notice New role identifier for whitelister role.
     bytes32 public constant WHITELISTER_ROLE = keccak256("WHITELISTER_ROLE");
 
-    /// @notice Constructor for ERC20Custody.
+    /// @notice Initializer for ERC20Custody.
     /// @dev Set admin as default admin and pauser, and tssAddress as tss role.
-    constructor(address gateway_, address tssAddress_, address admin_) {
+    function initialize(address gateway_, address tssAddress_, address admin_) public initializer {
         if (gateway_ == address(0) || tssAddress_ == address(0) || admin_ == address(0)) {
             revert ZeroAddress();
         }
+
+        __UUPSUpgradeable_init();
+        __ReentrancyGuard_init();
+        __AccessControl_init();
+        __Pausable_init();
+
         gateway = IGatewayEVM(gateway_);
         tssAddress = tssAddress_;
         _grantRole(DEFAULT_ADMIN_ROLE, admin_);
@@ -48,6 +64,10 @@ contract ERC20Custody is IERC20Custody, ReentrancyGuard, AccessControl, Pausable
         _grantRole(WHITELISTER_ROLE, tssAddress_);
     }
 
+    /// @dev Authorizes the upgrade of the contract, sender must be owner.
+    /// @param newImplementation Address of the new implementation.
+    function _authorizeUpgrade(address newImplementation) internal override onlyRole(DEFAULT_ADMIN_ROLE) { }
+
     /// @notice Pause contract.
     function pause() external onlyRole(PAUSER_ROLE) {
         _pause();

v2/scripts/deploy/deterministic/DeployERC20Custody.s.sol

@@ -2,27 +2,55 @@
 pragma solidity 0.8.26;
 
 import "forge-std/Script.sol";
-import "src/evm/ERC20Custody.sol";
+import "contracts/evm/ERC20Custody.sol";
+import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
 
 contract DeployERC20Custody is Script {
     function run() external {
         address payable tss = payable(vm.envAddress("TSS_ADDRESS"));
         address admin = vm.envAddress("ERC20_CUSTODY_ADMIN_ADDRESS_EVM");
         address gateway = vm.envAddress("GATEWAY_PROXY_EVM");
 
-        bytes32 salt = keccak256("ERC20Custody");
+        address expectedImplAddress;
+        address expectedProxyAddress;
+
+        bytes32 implSalt = keccak256("ERC20Custody");
+        bytes32 proxySalt = keccak256("ERC20CustodyProxy");
 
         vm.startBroadcast();
 
-        ERC20Custody custody = new ERC20Custody{salt: salt}(gateway, tss, admin);
-        require(address(custody) != address(0), "deployment failed");
+        expectedImplAddress = vm.computeCreate2Address(
+            implSalt,
+            hashInitCode(type(ERC20Custody).creationCode)
+        );
+
+        ERC20Custody erc20CustodyImpl = new ERC20Custody{salt: implSalt}();
+        require(address(erc20CustodyImpl) != address(0), "erc20CustodyImpl deployment failed");
+
+        require(expectedImplAddress == address(erc20CustodyImpl), "impl address doesn't match expected address");
 
-        address expectedAddr = vm.computeCreate2Address(
-            salt,
-            hashInitCode(type(ERC20Custody).creationCode, abi.encode(gateway, tss, admin))
+        expectedProxyAddress = vm.computeCreate2Address(
+            proxySalt,
+            hashInitCode(
+                type(ERC1967Proxy).creationCode,
+                abi.encode(
+                    address(erc20CustodyImpl),
+                    abi.encodeWithSelector(ERC20Custody.initialize.selector, gateway, tss, admin)
+                )
+            )
         );
 
-        require(expectedAddr == address(custody), "erc20 custody address doesn't match expected address");
+        ERC1967Proxy erc20CustodyProxy = new ERC1967Proxy{salt: proxySalt}(
+            address(erc20CustodyImpl),
+            abi.encodeWithSelector(ERC20Custody.initialize.selector, gateway, tss, admin)
+        );
+        require(address(erc20CustodyProxy) != address(0), "erc20CustodyProxy deployment failed");
+
+        require(expectedProxyAddress == address(erc20CustodyProxy), "proxy address doesn't match expected address");
+
+        ERC20Custody erc20Custody = ERC20Custody(address(erc20CustodyProxy));
+        require(erc20Custody.tssAddress() == tss, "tss not set");
+        require(address(erc20Custody.gateway()) == gateway, "gateway not set");
 
         vm.stopBroadcast();
     }

v2/scripts/deploy/deterministic/DeployGatewayEVM.s.sol

@@ -2,7 +2,7 @@
 pragma solidity 0.8.26;
 
 import "forge-std/Script.sol";
-import "src/evm/GatewayEVM.sol";
+import "contracts/evm/GatewayEVM.sol";
 import "test/utils/TestERC20.sol";
 import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
 
@@ -47,7 +47,6 @@ contract DeployGatewayEVM is Script {
         );
         require(address(gatewayProxy) != address(0), "gatewayProxy deployment failed");
 
-
         require(expectedProxyAddress == address(gatewayProxy), "proxy address doesn't match expected address");
 
         GatewayEVM gateway = GatewayEVM(address(gatewayProxy));

v2/scripts/deploy/deterministic/DeployZetaConnectorNonNative.s.sol

@@ -2,7 +2,7 @@
 pragma solidity 0.8.26;
 
 import "forge-std/Script.sol";
-import "src/evm/ZetaConnectorNonNative.sol";
+import "contracts/evm/ZetaConnectorNonNative.sol";
 
 contract DeployZetaConnectorNonNative is Script {
     function run() external {

v2/scripts/deploy/deterministic/UpgradeGatewayEVM.s.sol

@@ -2,7 +2,7 @@
 pragma solidity 0.8.26;
 
 import "forge-std/Script.sol";
-import "src/evm/GatewayEVM.sol";
+import "contracts/evm/GatewayEVM.sol";
 import "test/utils/GatewayEVMUpgradeTest.sol";
 import "test/utils/TestERC20.sol";
 import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";

v2/test/ERC20Custody.t.sol

@@ -23,7 +23,6 @@ import "./utils/IReceiverEVM.sol";
 contract ERC20CustodyTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiverEVMEvents, IERC20CustodyEvents {
     using SafeERC20 for IERC20;
 
-    address proxy;
     GatewayEVM gateway;
     ReceiverEVM receiver;
     ERC20Custody custody;
@@ -55,11 +54,14 @@ contract ERC20CustodyTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiv
         token = new TestERC20("test", "TTK");
         zeta = new TestERC20("zeta", "ZETA");
 
-        proxy = Upgrades.deployUUPSProxy(
+        address proxy = Upgrades.deployUUPSProxy(
             "GatewayEVM.sol", abi.encodeCall(GatewayEVM.initialize, (tssAddress, address(zeta), owner))
         );
         gateway = GatewayEVM(proxy);
-        custody = new ERC20Custody(address(gateway), tssAddress, owner);
+        proxy = Upgrades.deployUUPSProxy(
+            "ERC20Custody.sol", abi.encodeCall(ERC20Custody.initialize, (address(gateway), tssAddress, owner))
+        );
+        custody = ERC20Custody(proxy);
         zetaConnector = new ZetaConnectorNonNative(address(gateway), address(zeta), tssAddress, owner);
         receiver = new ReceiverEVM();
 
@@ -177,19 +179,6 @@ contract ERC20CustodyTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiv
         assertEq(false, whitelisted);
     }
 
-    function testNewCustodyFailsIfAddressesAreZero() public {
-        vm.expectRevert(ZeroAddress.selector);
-        ERC20Custody newCustody = new ERC20Custody(address(0), tssAddress, owner);
-
-        vm.expectRevert(ZeroAddress.selector);
-        newCustody = new ERC20Custody(address(gateway), address(0), owner);
-
-        vm.expectRevert(ZeroAddress.selector);
-        newCustody = new ERC20Custody(address(gateway), tssAddress, address(0));
-
-        newCustody = new ERC20Custody(address(gateway), tssAddress, owner);
-    }
-
     function testForwardCallToReceiveERC20ThroughCustody() public {
         uint256 amount = 100_000;
         bytes memory data =

v2/test/GatewayEVM.t.sol

@@ -58,7 +58,10 @@ contract GatewayEVMTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiver
             "GatewayEVM.sol", abi.encodeCall(GatewayEVM.initialize, (tssAddress, address(zeta), owner))
         );
         gateway = GatewayEVM(proxy);
-        custody = new ERC20Custody(address(gateway), tssAddress, owner);
+        proxy = Upgrades.deployUUPSProxy(
+            "ERC20Custody.sol", abi.encodeCall(ERC20Custody.initialize, (address(gateway), tssAddress, owner))
+        );
+        custody = ERC20Custody(proxy);
         zetaConnector = new ZetaConnectorNonNative(address(gateway), address(zeta), tssAddress, owner);
         vm.prank(tssAddress);
         zeta.updateTssAndConnectorAddresses(tssAddress, address(zetaConnector));
@@ -386,7 +389,11 @@ contract GatewayEVMInboundTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IR
             "GatewayEVM.sol", abi.encodeCall(GatewayEVM.initialize, (tssAddress, address(zeta), owner))
         );
         gateway = GatewayEVM(proxy);
-        custody = new ERC20Custody(address(gateway), tssAddress, owner);
+
+        proxy = Upgrades.deployUUPSProxy(
+            "ERC20Custody.sol", abi.encodeCall(ERC20Custody.initialize, (address(gateway), tssAddress, owner))
+        );
+        custody = ERC20Custody(proxy);
         zetaConnector = new ZetaConnectorNonNative(address(gateway), address(zeta), tssAddress, owner);
         vm.prank(tssAddress);
         zeta.updateTssAndConnectorAddresses(tssAddress, address(zetaConnector));

v2/test/GatewayEVMUpgrade.t.sol

@@ -52,7 +52,10 @@ contract GatewayEVMUUPSUpgradeTest is Test, IGatewayEVMErrors, IGatewayEVMEvents
         );
         gateway = GatewayEVM(proxy);
 
-        custody = new ERC20Custody(address(gateway), tssAddress, owner);
+        address erc20CustodyProxy = Upgrades.deployUUPSProxy(
+            "ERC20Custody.sol", abi.encodeCall(ERC20Custody.initialize, (address(gateway), tssAddress, owner))
+        );
+        custody = ERC20Custody(erc20CustodyProxy);
         zetaConnector = new ZetaConnectorNonNative(address(gateway), address(zeta), tssAddress, owner);
         receiver = new ReceiverEVM();
 

v2/test/GatewayEVMZEVM.t.sol

@@ -41,7 +41,6 @@ contract GatewayEVMZEVMTest is
     // evm
     using SafeERC20 for IERC20;
 
-    address proxyEVM;
     GatewayEVM gatewayEVM;
     ERC20Custody custody;
     ZetaConnectorNonNative zetaConnector;
@@ -72,11 +71,14 @@ contract GatewayEVMZEVMTest is
         token = new TestERC20("test", "TTK");
         zeta = new TestERC20("zeta", "ZETA");
 
-        proxyEVM = Upgrades.deployUUPSProxy(
+        address proxy = Upgrades.deployUUPSProxy(
             "GatewayEVM.sol", abi.encodeCall(GatewayEVM.initialize, (tssAddress, address(zeta), ownerEVM))
         );
-        gatewayEVM = GatewayEVM(proxyEVM);
-        custody = new ERC20Custody(address(gatewayEVM), tssAddress, ownerEVM);
+        gatewayEVM = GatewayEVM(proxy);
+        proxy = Upgrades.deployUUPSProxy(
+            "ERC20Custody.sol", abi.encodeCall(ERC20Custody.initialize, (address(gatewayEVM), tssAddress, ownerEVM))
+        );
+        custody = ERC20Custody(proxy);
         zetaConnector = new ZetaConnectorNonNative(address(gatewayEVM), address(zeta), tssAddress, ownerEVM);
 
         vm.deal(tssAddress, 1 ether);

v2/test/ZetaConnectorNative.t.sol

@@ -30,7 +30,6 @@ contract ZetaConnectorNativeTest is
 {
     using SafeERC20 for IERC20;
 
-    address proxy;
     GatewayEVM gateway;
     ReceiverEVM receiver;
     ERC20Custody custody;
@@ -56,11 +55,14 @@ contract ZetaConnectorNativeTest is
 
         zetaToken = new TestERC20("zeta", "ZETA");
 
-        proxy = Upgrades.deployUUPSProxy(
+        address proxy = Upgrades.deployUUPSProxy(
             "GatewayEVM.sol", abi.encodeCall(GatewayEVM.initialize, (tssAddress, address(zetaToken), owner))
         );
         gateway = GatewayEVM(proxy);
-        custody = new ERC20Custody(address(gateway), tssAddress, owner);
+        proxy = Upgrades.deployUUPSProxy(
+            "ERC20Custody.sol", abi.encodeCall(ERC20Custody.initialize, (address(gateway), tssAddress, owner))
+        );
+        custody = ERC20Custody(proxy);
         zetaConnector = new ZetaConnectorNative(address(gateway), address(zetaToken), tssAddress, owner);
 
         receiver = new ReceiverEVM();

v2/test/ZetaConnectorNonNative.t.sol

@@ -30,7 +30,6 @@ contract ZetaConnectorNonNativeTest is
 {
     using SafeERC20 for IERC20;
 
-    address proxy;
     GatewayEVM gateway;
     ReceiverEVM receiver;
     ERC20Custody custody;
@@ -56,12 +55,15 @@ contract ZetaConnectorNonNativeTest is
 
         zetaToken = new ZetaNonEth(tssAddress, tssAddress);
 
-        proxy = Upgrades.deployUUPSProxy(
+        address proxy = Upgrades.deployUUPSProxy(
             "GatewayEVM.sol", abi.encodeCall(GatewayEVM.initialize, (tssAddress, address(zetaToken), owner))
         );
         gateway = GatewayEVM(proxy);
 
-        custody = new ERC20Custody(address(gateway), tssAddress, owner);
+        proxy = Upgrades.deployUUPSProxy(
+            "ERC20Custody.sol", abi.encodeCall(ERC20Custody.initialize, (address(gateway), tssAddress, owner))
+        );
+        custody = ERC20Custody(proxy);
         zetaConnector = new ZetaConnectorNonNative(address(gateway), address(zetaToken), tssAddress, owner);
 
         vm.prank(tssAddress);