feat: added some info in the tokenManagerProxy, which will cause some gas saving and will allow for a more powerfull design in the future. (#170)

* added some info in the tokenManagerProxy, which will cause some gas saving and will allow for a more powerfull design in the future.

* add service approval on appropriate token managers

* address comment

* refactor imports

* update import

---------

Co-authored-by: Milap Sheth <[email protected]>

Author: Foivos

contracts/interfaces/IBaseTokenManager.sol

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+/**
+ * @title IBaseTokenManager
+ * @notice This contract is defines the base token manager interface implemented by all token managers.
+ */
+interface IBaseTokenManager {
+    /**
+     * @notice A function that returns the token id.
+     */
+    function interchainTokenId() external view returns (bytes32);
+
+    /**
+     * @notice A function that should return the address of the token.
+     * Must be overridden in the inheriting contract.
+     * @return address address of the token.
+     */
+    function tokenAddress() external view returns (address);
+
+    /**
+     * @notice A function that should return the implementation type of the token manager.
+     */
+    function implementationType() external pure returns (uint256);
+
+    /**
+     * @notice A function that should return the token address from the init params.
+     */
+    function getTokenAddressFromParams(bytes calldata params) external pure returns (address);
+}

contracts/interfaces/ITokenManager.sol

@@ -4,15 +4,15 @@ pragma solidity ^0.8.0;
 
 import { IImplementation } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/interfaces/IImplementation.sol';
 
-import { ITokenManagerType } from './ITokenManagerType.sol';
+import { IBaseTokenManager } from './IBaseTokenManager.sol';
 import { IOperatable } from './IOperatable.sol';
 import { IFlowLimit } from './IFlowLimit.sol';
 
 /**
  * @title ITokenManager
  * @notice This contract is responsible for handling tokens before initiating a cross chain token transfer, or after receiving one.
  */
-interface ITokenManager is ITokenManagerType, IOperatable, IFlowLimit, IImplementation {
+interface ITokenManager is IBaseTokenManager, IOperatable, IFlowLimit, IImplementation {
     error TokenLinkerZeroAddress();
     error NotService(address caller);
     error TakeTokenFailed();
@@ -22,23 +22,6 @@ interface ITokenManager is ITokenManagerType, IOperatable, IFlowLimit, IImplemen
     error AlreadyFlowLimiter(address flowLimiter);
     error NotFlowLimiter(address flowLimiter);
 
-    /**
-     * @notice A function that returns the token id.
-     */
-    function interchainTokenId() external view returns (bytes32);
-
-    /**
-     * @notice A function that should return the address of the token.
-     * Must be overridden in the inheriting contract.
-     * @return address address of the token.
-     */
-    function tokenAddress() external view returns (address);
-
-    /**
-     * @notice A function that should return the implementation type of the token manager.
-     */
-    function implementationType() external pure returns (uint256);
-
     /**
      * @notice Calls the service to initiate a cross-chain transfer after taking the appropriate amount of tokens from the user.
      * @param destinationChain the name of the chain to send tokens to.

contracts/interfaces/ITokenManagerProxy.sol

@@ -22,4 +22,14 @@ interface ITokenManagerProxy is IProxy {
      * @return bytes32 The interchain token ID of the token manager.
      */
     function interchainTokenId() external view returns (bytes32);
+
+    /**
+     * @notice Returns token address that this token manager manages.
+     */
+    function tokenAddress() external view returns (address);
+
+    /**
+     * @notice Returns implementation type and token address.
+     */
+    function getImplementationTypeAndTokenAddress() external view returns (uint256, address);
 }

contracts/proxies/TokenManagerProxy.sol

@@ -5,6 +5,7 @@ pragma solidity ^0.8.0;
 import { IProxy } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/interfaces/IProxy.sol';
 import { BaseProxy } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/upgradable/BaseProxy.sol';
 
+import { IBaseTokenManager } from '../interfaces/IBaseTokenManager.sol';
 import { ITokenManagerProxy } from '../interfaces/ITokenManagerProxy.sol';
 import { ITokenManagerImplementation } from '../interfaces/ITokenManagerImplementation.sol';
 
@@ -19,6 +20,7 @@ contract TokenManagerProxy is BaseProxy, ITokenManagerProxy {
     address public immutable interchainTokenService;
     uint256 public immutable implementationType;
     bytes32 public immutable interchainTokenId;
+    address public immutable tokenAddress;
 
     /**
      * @notice Constructs the TokenManagerProxy contract.
@@ -39,6 +41,8 @@ contract TokenManagerProxy is BaseProxy, ITokenManagerProxy {
 
         (bool success, ) = implementation_.delegatecall(abi.encodeWithSelector(IProxy.setup.selector, params));
         if (!success) revert SetupFailed();
+
+        tokenAddress = IBaseTokenManager(implementation_).getTokenAddressFromParams(params);
     }
 
     /**
@@ -49,6 +53,14 @@ contract TokenManagerProxy is BaseProxy, ITokenManagerProxy {
         return CONTRACT_ID;
     }
 
+    /**
+     * @notice Returns implementation type and token address.
+     */
+    function getImplementationTypeAndTokenAddress() external view returns (uint256 implementationType_, address tokenAddress_) {
+        implementationType_ = implementationType;
+        tokenAddress_ = tokenAddress;
+    }
+
     /**
      * @notice Returns the address of the current implementation.
      * @return implementation_ The address of the current implementation.

contracts/test/HardCodedConstantsTest.sol

@@ -14,7 +14,6 @@ contract TestTokenManager is TokenManagerLiquidityPool {
     string public constant NAME = 'TestTokenManager';
 
     constructor(address interchainTokenService_) TokenManagerLiquidityPool(interchainTokenService_) {
-        if (TOKEN_ADDRESS_SLOT != uint256(keccak256('token-address')) - 1) revert Invalid();
         if (LIQUIDITY_POOL_SLOT != uint256(keccak256('liquidity-pool-slot')) - 1) revert Invalid();
     }
 }

contracts/token-manager/TokenManager.sol

@@ -7,6 +7,7 @@ import { IImplementation } from '@axelar-network/axelar-gmp-sdk-solidity/contrac
 import { Implementation } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/upgradable/Implementation.sol';
 
 import { ITokenManager } from '../interfaces/ITokenManager.sol';
+import { ITokenManagerType } from '../interfaces/ITokenManagerType.sol';
 import { IInterchainTokenService } from '../interfaces/IInterchainTokenService.sol';
 
 import { Operatable } from '../utils/Operatable.sol';
@@ -16,16 +17,13 @@ import { FlowLimit } from '../utils/FlowLimit.sol';
  * @title The main functionality of TokenManagers.
  * @notice This contract is responsible for handling tokens before initiating a cross chain token transfer, or after receiving one.
  */
-abstract contract TokenManager is ITokenManager, Operatable, FlowLimit, Implementation {
+abstract contract TokenManager is ITokenManager, ITokenManagerType, Operatable, FlowLimit, Implementation {
     using AddressBytes for bytes;
 
     IInterchainTokenService public immutable interchainTokenService;
 
     bytes32 private constant CONTRACT_ID = keccak256('token-manager');
 
-    // uint256(keccak256('token-address')) - 1
-    uint256 internal constant TOKEN_ADDRESS_SLOT = 0xc4e632779a6a7838736dd7e5e6a0eadf171dd37dfb6230720e265576dfcf42ba;
-
     /**
      * @notice Constructs the TokenManager contract.
      * @param interchainTokenService_ The address of the interchain token service
@@ -48,7 +46,7 @@ abstract contract TokenManager is ITokenManager, Operatable, FlowLimit, Implemen
      * @dev A modifier that allows only the token to execute the function.
      */
     modifier onlyToken() {
-        if (msg.sender != tokenAddress()) revert NotToken(msg.sender);
+        if (msg.sender != this.tokenAddress()) revert NotToken(msg.sender);
         _;
     }
 
@@ -60,13 +58,12 @@ abstract contract TokenManager is ITokenManager, Operatable, FlowLimit, Implemen
     }
 
     /**
-     * @dev Reads the stored token address from the predetermined storage slot
+     * @dev Reads the token address from the proxy
      * @return tokenAddress_ The address of the token
      */
-    function tokenAddress() public view virtual returns (address tokenAddress_) {
-        assembly {
-            tokenAddress_ := sload(TOKEN_ADDRESS_SLOT)
-        }
+    function tokenAddress() external view virtual returns (address tokenAddress_) {
+        // Ask the proxy what the token address is.
+        tokenAddress_ = this.tokenAddress();
     }
 
     /**
@@ -78,6 +75,13 @@ abstract contract TokenManager is ITokenManager, Operatable, FlowLimit, Implemen
         return this.interchainTokenId();
     }
 
+    /**
+     * @notice A function that should return the token address from the setup params.
+     */
+    function getTokenAddressFromParams(bytes calldata params) external pure returns (address tokenAddress_) {
+        (, tokenAddress_) = abi.decode(params, (uint256, address));
+    }
+
     /**
      * @dev This function should only be called by the proxy, and only once from the proxy constructor
      * @param params the parameters to be used to initialize the TokenManager. The exact format depends
@@ -251,16 +255,6 @@ abstract contract TokenManager is ITokenManager, Operatable, FlowLimit, Implemen
         _setFlowLimit(flowLimit_, interchainTokenId());
     }
 
-    /**
-     * @dev Stores the token address in the predetermined storage slot
-     * @param tokenAddress_ The address of the token to store
-     */
-    function _setTokenAddress(address tokenAddress_) internal {
-        assembly {
-            sstore(TOKEN_ADDRESS_SLOT, tokenAddress_)
-        }
-    }
-
     /**
      * @notice Transfers tokens from a specific address to this contract.
      * Must be overridden in the inheriting contract.
@@ -284,5 +278,5 @@ abstract contract TokenManager is ITokenManager, Operatable, FlowLimit, Implemen
      * Must be overridden in the inheriting contract.
      * @param params The setup parameters
      */
-    function _setup(bytes calldata params) internal virtual;
+    function _setup(bytes calldata params) internal virtual {}
 }

contracts/token-manager/TokenManagerLiquidityPool.sol

@@ -3,7 +3,7 @@
 pragma solidity ^0.8.0;
 
 import { IERC20 } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/interfaces/IERC20.sol';
-import { SafeTokenTransferFrom } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/libs/SafeTransfer.sol';
+import { SafeTokenTransferFrom, SafeTokenCall } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/libs/SafeTransfer.sol';
 import { ReentrancyGuard } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/utils/ReentrancyGuard.sol';
 
 import { ITokenManagerLiquidityPool } from '../interfaces/ITokenManagerLiquidityPool.sol';
@@ -18,6 +18,7 @@ import { TokenManager } from './TokenManager.sol';
  */
 contract TokenManagerLiquidityPool is TokenManager, ReentrancyGuard, ITokenManagerLiquidityPool {
     using SafeTokenTransferFrom for IERC20;
+    using SafeTokenCall for IERC20;
 
     error NotSupported();
 
@@ -41,9 +42,10 @@ contract TokenManagerLiquidityPool is TokenManager, ReentrancyGuard, ITokenManag
      */
     function _setup(bytes calldata params_) internal override {
         // The first argument is reserved for the operator.
-        (, address tokenAddress_, address liquidityPool_) = abi.decode(params_, (bytes, address, address));
-        _setTokenAddress(tokenAddress_);
+        (, address tokenAddress_, address liquidityPool_) = abi.decode(params_, (uint256, address, address));
         _setLiquidityPool(liquidityPool_);
+
+        IERC20(tokenAddress_).safeCall(abi.encodeWithSelector(IERC20.approve.selector, interchainTokenService, type(uint256).max));
     }
 
     /**
@@ -81,7 +83,7 @@ contract TokenManagerLiquidityPool is TokenManager, ReentrancyGuard, ITokenManag
      * @return uint The actual amount of tokens transferred. This allows support for fee-on-transfer tokens.
      */
     function _takeToken(address from, uint256 amount) internal override noReEntrancy returns (uint256) {
-        IERC20 token = IERC20(tokenAddress());
+        IERC20 token = IERC20(this.tokenAddress());
         address liquidityPool_ = liquidityPool();
         uint256 balance = token.balanceOf(liquidityPool_);
 
@@ -101,7 +103,7 @@ contract TokenManagerLiquidityPool is TokenManager, ReentrancyGuard, ITokenManag
      * @return uint The actual amount of tokens transferred
      */
     function _giveToken(address to, uint256 amount) internal override noReEntrancy returns (uint256) {
-        IERC20 token = IERC20(tokenAddress());
+        IERC20 token = IERC20(this.tokenAddress());
         uint256 balance = token.balanceOf(to);
 
         // slither-disable-next-line arbitrary-send-erc20

contracts/token-manager/TokenManagerLockUnlock.sol

@@ -3,7 +3,7 @@
 pragma solidity ^0.8.0;
 
 import { IERC20 } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/interfaces/IERC20.sol';
-import { SafeTokenTransfer, SafeTokenTransferFrom } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/libs/SafeTransfer.sol';
+import { SafeTokenTransfer, SafeTokenTransferFrom, SafeTokenCall } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/libs/SafeTransfer.sol';
 
 import { ITokenManagerLockUnlock } from '..//interfaces/ITokenManagerLockUnlock.sol';
 import { TokenManager } from './TokenManager.sol';
@@ -17,6 +17,7 @@ import { TokenManager } from './TokenManager.sol';
 contract TokenManagerLockUnlock is TokenManager, ITokenManagerLockUnlock {
     using SafeTokenTransfer for IERC20;
     using SafeTokenTransferFrom for IERC20;
+    using SafeTokenCall for IERC20;
 
     /**
      * @dev Constructs an instance of TokenManagerLockUnlock. Calls the constructor
@@ -30,13 +31,14 @@ contract TokenManagerLockUnlock is TokenManager, ITokenManagerLockUnlock {
     }
 
     /**
-     * @dev Sets up the token address.
-     * @param params_ The setup parameters in bytes. Should be encoded with the token address.
+     * @dev Sets up the token address and liquidity pool address.
+     * @param params_ The setup parameters in bytes. Should be encoded with the token address and the liquidity pool address.
      */
     function _setup(bytes calldata params_) internal override {
         // The first argument is reserved for the operator.
-        (, address tokenAddress_) = abi.decode(params_, (bytes, address));
-        _setTokenAddress(tokenAddress_);
+        (, address tokenAddress_) = abi.decode(params_, (uint256, address));
+
+        IERC20(tokenAddress_).safeCall(abi.encodeWithSelector(IERC20.approve.selector, interchainTokenService, type(uint256).max));
     }
 
     /**
@@ -46,7 +48,7 @@ contract TokenManagerLockUnlock is TokenManager, ITokenManagerLockUnlock {
      * @return uint The actual amount of tokens transferred. This allows support for fee-on-transfer tokens.
      */
     function _takeToken(address from, uint256 amount) internal override returns (uint256) {
-        IERC20 token = IERC20(tokenAddress());
+        IERC20 token = IERC20(this.tokenAddress());
 
         token.safeTransferFrom(from, address(this), amount);
 
@@ -60,7 +62,7 @@ contract TokenManagerLockUnlock is TokenManager, ITokenManagerLockUnlock {
      * @return uint The actual amount of tokens transferred
      */
     function _giveToken(address to, uint256 amount) internal override returns (uint256) {
-        IERC20 token = IERC20(tokenAddress());
+        IERC20 token = IERC20(this.tokenAddress());
 
         token.safeTransfer(to, amount);
 

contracts/token-manager/TokenManagerLockUnlockFee.sol

@@ -3,7 +3,7 @@
 pragma solidity ^0.8.0;
 
 import { IERC20 } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/interfaces/IERC20.sol';
-import { SafeTokenTransferFrom, SafeTokenTransfer } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/libs/SafeTransfer.sol';
+import { SafeTokenTransferFrom, SafeTokenTransfer, SafeTokenCall } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/libs/SafeTransfer.sol';
 import { ReentrancyGuard } from '@axelar-network/axelar-gmp-sdk-solidity/contracts/utils/ReentrancyGuard.sol';
 
 import { ITokenManagerLockUnlock } from '../interfaces/ITokenManagerLockUnlock.sol';
@@ -18,6 +18,7 @@ import { TokenManager } from './TokenManager.sol';
 contract TokenManagerLockUnlockFee is TokenManager, ReentrancyGuard, ITokenManagerLockUnlock {
     using SafeTokenTransfer for IERC20;
     using SafeTokenTransferFrom for IERC20;
+    using SafeTokenCall for IERC20;
 
     /**
      * @dev Constructs an instance of TokenManagerLockUnlock. Calls the constructor
@@ -31,13 +32,14 @@ contract TokenManagerLockUnlockFee is TokenManager, ReentrancyGuard, ITokenManag
     }
 
     /**
-     * @dev Sets up the token address.
-     * @param params_ The setup parameters in bytes. Should be encoded with the token address.
+     * @dev Sets up the token address and liquidity pool address.
+     * @param params_ The setup parameters in bytes. Should be encoded with the token address and the liquidity pool address.
      */
     function _setup(bytes calldata params_) internal override {
         // The first argument is reserved for the operator.
-        (, address tokenAddress_) = abi.decode(params_, (bytes, address));
-        _setTokenAddress(tokenAddress_);
+        (, address tokenAddress_) = abi.decode(params_, (uint256, address));
+
+        IERC20(tokenAddress_).safeCall(abi.encodeWithSelector(IERC20.approve.selector, interchainTokenService, type(uint256).max));
     }
 
     /**
@@ -47,7 +49,7 @@ contract TokenManagerLockUnlockFee is TokenManager, ReentrancyGuard, ITokenManag
      * @return uint The actual amount of tokens transferred. This allows support for fee-on-transfer tokens.
      */
     function _takeToken(address from, uint256 amount) internal override noReEntrancy returns (uint256) {
-        IERC20 token = IERC20(tokenAddress());
+        IERC20 token = IERC20(this.tokenAddress());
         uint256 balanceBefore = token.balanceOf(address(this));
 
         token.safeTransferFrom(from, address(this), amount);
@@ -67,7 +69,7 @@ contract TokenManagerLockUnlockFee is TokenManager, ReentrancyGuard, ITokenManag
      * @return uint The actual amount of tokens transferred
      */
     function _giveToken(address to, uint256 amount) internal override noReEntrancy returns (uint256) {
-        IERC20 token = IERC20(tokenAddress());
+        IERC20 token = IERC20(this.tokenAddress());
         uint256 balanceBefore = token.balanceOf(to);
 
         token.safeTransfer(to, amount);