feat: v2 contracts access control (#247)

zeta-chain · Jul 25, 2024 · 9b5ff23 · 9b5ff23
1 parent a2da596
commit 9b5ff23
Show file tree

Hide file tree

Showing 60 changed files with 2,815 additions and 156 deletions.
diff --git a/contracts/prototypes/evm/ERC20CustodyNew.sol b/contracts/prototypes/evm/ERC20CustodyNew.sol
@@ -2,44 +2,47 @@
 pragma solidity ^0.8.0;
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
-import "./IGatewayEVM.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@openzeppelin/contracts/security/ReentrancyGuard.sol";
 
+import "./IGatewayEVM.sol";
+import "./IERC20CustodyNew.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 is ReentrancyGuard{
+contract ERC20CustodyNew is IERC20CustodyNewEvents, IERC20CustodyNewErrors, ReentrancyGuard {
     using SafeERC20 for IERC20;
-    error ZeroAddress();
 
     IGatewayEVM public gateway;
+    address public tssAddress;
 
-    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);
+    // @dev Only TSS address allowed modifier.
+    modifier onlyTSS() {
+        if (msg.sender != tssAddress) {
+            revert InvalidSender();
+        }
+        _;
+    }
 
-    constructor(address _gateway) {
-         if (_gateway == address(0)) {
+    constructor(address _gateway, address _tssAddress) {
+         if (_gateway == address(0) || _tssAddress == address(0)) {
             revert ZeroAddress();
         }
         gateway = IGatewayEVM(_gateway);
+        tssAddress = _tssAddress;
     }
-    
+
     // Withdraw is called by TSS address, it directly transfers the tokens to the destination address without contract call
-    // TODO: Finalize access control
-    // https://github.com/zeta-chain/protocol-contracts/issues/204
-    function withdraw(address token, address to, uint256 amount) external nonReentrant {
+    function withdraw(address token, address to, uint256 amount) external nonReentrant onlyTSS {
         IERC20(token).safeTransfer(to, amount);
 
         emit Withdraw(token, to, amount);
     }
 
     // WithdrawAndCall 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 withdrawAndCall(address token, address to, uint256 amount, bytes calldata data) public nonReentrant {
+    function withdrawAndCall(address token, address to, uint256 amount, bytes calldata data) public nonReentrant onlyTSS {
         // Transfer the tokens to the Gateway contract
         IERC20(token).safeTransfer(address(gateway), amount);
 
@@ -51,9 +54,7 @@ contract ERC20CustodyNew is ReentrancyGuard{
 
     // 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 {
+    function withdrawAndRevert(address token, address to, uint256 amount, bytes calldata data) public nonReentrant onlyTSS {
         // Transfer the tokens to the Gateway contract
         IERC20(token).safeTransfer(address(gateway), amount);
 

diff --git a/contracts/prototypes/evm/GatewayEVM.sol b/contracts/prototypes/evm/GatewayEVM.sol
@@ -27,6 +27,22 @@ contract GatewayEVM is Initializable, OwnableUpgradeable, UUPSUpgradeable, IGate
     /// @notice The address of the Zeta token contract.
     address public zetaToken;
 
+    // @dev Only TSS address allowed modifier.
+    modifier onlyTSS() {
+        if (msg.sender != tssAddress) {
+            revert InvalidSender();
+        }
+        _;
+    }
+
+    // @dev Only asset handler (custody, connector) allowed modifier.
+    modifier onlyAssetHandler() {
+        if (msg.sender != custody && msg.sender != zetaConnector) {
+            revert InvalidSender();
+        }
+        _;
+    }
+
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor() {
         _disableInitializers();
@@ -56,7 +72,7 @@ contract GatewayEVM is Initializable, OwnableUpgradeable, UUPSUpgradeable, IGate
 
     // Called by the TSS
     // Calling onRevert directly
-    function executeRevert(address destination, bytes calldata data) public payable {
+    function executeRevert(address destination, bytes calldata data) public payable onlyTSS {
         (bool success, bytes memory result) = destination.call{value: msg.value}("");
         if (!success) revert ExecutionFailed();
         Revertable(destination).onRevert(data);
@@ -67,7 +83,7 @@ contract GatewayEVM is Initializable, OwnableUpgradeable, UUPSUpgradeable, IGate
     // 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) {
+    function execute(address destination, bytes calldata data) external payable onlyTSS returns (bytes memory)  {
         bytes memory result = _execute(destination, data);
 
         emit Executed(destination, msg.value, data);
@@ -84,7 +100,7 @@ contract GatewayEVM is Initializable, OwnableUpgradeable, UUPSUpgradeable, IGate
         address to,
         uint256 amount,
         bytes calldata data
-    ) public nonReentrant {
+    ) public nonReentrant onlyAssetHandler {
         if (amount == 0) revert InsufficientERC20Amount();
         // Approve the target contract to spend the tokens
         if(!resetApproval(token, to)) revert ApprovalFailed();
@@ -111,7 +127,7 @@ contract GatewayEVM is Initializable, OwnableUpgradeable, UUPSUpgradeable, IGate
         address to,
         uint256 amount,
         bytes calldata data
-    ) external nonReentrant {
+    ) external nonReentrant onlyAssetHandler {
         if (amount == 0) revert InsufficientERC20Amount();
 
         IERC20(token).safeTransfer(address(to), amount);
@@ -163,14 +179,14 @@ contract GatewayEVM is Initializable, OwnableUpgradeable, UUPSUpgradeable, IGate
         emit Call(msg.sender, receiver, payload);
     }
 
-    function setCustody(address _custody) external {
+    function setCustody(address _custody) external onlyTSS {
         if (custody != address(0)) revert CustodyInitialized();
         if (_custody == address(0)) revert ZeroAddress();
 
         custody = _custody;
     }
 
-     function setConnector(address _zetaConnector) external {
+     function setConnector(address _zetaConnector) external onlyTSS {
         if (zetaConnector != address(0)) revert CustodyInitialized();
         if (_zetaConnector == address(0)) revert ZeroAddress();
 

diff --git a/contracts/prototypes/evm/IERC20CustodyNew.sol b/contracts/prototypes/evm/IERC20CustodyNew.sol
@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+interface IERC20CustodyNewEvents {
+    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);
+}
+
+interface IERC20CustodyNewErrors {
+    error ZeroAddress();
+    error InvalidSender();
+}
diff --git a/contracts/prototypes/evm/IGatewayEVM.sol b/contracts/prototypes/evm/IGatewayEVM.sol
@@ -18,6 +18,7 @@ interface IGatewayEVMErrors {
     error ZeroAddress();
     error ApprovalFailed();
     error CustodyInitialized();
+    error InvalidSender();
 }
 
 interface IGatewayEVM {

diff --git a/contracts/prototypes/evm/IZetaConnector.sol b/contracts/prototypes/evm/IZetaConnector.sol
@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+interface IZetaConnectorEvents {
+    event Withdraw(address indexed to, uint256 amount);
+    event WithdrawAndCall(address indexed to, uint256 amount, bytes data);
+    event WithdrawAndRevert(address indexed to, uint256 amount, bytes data);
+}
diff --git a/contracts/prototypes/evm/ZetaConnectorNative.sol b/contracts/prototypes/evm/ZetaConnectorNative.sol
@@ -8,18 +8,18 @@ import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 contract ZetaConnectorNative is ZetaConnectorNewBase {
     using SafeERC20 for IERC20;
 
-    constructor(address _gateway, address _zetaToken)
-        ZetaConnectorNewBase(_gateway, _zetaToken)
+    constructor(address _gateway, address _zetaToken, address _tssAddress)
+        ZetaConnectorNewBase(_gateway, _zetaToken, _tssAddress)
     {}
 
     // @dev withdraw is called by TSS address, it directly transfers zetaToken to the destination address without contract call
-    function withdraw(address to, uint256 amount, bytes32 internalSendHash) external override nonReentrant {
+    function withdraw(address to, uint256 amount, bytes32 internalSendHash) external override nonReentrant onlyTSS {
         IERC20(zetaToken).safeTransfer(to, amount);
         emit Withdraw(to, amount);
     }
 
     // @dev withdrawAndCall is called by TSS address, it transfers zetaToken to the gateway and calls a contract
-    function withdrawAndCall(address to, uint256 amount, bytes calldata data, bytes32 internalSendHash) external override nonReentrant {
+    function withdrawAndCall(address to, uint256 amount, bytes calldata data, bytes32 internalSendHash) external override nonReentrant onlyTSS {
         // Transfer zetaToken to the Gateway contract
         IERC20(zetaToken).safeTransfer(address(gateway), amount);
 
@@ -30,7 +30,7 @@ contract ZetaConnectorNative is ZetaConnectorNewBase {
     }
 
     // @dev withdrawAndRevert is called by TSS address, it transfers zetaToken to the gateway and calls onRevert on a contract
-    function withdrawAndRevert(address to, uint256 amount, bytes calldata data, bytes32 internalSendHash) external override nonReentrant {
+    function withdrawAndRevert(address to, uint256 amount, bytes calldata data, bytes32 internalSendHash) external override nonReentrant onlyTSS {
         // Transfer zetaToken to the Gateway contract
         IERC20(zetaToken).safeTransfer(address(gateway), amount);
 
@@ -40,7 +40,7 @@ contract ZetaConnectorNative is ZetaConnectorNewBase {
         emit WithdrawAndRevert(to, amount, data);
     }
 
-    // @dev receiveTokens handles token transfer and burn them
+    // @dev receiveTokens handles token transfer back to connector
     function receiveTokens(uint256 amount) external override {
         IERC20(zetaToken).safeTransferFrom(msg.sender, address(this), amount);
     }

diff --git a/contracts/prototypes/evm/ZetaConnectorNewBase.sol b/contracts/prototypes/evm/ZetaConnectorNewBase.sol
@@ -2,28 +2,37 @@
 pragma solidity 0.8.7;
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
-import "./IGatewayEVM.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@openzeppelin/contracts/security/ReentrancyGuard.sol";
 
-abstract contract ZetaConnectorNewBase is ReentrancyGuard {
+import "./IGatewayEVM.sol";
+import "./IZetaConnector.sol";
+
+abstract contract ZetaConnectorNewBase is IZetaConnectorEvents, ReentrancyGuard {
     using SafeERC20 for IERC20;
 
     error ZeroAddress();
+    error InvalidSender();
 
     IGatewayEVM public immutable gateway;
     address public immutable zetaToken;
+    address public tssAddress;
 
-    event Withdraw(address indexed to, uint256 amount);
-    event WithdrawAndCall(address indexed to, uint256 amount, bytes data);
-    event WithdrawAndRevert(address indexed to, uint256 amount, bytes data);
+    // @dev Only TSS address allowed modifier.
+    modifier onlyTSS() {
+        if (msg.sender != tssAddress) {
+            revert InvalidSender();
+        }
+        _;
+    }
 
-    constructor(address _gateway, address _zetaToken) {
-        if (_gateway == address(0) || _zetaToken == address(0)) {
+    constructor(address _gateway, address _zetaToken, address _tssAddress) {
+        if (_gateway == address(0) || _zetaToken == address(0) || _tssAddress == address(0)) {
             revert ZeroAddress();
         }
         gateway = IGatewayEVM(_gateway);
         zetaToken = _zetaToken;
+        tssAddress = _tssAddress;
     }
 
     function withdraw(address to, uint256 amount, bytes32 internalSendHash) external virtual;

diff --git a/contracts/prototypes/evm/ZetaConnectorNonNative.sol b/contracts/prototypes/evm/ZetaConnectorNonNative.sol
@@ -6,18 +6,18 @@ import "./IZetaNonEthNew.sol";
 import "@openzeppelin/contracts/token/ERC20/extensions/ERC20Burnable.sol";
 
 contract ZetaConnectorNonNative is ZetaConnectorNewBase {
-    constructor(address _gateway, address _zetaToken)
-        ZetaConnectorNewBase(_gateway, _zetaToken)
+    constructor(address _gateway, address _zetaToken, address _tssAddress)
+        ZetaConnectorNewBase(_gateway, _zetaToken, _tssAddress)
     {}
 
     // @dev withdraw is called by TSS address, it mints zetaToken to the destination address
-    function withdraw(address to, uint256 amount, bytes32 internalSendHash) external override nonReentrant {
+    function withdraw(address to, uint256 amount, bytes32 internalSendHash) external override nonReentrant onlyTSS {
         IZetaNonEthNew(zetaToken).mint(to, amount, internalSendHash);
         emit Withdraw(to, amount);
     }
 
     // @dev withdrawAndCall is called by TSS address, it mints zetaToken and calls a contract
-    function withdrawAndCall(address to, uint256 amount, bytes calldata data, bytes32 internalSendHash) external override nonReentrant {
+    function withdrawAndCall(address to, uint256 amount, bytes calldata data, bytes32 internalSendHash) external override nonReentrant onlyTSS {
         // Mint zetaToken to the Gateway contract
         IZetaNonEthNew(zetaToken).mint(address(gateway), amount, internalSendHash);
 
@@ -28,7 +28,7 @@ contract ZetaConnectorNonNative is ZetaConnectorNewBase {
     }
 
     // @dev withdrawAndRevert is called by TSS address, it mints zetaToken to the gateway and calls onRevert on a contract
-    function withdrawAndRevert(address to, uint256 amount, bytes calldata data, bytes32 internalSendHash) external override nonReentrant {
+    function withdrawAndRevert(address to, uint256 amount, bytes calldata data, bytes32 internalSendHash) external override nonReentrant onlyTSS {
         // Mint zetaToken to the Gateway contract
         IZetaNonEthNew(zetaToken).mint(address(gateway), amount, internalSendHash);