feat: revert specs

contracts/prototypes/evm/ERC20CustodyNew.sol

@@ -17,6 +17,7 @@ contract ERC20CustodyNew is ReentrancyGuard{
 
     event Withdraw(address indexed token, address indexed to, uint256 amount);
     event WithdrawAndCall(address indexed token, address indexed to, uint256 amount, bytes data);
+    event WithdrawAndRevert(address indexed token, address indexed to, uint256 amount, bytes data);
 
     constructor(address _gateway) {
          if (_gateway == address(0)) {
@@ -47,4 +48,18 @@ contract ERC20CustodyNew is ReentrancyGuard{
 
         emit WithdrawAndCall(token, to, amount, data);
     }
+
+    // WithdrawAndRevert is called by TSS address, it transfers the tokens and call a contract
+    // For this, it passes through the Gateway contract, it transfers the tokens to the Gateway contract and then calls the contract
+    // TODO: Finalize access control
+    // https://github.com/zeta-chain/protocol-contracts/issues/204
+    function withdrawAndRevert(address token, address to, uint256 amount, bytes calldata data) public nonReentrant {
+        // 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);
+
+        emit WithdrawAndRevert(token, to, amount, data);
+    }
 }

contracts/prototypes/evm/GatewayEVM.sol

@@ -6,6 +6,7 @@
 import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/security/ReentrancyGuardUpgradeable.sol";
 import "./IGatewayEVM.sol";
 import "./ZetaConnectorNewBase.sol";
 
@@ -14,47 +15,56 @@
  * @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 {
+contract GatewayEVM is Initializable, OwnableUpgradeable, UUPSUpgradeable, IGatewayEVMErrors, IGatewayEVMEvents, ReentrancyGuardUpgradeable {
     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, address _zetaToken) public initializer {
-        __Ownable_init();
-        __UUPSUpgradeable_init();
-
         if (_tssAddress == address(0) || _zetaToken == address(0)) {
             revert ZeroAddress();
         }
+
+        __Ownable_init();
+        __UUPSUpgradeable_init();
+        __ReentrancyGuard_init();
 
         tssAddress = _tssAddress;
         zetaToken = _zetaToken;
     }
 
     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
+    // Calling onRevert directly
+    function executeRevert(address destination, bytes calldata data) public payable {
+        (bool success, bytes memory result) = destination.call{value: msg.value}("");
+        if (!success) revert ExecutionFailed();
+        Revertable(destination).onRevert(data);
+
+        emit Reverted(destination, msg.value, data);
+    }
+
     // 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
@@ -75,7 +85,7 @@
         address to,
         uint256 amount,
         bytes calldata data
-    ) public {
+    ) public nonReentrant {
         if (amount == 0) revert InsufficientETHAmount();
         // Approve the target contract to spend the tokens
         if(!resetApproval(token, to)) revert ApprovalFailed();
@@ -96,6 +106,22 @@
         emit ExecutedWithERC20(token, to, amount, data);
     }
 
+    // Called by the ERC20Custody contract
+    // Directly transfers ERC20 and calls onRevert
+    function revertWithERC20(
+        address token,
+        address to,
+        uint256 amount,
+        bytes calldata data
+    ) external nonReentrant {
+        if (amount == 0) revert InsufficientERC20Amount();
+
+        IERC20(token).safeTransfer(address(to), amount);
+        Revertable(to).onRevert(data);
+
+        emit RevertedWithERC20(token, to, amount, data);
+    }
+
     // Deposit ETH to tss
     function deposit(address receiver) external payable {
         if (msg.value == 0) revert InsufficientETHAmount();

contracts/prototypes/evm/GatewayEVMUpgradeTest.sol

@@ -1,49 +1,68 @@
 // SPDX-License-Identifier: MIT
-pragma solidity 0.8.7;
+pragma solidity ^0.8.0;
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/security/ReentrancyGuardUpgradeable.sol";
 import "./IGatewayEVM.sol";
+import "./ZetaConnectorNewBase.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, IGatewayEVMErrors, IGatewayEVMEvents {
+contract GatewayEVMUpgradeTest is Initializable, OwnableUpgradeable, UUPSUpgradeable, IGatewayEVMErrors, IGatewayEVMEvents, ReentrancyGuardUpgradeable {
     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;
 
     event ExecutedV2(address indexed destination, uint256 value, bytes data);
 
     constructor() {}
 
     function initialize(address _tssAddress, address _zetaToken) public initializer {
+        if (_tssAddress == address(0) || _zetaToken == address(0)) {
+            revert ZeroAddress();
+        }
+
         __Ownable_init();
         __UUPSUpgradeable_init();
-
-        if (_tssAddress == address(0)) revert ZeroAddress();
+        __ReentrancyGuard_init();
 
         tssAddress = _tssAddress;
         zetaToken = _zetaToken;
     }
 
     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
+    // Calling onRevert directly
+    function executeRevert(address destination, bytes calldata data) public payable {
+        (bool success, bytes memory result) = destination.call{value: msg.value}("");
+        if (!success) revert ExecutionFailed();
+        Revertable(destination).onRevert(data);
+
+        emit Reverted(destination, msg.value, data);
+    }
+
     // 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
@@ -64,7 +83,7 @@
         address to,
         uint256 amount,
         bytes calldata data
-    ) public returns (bytes memory) {
+    ) public nonReentrant {
         if (amount == 0) revert InsufficientETHAmount();
         // Approve the target contract to spend the tokens
         if(!resetApproval(token, to)) revert ApprovalFailed();
@@ -79,16 +98,26 @@
         // Transfer any remaining tokens back to the custody/connector contract
         uint256 remainingBalance = IERC20(token).balanceOf(address(this));
         if (remainingBalance > 0) {
-             address destination = address(custody);
-            if (token == zetaToken) {
-                destination = address(zetaConnector);
-            }
-            IERC20(token).safeTransfer(address(destination), remainingBalance);
+            transferToAssetHandler(token, amount);
         }
 
         emit ExecutedWithERC20(token, to, amount, data);
+    }
 
-        return result;
+    // Called by the ERC20Custody contract
+    // Directly transfers ERC20 and calls onRevert
+    function revertWithERC20(
+        address token,
+        address to,
+        uint256 amount,
+        bytes calldata data
+    ) external nonReentrant {
+        if (amount == 0) revert InsufficientERC20Amount();
+
+        IERC20(token).safeTransfer(address(to), amount);
+        Revertable(to).onRevert(data);
+
+        emit RevertedWithERC20(token, to, amount, data);
     }
 
     // Deposit ETH to tss
@@ -101,15 +130,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();
 
-        address destination = address(custody);
-        if (asset == zetaToken) {
-            destination = address(zetaConnector);
-        }
-        IERC20(asset).safeTransferFrom(msg.sender, address(destination), amount);
+        transferFromToAssetHandler(msg.sender, asset, amount);
 
         emit Deposit(msg.sender, receiver, amount, asset, "");
     }
@@ -124,15 +149,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();
 
-        address destination = address(custody);
-        if (asset == zetaToken) {
-            destination = address(zetaConnector);
-        }
-        IERC20(asset).safeTransferFrom(msg.sender, address(destination), amount);
+        transferFromToAssetHandler(msg.sender, asset, amount);
 
         emit Deposit(msg.sender, receiver, amount, asset, payload);
     }
@@ -144,15 +165,43 @@
 
     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 transferFromToAssetHandler(address from, address token, uint256 amount) private {
+        if (token == zetaToken) { // transfer to connector
+            // transfer amount to gateway
+            IERC20(token).safeTransferFrom(from, address(this), amount);
+            // approve connector to handle tokens depending on connector version (eg. lock or burn)
+            IERC20(token).approve(zetaConnector, amount);
+            // send tokens to connector
+            ZetaConnectorNewBase(zetaConnector).receiveTokens(amount);
+        } else { // transfer to custody
+            IERC20(token).safeTransferFrom(from, custody, amount);
+        }
+    }
+
+    function transferToAssetHandler(address token, uint256 amount) private {
+        if (token == zetaToken) { // transfer to connector
+            // approve connector to handle tokens depending on connector version (eg. lock or burn)
+            IERC20(token).approve(zetaConnector, amount);
+            // send tokens to connector
+            ZetaConnectorNewBase(zetaConnector).receiveTokens(amount);
+        } else { // transfer to custody
+            IERC20(token).safeTransfer(custody, amount);
+        }
+    }
 }

contracts/prototypes/evm/IGatewayEVM.sol

@@ -3,7 +3,9 @@ pragma solidity ^0.8.0;
 
 interface IGatewayEVMEvents {
     event Executed(address indexed destination, uint256 value, bytes data);
+    event Reverted(address indexed destination, uint256 value, bytes data);
     event ExecutedWithERC20(address indexed token, address indexed to, uint256 amount, bytes data);
+    event RevertedWithERC20(address indexed token, address indexed to, uint256 amount, bytes data);
     event Deposit(address indexed sender, address indexed receiver, uint256 amount, address asset, bytes payload);
     event Call(address indexed sender, address indexed receiver, bytes payload);
 }
@@ -27,4 +29,15 @@ interface IGatewayEVM {
     ) external;
 
     function execute(address destination, bytes calldata data) external payable returns (bytes memory);
-}
+
+    function revertWithERC20(
+        address token,
+        address to,
+        uint256 amount,
+        bytes calldata data
+    ) external;
+}
+
+interface Revertable {
+    function onRevert(bytes calldata data) external;
+}