From ec2fd2afc191922ecd1aea1903a837977ec7967e Mon Sep 17 00:00:00 2001 From: skosito Date: Fri, 11 Oct 2024 20:49:42 +0200 Subject: [PATCH] add upgrade unit test for every upgradable contract --- v2/test/ERC20Custody.t.sol | 33 ++ v2/test/GatewayEVM.t.sol | 32 +- v2/test/GatewayEVMUpgrade.t.sol | 105 ---- v2/test/GatewayZEVM.t.sol | 41 ++ v2/test/ZetaConnectorNative.t.sol | 23 + v2/test/ZetaConnectorNonNative.t.sol | 24 + .../upgrades/ERC20CustodyUpgradeTest.sol | 224 ++++++++ .../{ => upgrades}/GatewayEVMUpgradeTest.sol | 8 +- .../utils/upgrades/GatewayZEVMUpgradeTest.sol | 532 ++++++++++++++++++ .../ZetaConnectorNativeUpgradeTest.sol | 112 ++++ .../ZetaConnectorNonNativeUpgradeTest.sol | 134 +++++ 11 files changed, 1157 insertions(+), 111 deletions(-) delete mode 100644 v2/test/GatewayEVMUpgrade.t.sol create mode 100644 v2/test/utils/upgrades/ERC20CustodyUpgradeTest.sol rename v2/test/utils/{ => upgrades}/GatewayEVMUpgradeTest.sol (98%) create mode 100644 v2/test/utils/upgrades/GatewayZEVMUpgradeTest.sol create mode 100644 v2/test/utils/upgrades/ZetaConnectorNativeUpgradeTest.sol create mode 100644 v2/test/utils/upgrades/ZetaConnectorNonNativeUpgradeTest.sol diff --git a/v2/test/ERC20Custody.t.sol b/v2/test/ERC20Custody.t.sol index 312b2836..903c4e41 100644 --- a/v2/test/ERC20Custody.t.sol +++ b/v2/test/ERC20Custody.t.sol @@ -7,6 +7,7 @@ import "forge-std/Vm.sol"; import "./utils/ReceiverEVM.sol"; import "./utils/TestERC20.sol"; +import "./utils/upgrades/ERC20CustodyUpgradeTest.sol"; import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; import "@openzeppelin/contracts/token/ERC20/IERC20.sol"; @@ -39,6 +40,8 @@ contract ERC20CustodyTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiv error AccessControlUnauthorizedAccount(address account, bytes32 neededRole); error LegacyMethodsNotSupported(); + event WithdrawnV2(address indexed to, address indexed token, uint256 amount); + bytes32 public constant TSS_ROLE = keccak256("TSS_ROLE"); bytes32 public constant WITHDRAWER_ROLE = keccak256("WITHDRAWER_ROLE"); bytes32 public constant ASSET_HANDLER_ROLE = keccak256("ASSET_HANDLER_ROLE"); @@ -552,4 +555,34 @@ contract ERC20CustodyTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiv vm.expectRevert(NotWhitelisted.selector); custody.deposit(abi.encodePacked(destination), token, 1000, message); } + + function testUpgradeAndWithdraw() public { + // upgrade + Upgrades.upgradeProxy(address(custody), "ERC20CustodyUpgradeTest.sol", "", owner); + ERC20CustodyUpgradeTest custodyV2 = ERC20CustodyUpgradeTest(address(custody)); + // withdraw + uint256 amount = 100_000; + uint256 balanceBefore = token.balanceOf(destination); + assertEq(balanceBefore, 0); + uint256 balanceBeforeCustody = token.balanceOf(address(custodyV2)); + + bytes memory transferData = abi.encodeWithSignature("transfer(address,uint256)", address(destination), amount); + vm.expectCall(address(token), 0, transferData); + vm.expectEmit(true, true, true, true, address(custodyV2)); + emit WithdrawnV2(destination, address(token), amount); + vm.prank(tssAddress); + custodyV2.withdraw(destination, address(token), amount); + + // Verify that the tokens were transferred to the destination address + uint256 balanceAfter = token.balanceOf(destination); + assertEq(balanceAfter, amount); + + // Verify that the tokens were substracted from custody + uint256 balanceAfterCustody = token.balanceOf(address(custodyV2)); + assertEq(balanceAfterCustody, balanceBeforeCustody - amount); + + // Verify that gateway doesn't hold any tokens + uint256 balanceGateway = token.balanceOf(address(gateway)); + assertEq(balanceGateway, 0); + } } diff --git a/v2/test/GatewayEVM.t.sol b/v2/test/GatewayEVM.t.sol index ec1b2d7b..b9c928e6 100644 --- a/v2/test/GatewayEVM.t.sol +++ b/v2/test/GatewayEVM.t.sol @@ -7,6 +7,7 @@ import "forge-std/Vm.sol"; import "./utils/ReceiverEVM.sol"; import "./utils/TestERC20.sol"; +import "./utils/upgrades/GatewayEVMUpgradeTest.sol"; import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; import "@openzeppelin/contracts/token/ERC20/IERC20.sol"; @@ -24,7 +25,6 @@ import "./utils/Zeta.non-eth.sol"; contract GatewayEVMTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiverEVMEvents, IERC20CustodyEvents { using SafeERC20 for IERC20; - address proxy; GatewayEVM gateway; ReceiverEVM receiver; ERC20Custody custody; @@ -40,6 +40,8 @@ contract GatewayEVMTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiver error EnforcedPause(); error AccessControlUnauthorizedAccount(address account, bytes32 neededRole); + event ExecutedV2(address indexed destination, uint256 value, bytes data); + bytes32 public constant TSS_ROLE = keccak256("TSS_ROLE"); bytes32 public constant WITHDRAWER_ROLE = keccak256("WITHDRAWER_ROLE"); bytes32 public constant ASSET_HANDLER_ROLE = keccak256("ASSET_HANDLER_ROLE"); @@ -54,7 +56,7 @@ contract GatewayEVMTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiver token = new TestERC20("test", "TTK"); zeta = new ZetaNonEth(tssAddress, tssAddress); - proxy = Upgrades.deployUUPSProxy( + address proxy = Upgrades.deployUUPSProxy( "GatewayEVM.sol", abi.encodeCall(GatewayEVM.initialize, (tssAddress, address(zeta), owner)) ); gateway = GatewayEVM(proxy); @@ -364,6 +366,32 @@ contract GatewayEVMTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiver vm.expectRevert(ZeroAddress.selector); gateway.executeRevert{ value: value }(address(0), data, revertContext); } + + function testUpgradeAndForwardCallToReceivePayable() public { + // upgrade + Upgrades.upgradeProxy(address(gateway), "GatewayEVMUpgradeTest.sol", "", owner); + GatewayEVMUpgradeTest gatewayUpgradeTest = GatewayEVMUpgradeTest(address(gateway)); + // call + address custodyBeforeUpgrade = gateway.custody(); + address tssBeforeUpgrade = gateway.tssAddress(); + + string memory str = "Hello, Foundry!"; + uint256 num = 42; + bool flag = true; + uint256 value = 1 ether; + + bytes memory data = abi.encodeWithSignature("receivePayable(string,uint256,bool)", str, num, flag); + vm.expectCall(address(receiver), value, data); + vm.expectEmit(true, true, true, true, address(receiver)); + emit ReceivedPayable(address(gateway), value, str, num, flag); + vm.expectEmit(true, true, true, true, address(gateway)); + emit ExecutedV2(address(receiver), value, data); + vm.prank(tssAddress); + gatewayUpgradeTest.execute{ value: value }(address(receiver), data); + + assertEq(custodyBeforeUpgrade, gateway.custody()); + assertEq(tssBeforeUpgrade, gateway.tssAddress()); + } } contract GatewayEVMInboundTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiverEVMEvents { diff --git a/v2/test/GatewayEVMUpgrade.t.sol b/v2/test/GatewayEVMUpgrade.t.sol deleted file mode 100644 index cf084d23..00000000 --- a/v2/test/GatewayEVMUpgrade.t.sol +++ /dev/null @@ -1,105 +0,0 @@ -// SPDX-License-Identifier: MIT -pragma solidity 0.8.26; - -import "forge-std/Test.sol"; -import "forge-std/Vm.sol"; - -import "./utils/GatewayEVMUpgradeTest.sol"; - -import "./utils/IReceiverEVM.sol"; -import "./utils/ReceiverEVM.sol"; -import "./utils/TestERC20.sol"; -import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; - -import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; -import "@openzeppelin/contracts/token/ERC20/IERC20.sol"; -import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; -import { Upgrades } from "openzeppelin-foundry-upgrades/Upgrades.sol"; - -import "./utils/IReceiverEVM.sol"; - -import "../contracts/evm/ERC20Custody.sol"; -import "../contracts/evm/GatewayEVM.sol"; -import "../contracts/evm/ZetaConnectorNonNative.sol"; -import "../contracts/evm/interfaces/IGatewayEVM.sol"; - -contract GatewayEVMUUPSUpgradeTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiverEVMEvents { - using SafeERC20 for IERC20; - - event ExecutedV2(address indexed destination, uint256 value, bytes data); - - address proxy; - GatewayEVM gateway; - ReceiverEVM receiver; - ERC20Custody custody; - ZetaConnectorNonNative zetaConnector; - TestERC20 token; - TestERC20 zeta; - address owner; - address destination; - address tssAddress; - - function setUp() public { - owner = address(this); - destination = address(0x1234); - tssAddress = address(0x5678); - - token = new TestERC20("test", "TTK"); - zeta = new TestERC20("zeta", "ZETA"); - - proxy = Upgrades.deployUUPSProxy( - "GatewayEVM.sol", abi.encodeCall(GatewayEVM.initialize, (tssAddress, address(zeta), owner)) - ); - gateway = GatewayEVM(proxy); - - address erc20CustodyProxy = Upgrades.deployUUPSProxy( - "ERC20Custody.sol", abi.encodeCall(ERC20Custody.initialize, (address(gateway), tssAddress, owner)) - ); - custody = ERC20Custody(erc20CustodyProxy); - address connectorProxy = Upgrades.deployUUPSProxy( - "ZetaConnectorNonNative.sol", - abi.encodeCall(ZetaConnectorNonNative.initialize, (address(gateway), address(zeta), tssAddress, owner)) - ); - zetaConnector = ZetaConnectorNonNative(connectorProxy); - - receiver = new ReceiverEVM(); - - vm.deal(tssAddress, 1 ether); - - vm.startPrank(owner); - gateway.setCustody(address(custody)); - gateway.setConnector(address(zetaConnector)); - vm.stopPrank(); - - token.mint(owner, 1_000_000); - token.transfer(address(custody), 500_000); - - vm.deal(tssAddress, 1 ether); - } - - function testUpgradeAndForwardCallToReceivePayable() public { - address custodyBeforeUpgrade = gateway.custody(); - address tssBeforeUpgrade = gateway.tssAddress(); - - string memory str = "Hello, Foundry!"; - uint256 num = 42; - bool flag = true; - uint256 value = 1 ether; - - Upgrades.upgradeProxy(proxy, "GatewayEVMUpgradeTest.sol", "", owner); - - bytes memory data = abi.encodeWithSignature("receivePayable(string,uint256,bool)", str, num, flag); - GatewayEVMUpgradeTest gatewayUpgradeTest = GatewayEVMUpgradeTest(proxy); - - vm.expectCall(address(receiver), value, data); - vm.expectEmit(true, true, true, true, address(receiver)); - emit ReceivedPayable(address(gateway), value, str, num, flag); - vm.expectEmit(true, true, true, true, address(gateway)); - emit ExecutedV2(address(receiver), value, data); - vm.prank(tssAddress); - gatewayUpgradeTest.execute{ value: value }(address(receiver), data); - - assertEq(custodyBeforeUpgrade, gateway.custody()); - assertEq(tssBeforeUpgrade, gateway.tssAddress()); - } -} diff --git a/v2/test/GatewayZEVM.t.sol b/v2/test/GatewayZEVM.t.sol index 0016d004..1538b961 100644 --- a/v2/test/GatewayZEVM.t.sol +++ b/v2/test/GatewayZEVM.t.sol @@ -9,6 +9,7 @@ import "../contracts/zevm/SystemContract.sol"; import "./utils/TestUniversalContract.sol"; import "./utils/WZETA.sol"; +import "./utils/upgrades/GatewayZEVMUpgradeTest.sol"; import "../contracts/zevm/GatewayZEVM.sol"; import "../contracts/zevm/ZRC20.sol"; @@ -32,6 +33,19 @@ contract GatewayZEVMInboundTest is Test, IGatewayZEVMEvents, IGatewayZEVMErrors error ZeroAddress(); error LowBalance(); + event WithdrawnV2( + address indexed sender, + uint256 indexed chainId, + bytes receiver, + address zrc20, + uint256 value, + uint256 gasfee, + uint256 protocolFlatFee, + bytes message, + CallOptions callOptions, + RevertOptions revertOptions + ); + function setUp() public { owner = address(this); addr1 = address(0x1234); @@ -585,6 +599,33 @@ contract GatewayZEVMInboundTest is Test, IGatewayZEVMEvents, IGatewayZEVMErrors emit Called(owner, address(zrc20), abi.encodePacked(addr1), message, callOptions, revertOptions); gateway.call(abi.encodePacked(addr1), address(zrc20), message, callOptions, revertOptions); } + + function testUpgradeAndWithdrawZRC20() public { + // upgrade + Upgrades.upgradeProxy(proxy, "GatewayZEVMUpgradeTest.sol", "", owner); + GatewayZEVMUpgradeTest gatewayUpgradeTest = GatewayZEVMUpgradeTest(proxy); + // withdraw + uint256 amount = 1; + uint256 ownerBalanceBefore = zrc20.balanceOf(owner); + + vm.expectEmit(true, true, true, true, address(gatewayUpgradeTest)); + emit WithdrawnV2( + owner, + 0, + abi.encodePacked(addr1), + address(zrc20), + amount, + 0, + zrc20.PROTOCOL_FLAT_FEE(), + "", + CallOptions({ gasLimit: 0, isArbitraryCall: true }), + revertOptions + ); + gatewayUpgradeTest.withdraw(abi.encodePacked(addr1), amount, address(zrc20), revertOptions); + + uint256 ownerBalanceAfter = zrc20.balanceOf(owner); + assertEq(ownerBalanceBefore - amount, ownerBalanceAfter); + } } contract GatewayZEVMOutboundTest is Test, IGatewayZEVMEvents, IGatewayZEVMErrors { diff --git a/v2/test/ZetaConnectorNative.t.sol b/v2/test/ZetaConnectorNative.t.sol index 3bd79bba..6f754301 100644 --- a/v2/test/ZetaConnectorNative.t.sol +++ b/v2/test/ZetaConnectorNative.t.sol @@ -7,6 +7,7 @@ import "forge-std/Vm.sol"; import "./utils/ReceiverEVM.sol"; import "./utils/TestERC20.sol"; +import "./utils/upgrades/ZetaConnectorNativeUpgradeTest.sol"; import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; import "@openzeppelin/contracts/token/ERC20/IERC20.sol"; @@ -43,6 +44,8 @@ contract ZetaConnectorNativeTest is error EnforcedPause(); error AccessControlUnauthorizedAccount(address account, bytes32 neededRole); + event WithdrawnV2(address indexed to, uint256 amount); + bytes32 public constant WITHDRAWER_ROLE = keccak256("WITHDRAWER_ROLE"); bytes32 public constant PAUSER_ROLE = keccak256("PAUSER_ROLE"); bytes32 public constant TSS_ROLE = keccak256("TSS_ROLE"); @@ -350,4 +353,24 @@ contract ZetaConnectorNativeTest is vm.expectRevert(abi.encodeWithSelector(AccessControlUnauthorizedAccount.selector, owner, WITHDRAWER_ROLE)); zetaConnector.withdrawAndRevert(address(receiver), amount, data, internalSendHash, revertContext); } + + function testUpgradeAndWithdraw() public { + // upgrade + Upgrades.upgradeProxy(address(zetaConnector), "ZetaConnectorNativeUpgradeTest.sol", "", owner); + ZetaConnectorNativeUpgradeTest zetaConnectorV2 = ZetaConnectorNativeUpgradeTest(address(zetaConnector)); + // withdraw + uint256 amount = 100_000; + bytes32 internalSendHash = ""; + uint256 balanceBefore = zetaToken.balanceOf(destination); + assertEq(balanceBefore, 0); + + bytes memory data = abi.encodeWithSignature("transfer(address,uint256)", destination, amount); + vm.expectCall(address(zetaToken), 0, data); + vm.expectEmit(true, true, true, true, address(zetaConnectorV2)); + emit WithdrawnV2(destination, amount); + vm.prank(tssAddress); + zetaConnectorV2.withdraw(destination, amount, internalSendHash); + uint256 balanceAfter = zetaToken.balanceOf(destination); + assertEq(balanceAfter, amount); + } } diff --git a/v2/test/ZetaConnectorNonNative.t.sol b/v2/test/ZetaConnectorNonNative.t.sol index e8fc1ec7..c93e1b02 100644 --- a/v2/test/ZetaConnectorNonNative.t.sol +++ b/v2/test/ZetaConnectorNonNative.t.sol @@ -7,6 +7,7 @@ import "forge-std/Vm.sol"; import "./utils/ReceiverEVM.sol"; import "./utils/TestERC20.sol"; +import "./utils/upgrades/ZetaConnectorNonNativeUpgradeTest.sol"; import "../contracts/evm/ERC20Custody.sol"; import "../contracts/evm/GatewayEVM.sol"; @@ -44,6 +45,8 @@ contract ZetaConnectorNonNativeTest is error ExceedsMaxSupply(); error EnforcedPause(); + event WithdrawnV2(address indexed to, uint256 amount); + bytes32 public constant WITHDRAWER_ROLE = keccak256("WITHDRAWER_ROLE"); bytes32 public constant TSS_ROLE = keccak256("TSS_ROLE"); bytes32 public constant PAUSER_ROLE = keccak256("PAUSER_ROLE"); @@ -357,4 +360,25 @@ contract ZetaConnectorNonNativeTest is vm.expectRevert(abi.encodeWithSelector(AccessControlUnauthorizedAccount.selector, owner, TSS_ROLE)); zetaConnector.setMaxSupply(10_000); } + + function testUpgradeAndWithdraw() public { + // upgrade + Upgrades.upgradeProxy(address(zetaConnector), "ZetaConnectorNonNativeUpgradeTest.sol", "", owner); + ZetaConnectorNonNativeUpgradeTest zetaConnectorV2 = ZetaConnectorNonNativeUpgradeTest(address(zetaConnector)); + // withdraw + uint256 amount = 100_000; + uint256 balanceBefore = zetaToken.balanceOf(destination); + assertEq(balanceBefore, 0); + bytes32 internalSendHash = ""; + + bytes memory data = + abi.encodeWithSignature("mint(address,uint256,bytes32)", destination, amount, internalSendHash); + vm.expectCall(address(zetaToken), 0, data); + vm.expectEmit(true, true, true, true, address(zetaConnectorV2)); + emit WithdrawnV2(destination, amount); + vm.prank(tssAddress); + zetaConnectorV2.withdraw(destination, amount, internalSendHash); + uint256 balanceAfter = zetaToken.balanceOf(destination); + assertEq(balanceAfter, amount); + } } diff --git a/v2/test/utils/upgrades/ERC20CustodyUpgradeTest.sol b/v2/test/utils/upgrades/ERC20CustodyUpgradeTest.sol new file mode 100644 index 00000000..f2118fd7 --- /dev/null +++ b/v2/test/utils/upgrades/ERC20CustodyUpgradeTest.sol @@ -0,0 +1,224 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.26; + +import { IERC20Custody } from "../../../contracts/evm/interfaces/IERC20Custody.sol"; +import { IGatewayEVM } from "../../../contracts/evm/interfaces/IGatewayEVM.sol"; + +import { RevertContext } from "../../../contracts/Revert.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"; + +/// @title ERC20CustodyUpgradeTest +/// @notice Modified ERC20Custody contract for testing upgrades +/// @dev The only difference is in event naming +/// @custom:oz-upgrades-from ERC20Custody +contract ERC20CustodyUpgradeTest is + Initializable, + UUPSUpgradeable, + IERC20Custody, + ReentrancyGuardUpgradeable, + AccessControlUpgradeable, + PausableUpgradeable +{ + using SafeERC20 for IERC20; + + /// @notice Gateway contract. + 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. + address public tssAddress; + /// @notice Used to flag if contract supports legacy methods (eg. deposit). + bool public supportsLegacy; + /// @notice New role identifier for pauser role. + bytes32 public constant PAUSER_ROLE = keccak256("PAUSER_ROLE"); + /// @notice New role identifier for withdrawer role. + bytes32 public constant WITHDRAWER_ROLE = keccak256("WITHDRAWER_ROLE"); + /// @notice New role identifier for whitelister role. + bytes32 public constant WHITELISTER_ROLE = keccak256("WHITELISTER_ROLE"); + + /// @dev Modified event for testing upgrade. + event WithdrawnV2(address indexed to, address indexed token, uint256 amount); + + /// @notice Initializer for ERC20Custody. + /// @dev Set admin as default admin and pauser, and tssAddress as tss role. + 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_); + _grantRole(PAUSER_ROLE, admin_); + _grantRole(WITHDRAWER_ROLE, tssAddress_); + _grantRole(WHITELISTER_ROLE, admin_); + _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(); + } + + /// @notice Unpause contract. + function unpause() external onlyRole(PAUSER_ROLE) { + _unpause(); + } + + /// @notice Update tss address + /// @param newTSSAddress new tss address + function updateTSSAddress(address newTSSAddress) external onlyRole(DEFAULT_ADMIN_ROLE) { + if (newTSSAddress == address(0)) revert ZeroAddress(); + + _revokeRole(WITHDRAWER_ROLE, tssAddress); + _revokeRole(WHITELISTER_ROLE, tssAddress); + + _grantRole(WITHDRAWER_ROLE, newTSSAddress); + _grantRole(WHITELISTER_ROLE, newTSSAddress); + + tssAddress = newTSSAddress; + + emit UpdatedCustodyTSSAddress(newTSSAddress); + } + + /// @notice Unpause contract. + function setSupportsLegacy(bool _supportsLegacy) external onlyRole(DEFAULT_ADMIN_ROLE) { + supportsLegacy = _supportsLegacy; + } + + /// @notice Whitelist ERC20 token. + /// @param token address of ERC20 token + function whitelist(address token) external onlyRole(WHITELISTER_ROLE) { + if (token == address(0)) revert ZeroAddress(); + whitelisted[token] = true; + emit Whitelisted(token); + } + + /// @notice Unwhitelist ERC20 token. + /// @param token address of ERC20 token + function unwhitelist(address token) external onlyRole(WHITELISTER_ROLE) { + if (token == address(0)) revert ZeroAddress(); + whitelisted[token] = false; + emit Unwhitelisted(token); + } + + /// @notice Withdraw directly transfers the tokens to the destination address without contract call. + /// @dev This function can only be called by the TSS address. + /// @param to Destination address for the tokens. + /// @param token Address of the ERC20 token. + /// @param amount Amount of tokens to withdraw. + function withdraw( + address to, + address token, + uint256 amount + ) + external + nonReentrant + onlyRole(WITHDRAWER_ROLE) + whenNotPaused + { + if (!whitelisted[token]) revert NotWhitelisted(); + + IERC20(token).safeTransfer(to, amount); + + emit WithdrawnV2(to, token, amount); + } + + /// @notice WithdrawAndCall transfers tokens to Gateway and call a contract through the Gateway. + /// @dev This function can only be called by the TSS address. + /// @param to Address of the contract to call. + /// @param token Address of the ERC20 token. + /// @param amount Amount of tokens to withdraw. + /// @param data Calldata to pass to the contract call. + function withdrawAndCall( + address to, + address token, + uint256 amount, + bytes calldata data + ) + public + nonReentrant + onlyRole(WITHDRAWER_ROLE) + whenNotPaused + { + if (!whitelisted[token]) revert NotWhitelisted(); + + // Transfer the tokens to the Gateway contract + IERC20(token).safeTransfer(address(gateway), amount); + + // Forward the call to the Gateway contract + gateway.executeWithERC20(token, to, amount, data); + + emit WithdrawnAndCalled(to, token, amount, data); + } + + /// @notice WithdrawAndRevert transfers tokens to Gateway and call a contract with a revert functionality through + /// the Gateway. + /// @dev This function can only be called by the TSS address. + /// @param to Address of the contract to call. + /// @param token Address of the ERC20 token. + /// @param amount Amount of tokens to withdraw. + /// @param data Calldata to pass to the contract call. + /// @param revertContext Revert context to pass to onRevert. + function withdrawAndRevert( + address to, + address token, + uint256 amount, + bytes calldata data, + RevertContext calldata revertContext + ) + public + nonReentrant + onlyRole(WITHDRAWER_ROLE) + whenNotPaused + { + if (!whitelisted[token]) revert NotWhitelisted(); + + // Transfer the tokens to the Gateway contract + IERC20(token).safeTransfer(address(gateway), amount); + + // Forward the call to the Gateway contract + gateway.revertWithERC20(token, to, amount, data, revertContext); + + emit WithdrawnAndReverted(to, token, amount, data, revertContext); + } + + /// @notice Deposits asset to custody and pay fee in zeta erc20. + /// @custom:deprecated This method is deprecated. + function deposit( + bytes calldata recipient, + IERC20 asset, + uint256 amount, + bytes calldata message + ) + external + nonReentrant + whenNotPaused + { + if (!supportsLegacy) revert LegacyMethodsNotSupported(); + if (!whitelisted[address(asset)]) revert NotWhitelisted(); + uint256 oldBalance = asset.balanceOf(address(this)); + asset.safeTransferFrom(msg.sender, address(this), amount); + // In case if there is a fee on a token transfer, we might not receive a full expected amount + // and we need to correctly process that, we subtract an old balance from a new balance, which should be an + // actual received amount. + emit Deposited(recipient, asset, asset.balanceOf(address(this)) - oldBalance, message); + } +} diff --git a/v2/test/utils/GatewayEVMUpgradeTest.sol b/v2/test/utils/upgrades/GatewayEVMUpgradeTest.sol similarity index 98% rename from v2/test/utils/GatewayEVMUpgradeTest.sol rename to v2/test/utils/upgrades/GatewayEVMUpgradeTest.sol index bba9f692..a1cf5efc 100644 --- a/v2/test/utils/GatewayEVMUpgradeTest.sol +++ b/v2/test/utils/upgrades/GatewayEVMUpgradeTest.sol @@ -1,10 +1,10 @@ // SPDX-License-Identifier: MIT pragma solidity 0.8.26; -import { RevertContext, RevertOptions, Revertable } from "../../contracts/Revert.sol"; -import "../../contracts/evm/ZetaConnectorBase.sol"; -import "../../contracts/evm/interfaces/IERC20Custody.sol"; -import "../../contracts/evm/interfaces/IGatewayEVM.sol"; +import { RevertContext, RevertOptions, Revertable } from "../../../contracts/Revert.sol"; +import "../../../contracts/evm/ZetaConnectorBase.sol"; +import "../../../contracts/evm/interfaces/IERC20Custody.sol"; +import "../../../contracts/evm/interfaces/IGatewayEVM.sol"; import "@openzeppelin/contracts-upgradeable/access/AccessControlUpgradeable.sol"; import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol"; diff --git a/v2/test/utils/upgrades/GatewayZEVMUpgradeTest.sol b/v2/test/utils/upgrades/GatewayZEVMUpgradeTest.sol new file mode 100644 index 00000000..6c5a5ffe --- /dev/null +++ b/v2/test/utils/upgrades/GatewayZEVMUpgradeTest.sol @@ -0,0 +1,532 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.26; + +import { CallOptions, IGatewayZEVM } from "../../../contracts/zevm/interfaces/IGatewayZEVM.sol"; + +import { RevertContext, RevertOptions, Revertable } from "../../../contracts/Revert.sol"; +import "../../../contracts/zevm/interfaces/IWZETA.sol"; +import { IZRC20 } from "../../../contracts/zevm/interfaces/IZRC20.sol"; +import { UniversalContract, zContext } from "../../../contracts/zevm/interfaces/UniversalContract.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"; + +/// @title GatewayZEVMUpgradeTest +/// @notice Modified GatewayZEVM contract for testing upgrades +/// @dev The only difference is in event naming +/// @custom:oz-upgrades-from GatewayZEVM +contract GatewayZEVMUpgradeTest is + IGatewayZEVM, + Initializable, + AccessControlUpgradeable, + UUPSUpgradeable, + ReentrancyGuardUpgradeable, + PausableUpgradeable +{ + /// @notice Error indicating a zero address was provided. + error ZeroAddress(); + + /// @notice The constant address of the protocol + address public constant PROTOCOL_ADDRESS = 0x735b14BB79463307AAcBED86DAf3322B1e6226aB; + /// @notice The address of the Zeta token. + address public zetaToken; + + /// @notice New role identifier for pauser role. + bytes32 public constant PAUSER_ROLE = keccak256("PAUSER_ROLE"); + + /// @notice Max size of message + revertOptions revert message. + uint256 public constant MAX_MESSAGE_SIZE = 1024; + + /// @dev Modified event for testing upgrade. + event WithdrawnV2( + address indexed sender, + uint256 indexed chainId, + bytes receiver, + address zrc20, + uint256 value, + uint256 gasfee, + uint256 protocolFlatFee, + bytes message, + CallOptions callOptions, + RevertOptions revertOptions + ); + + /// @dev Only protocol address allowed modifier. + modifier onlyProtocol() { + if (msg.sender != PROTOCOL_ADDRESS) { + revert CallerIsNotProtocol(); + } + _; + } + + /// @custom:oz-upgrades-unsafe-allow constructor + constructor() { + _disableInitializers(); + } + + /// @notice Initialize with address of zeta token and admin account set as DEFAULT_ADMIN_ROLE. + /// @dev Using admin to authorize upgrades and pause. + function initialize(address zetaToken_, address admin_) public initializer { + if (zetaToken_ == address(0) || admin_ == address(0)) { + revert ZeroAddress(); + } + __UUPSUpgradeable_init(); + __AccessControl_init(); + __Pausable_init(); + __ReentrancyGuard_init(); + + _grantRole(DEFAULT_ADMIN_ROLE, admin_); + _grantRole(PAUSER_ROLE, admin_); + zetaToken = zetaToken_; + } + + /// @dev Authorizes the upgrade of the contract. + /// @param newImplementation The address of the new implementation. + function _authorizeUpgrade(address newImplementation) internal override onlyRole(DEFAULT_ADMIN_ROLE) { } + + /// @dev Receive function to receive ZETA from WETH9.withdraw(). + receive() external payable whenNotPaused { + if (msg.sender != zetaToken && msg.sender != PROTOCOL_ADDRESS) revert OnlyWZETAOrProtocol(); + } + + /// @notice Pause contract. + function pause() external onlyRole(PAUSER_ROLE) { + _pause(); + } + + /// @notice Unpause contract. + function unpause() external onlyRole(PAUSER_ROLE) { + _unpause(); + } + + /// @dev Internal function to withdraw ZRC20 tokens. + /// @param amount The amount of tokens to withdraw. + /// @param zrc20 The address of the ZRC20 token. + /// @return The gas fee for the withdrawal. + function _withdrawZRC20(uint256 amount, address zrc20) internal returns (uint256) { + // Use gas limit from zrc20 + return _withdrawZRC20WithGasLimit(amount, zrc20, IZRC20(zrc20).GAS_LIMIT()); + } + + /// @dev Internal function to withdraw ZRC20 tokens with gas limit. + /// @param amount The amount of tokens to withdraw. + /// @param zrc20 The address of the ZRC20 token. + /// @param gasLimit Gas limit. + /// @return The gas fee for the withdrawal. + function _withdrawZRC20WithGasLimit(uint256 amount, address zrc20, uint256 gasLimit) internal returns (uint256) { + (address gasZRC20, uint256 gasFee) = IZRC20(zrc20).withdrawGasFeeWithGasLimit(gasLimit); + if (!IZRC20(gasZRC20).transferFrom(msg.sender, PROTOCOL_ADDRESS, gasFee)) { + revert GasFeeTransferFailed(); + } + + if (!IZRC20(zrc20).transferFrom(msg.sender, address(this), amount)) { + revert ZRC20TransferFailed(); + } + + if (!IZRC20(zrc20).burn(amount)) revert ZRC20BurnFailed(); + + return gasFee; + } + + /// @dev Internal function to transfer ZETA tokens. + /// @param amount The amount of tokens to transfer. + /// @param to The address to transfer the tokens to. + function _transferZETA(uint256 amount, address to) internal { + if (!IWETH9(zetaToken).transferFrom(msg.sender, address(this), amount)) revert FailedZetaSent(); + IWETH9(zetaToken).withdraw(amount); + (bool sent,) = to.call{ value: amount }(""); + if (!sent) revert FailedZetaSent(); + } + + /// @notice Withdraw ZRC20 tokens to an external chain. + /// @param receiver The receiver address on the external chain. + /// @param amount The amount of tokens to withdraw. + /// @param zrc20 The address of the ZRC20 token. + /// @param revertOptions Revert options. + function withdraw( + bytes memory receiver, + uint256 amount, + address zrc20, + RevertOptions calldata revertOptions + ) + external + nonReentrant + whenNotPaused + { + if (receiver.length == 0) revert ZeroAddress(); + if (amount == 0) revert InsufficientZRC20Amount(); + + uint256 gasFee = _withdrawZRC20(amount, zrc20); + emit WithdrawnV2( + msg.sender, + 0, + receiver, + zrc20, + amount, + gasFee, + IZRC20(zrc20).PROTOCOL_FLAT_FEE(), + "", + CallOptions({ gasLimit: IZRC20(zrc20).GAS_LIMIT(), isArbitraryCall: true }), + revertOptions + ); + } + + /// @notice Withdraw ZRC20 tokens and call a smart contract on an external chain. + /// @param receiver The receiver address on the external chain. + /// @param amount The amount of tokens to withdraw. + /// @param zrc20 The address of the ZRC20 token. + /// @param message The calldata to pass to the contract call. + /// @param gasLimit Gas limit. + /// @param revertOptions Revert options. + function withdrawAndCall( + bytes memory receiver, + uint256 amount, + address zrc20, + bytes calldata message, + uint256 gasLimit, + RevertOptions calldata revertOptions + ) + external + nonReentrant + whenNotPaused + { + if (receiver.length == 0) revert ZeroAddress(); + if (amount == 0) revert InsufficientZRC20Amount(); + if (gasLimit == 0) revert InsufficientGasLimit(); + if (message.length + revertOptions.revertMessage.length >= MAX_MESSAGE_SIZE) revert MessageSizeExceeded(); + + uint256 gasFee = _withdrawZRC20WithGasLimit(amount, zrc20, gasLimit); + emit Withdrawn( + msg.sender, + 0, + receiver, + zrc20, + amount, + gasFee, + IZRC20(zrc20).PROTOCOL_FLAT_FEE(), + message, + CallOptions({ gasLimit: gasLimit, isArbitraryCall: true }), + revertOptions + ); + } + + /// @notice Withdraw ZRC20 tokens and call a smart contract on an external chain. + /// @param receiver The receiver address on the external chain. + /// @param amount The amount of tokens to withdraw. + /// @param zrc20 The address of the ZRC20 token. + /// @param message The calldata to pass to the contract call. + /// @param callOptions Call options including gas limit and arbirtrary call flag. + /// @param revertOptions Revert options. + function withdrawAndCall( + bytes memory receiver, + uint256 amount, + address zrc20, + bytes calldata message, + CallOptions calldata callOptions, + RevertOptions calldata revertOptions + ) + external + nonReentrant + whenNotPaused + { + if (receiver.length == 0) revert ZeroAddress(); + if (amount == 0) revert InsufficientZRC20Amount(); + if (callOptions.gasLimit == 0) revert InsufficientGasLimit(); + if (message.length + revertOptions.revertMessage.length >= MAX_MESSAGE_SIZE) revert MessageSizeExceeded(); + + uint256 gasFee = _withdrawZRC20WithGasLimit(amount, zrc20, callOptions.gasLimit); + emit Withdrawn( + msg.sender, + 0, + receiver, + zrc20, + amount, + gasFee, + IZRC20(zrc20).PROTOCOL_FLAT_FEE(), + message, + callOptions, + revertOptions + ); + } + + /// @notice Withdraw ZETA tokens to an external chain. + /// @param receiver The receiver address on the external chain. + /// @param amount The amount of tokens to withdraw. + /// @param revertOptions Revert options. + function withdraw( + bytes memory receiver, + uint256 amount, + uint256 chainId, + RevertOptions calldata revertOptions + ) + external + nonReentrant + whenNotPaused + { + if (receiver.length == 0) revert ZeroAddress(); + if (amount == 0) revert InsufficientZetaAmount(); + + _transferZETA(amount, PROTOCOL_ADDRESS); + emit Withdrawn( + msg.sender, + chainId, + receiver, + address(zetaToken), + amount, + 0, + 0, + "", + CallOptions({ gasLimit: 0, isArbitraryCall: true }), + revertOptions + ); + } + + /// @notice Withdraw ZETA tokens and call a smart contract on an external chain. + /// @param receiver The receiver address on the external chain. + /// @param amount The amount of tokens to withdraw. + /// @param chainId Chain id of the external chain. + /// @param message The calldata to pass to the contract call. + /// @param revertOptions Revert options. + function withdrawAndCall( + bytes memory receiver, + uint256 amount, + uint256 chainId, + bytes calldata message, + RevertOptions calldata revertOptions + ) + external + nonReentrant + whenNotPaused + { + if (receiver.length == 0) revert ZeroAddress(); + if (amount == 0) revert InsufficientZetaAmount(); + if (message.length + revertOptions.revertMessage.length >= MAX_MESSAGE_SIZE) revert MessageSizeExceeded(); + + _transferZETA(amount, PROTOCOL_ADDRESS); + emit Withdrawn( + msg.sender, + chainId, + receiver, + address(zetaToken), + amount, + 0, + 0, + message, + CallOptions({ gasLimit: 0, isArbitraryCall: true }), + revertOptions + ); + } + + /// @notice Withdraw ZETA tokens and call a smart contract on an external chain. + /// @param receiver The receiver address on the external chain. + /// @param amount The amount of tokens to withdraw. + /// @param chainId Chain id of the external chain. + /// @param message The calldata to pass to the contract call. + /// @param callOptions Call options including gas limit and arbirtrary call flag. + /// @param revertOptions Revert options. + function withdrawAndCall( + bytes memory receiver, + uint256 amount, + uint256 chainId, + bytes calldata message, + CallOptions calldata callOptions, + RevertOptions calldata revertOptions + ) + external + nonReentrant + whenNotPaused + { + if (receiver.length == 0) revert ZeroAddress(); + if (amount == 0) revert InsufficientZetaAmount(); + if (callOptions.gasLimit == 0) revert InsufficientGasLimit(); + if (message.length + revertOptions.revertMessage.length >= MAX_MESSAGE_SIZE) revert MessageSizeExceeded(); + + _transferZETA(amount, PROTOCOL_ADDRESS); + emit Withdrawn( + msg.sender, chainId, receiver, address(zetaToken), amount, 0, 0, message, callOptions, revertOptions + ); + } + + /// @notice Call a smart contract on an external chain without asset transfer. + /// @param receiver The receiver address on the external chain. + /// @param zrc20 Address of zrc20 to pay fees. + /// @param message The calldata to pass to the contract call. + /// @param callOptions Call options including gas limit and arbirtrary call flag. + /// @param revertOptions Revert options. + function call( + bytes memory receiver, + address zrc20, + bytes calldata message, + CallOptions calldata callOptions, + RevertOptions calldata revertOptions + ) + external + nonReentrant + whenNotPaused + { + if (callOptions.gasLimit == 0) revert InsufficientGasLimit(); + if (message.length + revertOptions.revertMessage.length >= MAX_MESSAGE_SIZE) revert MessageSizeExceeded(); + + _call(receiver, zrc20, message, callOptions, revertOptions); + } + + /// @notice Call a smart contract on an external chain without asset transfer. + /// @param receiver The receiver address on the external chain. + /// @param zrc20 Address of zrc20 to pay fees. + /// @param message The calldata to pass to the contract call. + /// @param gasLimit Gas limit. + /// @param revertOptions Revert options. + function call( + bytes memory receiver, + address zrc20, + bytes calldata message, + uint256 gasLimit, + RevertOptions calldata revertOptions + ) + external + nonReentrant + whenNotPaused + { + if (gasLimit == 0) revert InsufficientGasLimit(); + if (message.length + revertOptions.revertMessage.length >= MAX_MESSAGE_SIZE) revert MessageSizeExceeded(); + + _call(receiver, zrc20, message, CallOptions({ gasLimit: gasLimit, isArbitraryCall: true }), revertOptions); + } + + function _call( + bytes memory receiver, + address zrc20, + bytes calldata message, + CallOptions memory callOptions, + RevertOptions memory revertOptions + ) + internal + { + if (receiver.length == 0) revert ZeroAddress(); + + (address gasZRC20, uint256 gasFee) = IZRC20(zrc20).withdrawGasFeeWithGasLimit(callOptions.gasLimit); + if (!IZRC20(gasZRC20).transferFrom(msg.sender, PROTOCOL_ADDRESS, gasFee)) { + revert GasFeeTransferFailed(); + } + + emit Called(msg.sender, zrc20, receiver, message, callOptions, revertOptions); + } + + /// @notice Deposit foreign coins into ZRC20. + /// @param zrc20 The address of the ZRC20 token. + /// @param amount The amount of tokens to deposit. + /// @param target The target address to receive the deposited tokens. + function deposit(address zrc20, uint256 amount, address target) external onlyProtocol whenNotPaused { + if (zrc20 == address(0) || target == address(0)) revert ZeroAddress(); + if (amount == 0) revert InsufficientZRC20Amount(); + + if (target == PROTOCOL_ADDRESS || target == address(this)) revert InvalidTarget(); + + if (!IZRC20(zrc20).deposit(target, amount)) revert ZRC20DepositFailed(); + } + + /// @notice Execute a user-specified contract on ZEVM. + /// @param context The context of the cross-chain call. + /// @param zrc20 The address of the ZRC20 token. + /// @param amount The amount of tokens to transfer. + /// @param target The target contract to call. + /// @param message The calldata to pass to the contract call. + function execute( + zContext calldata context, + address zrc20, + uint256 amount, + address target, + bytes calldata message + ) + external + onlyProtocol + whenNotPaused + { + if (zrc20 == address(0) || target == address(0)) revert ZeroAddress(); + + UniversalContract(target).onCrossChainCall(context, zrc20, amount, message); + } + + /// @notice Deposit foreign coins into ZRC20 and call a user-specified contract on ZEVM. + /// @param context The context of the cross-chain call. + /// @param zrc20 The address of the ZRC20 token. + /// @param amount The amount of tokens to transfer. + /// @param target The target contract to call. + /// @param message The calldata to pass to the contract call. + function depositAndCall( + zContext calldata context, + address zrc20, + uint256 amount, + address target, + bytes calldata message + ) + external + onlyProtocol + whenNotPaused + { + if (zrc20 == address(0) || target == address(0)) revert ZeroAddress(); + if (amount == 0) revert InsufficientZRC20Amount(); + if (target == PROTOCOL_ADDRESS || target == address(this)) revert InvalidTarget(); + + if (!IZRC20(zrc20).deposit(target, amount)) revert ZRC20DepositFailed(); + UniversalContract(target).onCrossChainCall(context, zrc20, amount, message); + } + + /// @notice Deposit ZETA and call a user-specified contract on ZEVM. + /// @param context The context of the cross-chain call. + /// @param amount The amount of tokens to transfer. + /// @param target The target contract to call. + /// @param message The calldata to pass to the contract call. + function depositAndCall( + zContext calldata context, + uint256 amount, + address target, + bytes calldata message + ) + external + onlyProtocol + whenNotPaused + { + if (target == address(0)) revert ZeroAddress(); + if (amount == 0) revert InsufficientZetaAmount(); + if (target == PROTOCOL_ADDRESS || target == address(this)) revert InvalidTarget(); + + _transferZETA(amount, target); + UniversalContract(target).onCrossChainCall(context, zetaToken, amount, message); + } + + /// @notice Revert a user-specified contract on ZEVM. + /// @param target The target contract to call. + /// @param revertContext Revert context to pass to onRevert. + function executeRevert(address target, RevertContext calldata revertContext) external onlyProtocol whenNotPaused { + if (target == address(0)) revert ZeroAddress(); + + Revertable(target).onRevert(revertContext); + } + + /// @notice Deposit foreign coins into ZRC20 and revert a user-specified contract on ZEVM. + /// @param zrc20 The address of the ZRC20 token. + /// @param amount The amount of tokens to revert. + /// @param target The target contract to call. + /// @param revertContext Revert context to pass to onRevert. + function depositAndRevert( + address zrc20, + uint256 amount, + address target, + RevertContext calldata revertContext + ) + external + onlyProtocol + whenNotPaused + { + if (zrc20 == address(0) || target == address(0)) revert ZeroAddress(); + if (amount == 0) revert InsufficientZRC20Amount(); + if (target == PROTOCOL_ADDRESS || target == address(this)) revert InvalidTarget(); + + if (!IZRC20(zrc20).deposit(target, amount)) revert ZRC20DepositFailed(); + Revertable(target).onRevert(revertContext); + } +} diff --git a/v2/test/utils/upgrades/ZetaConnectorNativeUpgradeTest.sol b/v2/test/utils/upgrades/ZetaConnectorNativeUpgradeTest.sol new file mode 100644 index 00000000..f4a56b2e --- /dev/null +++ b/v2/test/utils/upgrades/ZetaConnectorNativeUpgradeTest.sol @@ -0,0 +1,112 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.26; + +import "../../../contracts/evm/ZetaConnectorBase.sol"; +import "@openzeppelin/contracts/token/ERC20/IERC20.sol"; +import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; + +/// @title ZetaConnectorNativeUpgradeTest +/// @notice Modified ZetaConnectorNative contract for testing upgrades +/// @dev The only difference is in event naming +/// @custom:oz-upgrades-from ZetaConnectorNative +contract ZetaConnectorNativeUpgradeTest is ZetaConnectorBase { + using SafeERC20 for IERC20; + + /// @dev Modified event for testing upgrade. + event WithdrawnV2(address indexed to, uint256 amount); + + function initialize( + address gateway_, + address zetaToken_, + address tssAddress_, + address admin_ + ) + public + override + initializer + { + super.initialize(gateway_, zetaToken_, tssAddress_, admin_); + } + + /// @notice Withdraw tokens to a specified address. + /// @param to The address to withdraw tokens to. + /// @param amount The amount of tokens to withdraw. + /// @param internalSendHash A hash used for internal tracking of the transaction. + /// @dev This function can only be called by the TSS address. + function withdraw( + address to, + uint256 amount, + bytes32 internalSendHash + ) + external + override + nonReentrant + onlyRole(WITHDRAWER_ROLE) + whenNotPaused + { + IERC20(zetaToken).safeTransfer(to, amount); + emit WithdrawnV2(to, amount); + } + + /// @notice Withdraw tokens and call a contract through Gateway. + /// @param to The address to withdraw tokens to. + /// @param amount The amount of tokens to withdraw. + /// @param data The calldata to pass to the contract call. + /// @param internalSendHash A hash used for internal tracking of the transaction. + /// @dev This function can only be called by the TSS address. + function withdrawAndCall( + address to, + uint256 amount, + bytes calldata data, + bytes32 internalSendHash + ) + external + override + nonReentrant + onlyRole(WITHDRAWER_ROLE) + whenNotPaused + { + // Transfer zetaToken to the Gateway contract + IERC20(zetaToken).safeTransfer(address(gateway), amount); + + // Forward the call to the Gateway contract + gateway.executeWithERC20(address(zetaToken), to, amount, data); + + emit WithdrawnAndCalled(to, amount, data); + } + + /// @notice Withdraw tokens and call a contract with a revert callback through Gateway. + /// @param to The address to withdraw tokens to. + /// @param amount The amount of tokens to withdraw. + /// @param data The calldata to pass to the contract call. + /// @param internalSendHash A hash used for internal tracking of the transaction. + /// @dev This function can only be called by the TSS address. + /// @param revertContext Revert context to pass to onRevert. + function withdrawAndRevert( + address to, + uint256 amount, + bytes calldata data, + bytes32 internalSendHash, + RevertContext calldata revertContext + ) + external + override + nonReentrant + onlyRole(WITHDRAWER_ROLE) + whenNotPaused + { + // 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, revertContext); + + emit WithdrawnAndReverted(to, amount, data, revertContext); + } + + /// @notice Handle received tokens. + /// @param amount The amount of tokens received. + function receiveTokens(uint256 amount) external override whenNotPaused { + IERC20(zetaToken).safeTransferFrom(msg.sender, address(this), amount); + } +} diff --git a/v2/test/utils/upgrades/ZetaConnectorNonNativeUpgradeTest.sol b/v2/test/utils/upgrades/ZetaConnectorNonNativeUpgradeTest.sol new file mode 100644 index 00000000..adbec1f4 --- /dev/null +++ b/v2/test/utils/upgrades/ZetaConnectorNonNativeUpgradeTest.sol @@ -0,0 +1,134 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.26; + +import "../../../contracts/evm/ZetaConnectorBase.sol"; +import "../../../contracts/evm/interfaces/IZetaNonEthNew.sol"; + +/// @title ZetaConnectorNonNativeUpgradeTest +/// @notice Modified ZetaConnectorNonNative contract for testing upgrades +/// @dev The only difference is in event naming +/// @custom:oz-upgrades-from ZetaConnectorNonNative +contract ZetaConnectorNonNativeUpgradeTest is ZetaConnectorBase { + /// @notice Event triggered when max supply is updated. + /// @param maxSupply New max supply. + event MaxSupplyUpdated(uint256 maxSupply); + + error ExceedsMaxSupply(); + + /// @notice Max supply for minting. + uint256 public maxSupply; + + /// @dev Modified event for testing upgrade. + event WithdrawnV2(address indexed to, uint256 amount); + + function initialize( + address gateway_, + address zetaToken_, + address tssAddress_, + address admin_ + ) + public + override + initializer + { + super.initialize(gateway_, zetaToken_, tssAddress_, admin_); + + maxSupply = type(uint256).max; + } + + /// @notice Set max supply for minting. + /// @param maxSupply_ New max supply. + /// @dev This function can only be called by the TSS address. + function setMaxSupply(uint256 maxSupply_) external onlyRole(TSS_ROLE) whenNotPaused { + maxSupply = maxSupply_; + emit MaxSupplyUpdated(maxSupply_); + } + + /// @notice Withdraw tokens to a specified address. + /// @param to The address to withdraw tokens to. + /// @param amount The amount of tokens to withdraw. + /// @param internalSendHash A hash used for internal tracking of the transaction. + /// @dev This function can only be called by the TSS address. + function withdraw( + address to, + uint256 amount, + bytes32 internalSendHash + ) + external + override + nonReentrant + onlyRole(WITHDRAWER_ROLE) + whenNotPaused + { + _mintTo(to, amount, internalSendHash); + emit WithdrawnV2(to, amount); + } + + /// @notice Withdraw tokens and call a contract through Gateway. + /// @param to The address to withdraw tokens to. + /// @param amount The amount of tokens to withdraw. + /// @param data The calldata to pass to the contract call. + /// @param internalSendHash A hash used for internal tracking of the transaction. + /// @dev This function can only be called by the TSS address, and mints if supply is not reached. + function withdrawAndCall( + address to, + uint256 amount, + bytes calldata data, + bytes32 internalSendHash + ) + external + override + nonReentrant + onlyRole(WITHDRAWER_ROLE) + whenNotPaused + { + // Mint zetaToken to the Gateway contract + _mintTo(address(gateway), amount, internalSendHash); + + // Forward the call to the Gateway contract + gateway.executeWithERC20(address(zetaToken), to, amount, data); + + emit WithdrawnAndCalled(to, amount, data); + } + + /// @notice Withdraw tokens and call a contract with a revert callback through Gateway. + /// @param to The address to withdraw tokens to. + /// @param amount The amount of tokens to withdraw. + /// @param data The calldata to pass to the contract call. + /// @param internalSendHash A hash used for internal tracking of the transaction. + /// @dev This function can only be called by the TSS address, and mints if supply is not reached. + /// @param revertContext Revert context to pass to onRevert. + function withdrawAndRevert( + address to, + uint256 amount, + bytes calldata data, + bytes32 internalSendHash, + RevertContext calldata revertContext + ) + external + override + nonReentrant + onlyRole(WITHDRAWER_ROLE) + whenNotPaused + { + // Mint zetaToken to the Gateway contract + _mintTo(address(gateway), amount, internalSendHash); + + // Forward the call to the Gateway contract + gateway.revertWithERC20(address(zetaToken), to, amount, data, revertContext); + + emit WithdrawnAndReverted(to, amount, data, revertContext); + } + + /// @notice Handle received tokens and burn them. + /// @param amount The amount of tokens received. + function receiveTokens(uint256 amount) external override whenNotPaused { + IZetaNonEthNew(zetaToken).burnFrom(msg.sender, amount); + } + + /// @dev mints to provided account and checks if totalSupply will be exceeded + function _mintTo(address to, uint256 amount, bytes32 internalSendHash) internal { + if (amount + IERC20(zetaToken).totalSupply() > maxSupply) revert ExceedsMaxSupply(); + IZetaNonEthNew(zetaToken).mint(address(to), amount, internalSendHash); + } +}