Merge pull request #34 from tanim0la/feature/eng-19-strategy-module-d…

…evelopment

Addressed a possible attack vector

src/curators/engine.sol

@@ -1,10 +1,15 @@
 // SPDX-License-Identifier: GNU
 pragma solidity ^0.8.20;
 
-import "./interface/IStrategy.sol";
+import "@openzeppelin/contracts/access/Ownable2Step.sol";
 import {ERC4626, ERC20, IERC20} from "@openzeppelin/contracts/token/ERC20/extensions/ERC4626.sol";
 
-contract Engine is ERC4626 {
+import "./interface/IStrategy.sol";
+
+contract Engine is ERC4626, Ownable2Step {
+    // assetOut => true/false
+    mapping(address => bool) public verifyAssetOut;
+
     /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
     /*                          EVENTS                            */
     /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
@@ -29,19 +34,46 @@ contract Engine is ERC4626 {
     /*                           ERROR                            */
     /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
 
+    /// @dev A step execution fails.
+    error ExecuteStepFailed(string reason);
+
+    /// @dev The action type is invalid.
     error InvalidActionType();
 
-    constructor() ERC4626(IERC20(address(this))) ERC20("LIQUID", "LLP") {}
+    constructor() ERC4626(IERC20(address(this))) Ownable(msg.sender) ERC20("LIQUID", "LLP") {}
+
+    /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+    /*                     ONLY OWNER FUNCTIONS                   */
+    /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
+
+    /// @notice Withdraw user asset
+    function toggleAssetOut(address _assetOut) external onlyOwner {
+        verifyAssetOut[_assetOut] = !verifyAssetOut[_assetOut];
+    }
+
+    /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+    /*                       PUBLIC FUNCTIONS                     */
+    /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
 
     function join(bytes32 _strategyId, address _strategyModule, uint256[] memory _amounts) public {
         // Fetch the strategy
         ILiquidStrategy.Strategy memory strategy = ILiquidStrategy(_strategyModule).getStrategy(_strategyId);
 
         // Validate strategy - not necessary single we validate the steps before strategy creation
 
-        // Transfer initial deposit(s) from caller
+        // AssetsIn
+
         uint256 initialAssetsInLength = strategy.steps[0].assetsIn.length;
 
+        uint256[] memory assetsInBalBefore = new uint256[](initialAssetsInLength);
+
+        for (uint256 i; i < initialAssetsInLength; i++) {
+            address asset = strategy.steps[0].assetsIn[i];
+
+            assetsInBalBefore[i] = IERC20(asset).balanceOf(_strategyModule);
+        }
+
+        // Transfer initial deposit(s) from caller
         for (uint256 i; i < initialAssetsInLength; i++) {
             address asset = strategy.steps[0].assetsIn[i];
             // approve `this` as spender in client first
@@ -89,9 +121,6 @@ contract Engine is ERC4626 {
                 address[] memory underlyingTokens,
                 uint256[] memory underlyingAmounts
             ) {
-                // Verify result
-                // require(_verifyResult(shareAmount, step.assetOut, _strategyModule), "Invalid shareAmount");
-
                 // update user info
                 ILiquidStrategy(_strategyModule).updateUserStats(
                     _strategyId,
@@ -106,7 +135,23 @@ contract Engine is ERC4626 {
                     i
                 );
             } catch Error(string memory reason) {
-                revert(string(abi.encodePacked("Step ", i, " failed: ", reason)));
+                revert ExecuteStepFailed(string(abi.encodePacked("Step ", i, " failed: ", reason)));
+            }
+        }
+
+        // Return the remaining deposit balance
+        for (uint256 i; i < initialAssetsInLength; i++) {
+            address asset = strategy.steps[0].assetsIn[i];
+
+            uint256 remainingDeposit = IERC20(asset).balanceOf(_strategyModule) - assetsInBalBefore[i];
+
+            if (remainingDeposit > 0) {
+                ILiquidStrategy(_strategyModule).updateUserTokenBalanceEngine(
+                    _strategyId, msg.sender, asset, remainingDeposit, 1
+                );
+                _amounts[i] -= remainingDeposit;
+
+                ILiquidStrategy(_strategyModule).transferToken(msg.sender, asset, remainingDeposit);
             }
         }
 
@@ -129,7 +174,7 @@ contract Engine is ERC4626 {
     }
 
     function exit(bytes32 _strategyId, address _strategyModule) public {
-        // check and burn user's liquid share token (also prevent re-enterancy)
+        // check and burn user's liquid share token
 
         // Fetch the strategy
         ILiquidStrategy.Step[] memory steps = ILiquidStrategy(_strategyModule).getStrategy(_strategyId).steps;
@@ -168,7 +213,7 @@ contract Engine is ERC4626 {
             try IConnector(step.connector).execute(
                 actionType, assetsIn, assetOut, i - 1, 0, _strategyId, msg.sender, step.data
             ) returns (
-                address,
+                address protocol,
                 address[] memory,
                 uint256[] memory,
                 address assetOut,
@@ -180,27 +225,29 @@ contract Engine is ERC4626 {
 
                 // Transfer token to user
                 if (i == 1) {
+                    ILiquidStrategy.ShareBalance memory userShareBalance = ILiquidStrategy(_strategyModule)
+                        .getUserShareBalance(_strategyId, msg.sender, protocol, step.assetOut, 0);
+
+                    // zero free for now
+                    uint256 _fee = 0;
+
+                    // update strategy stats
+                    ILiquidStrategy(_strategyModule).updateStrategyStats(
+                        _strategyId, step.assetsIn, userShareBalance.underlyingAmounts, msg.sender, _fee, 1
+                    );
+
                     // todo: get all the assetout then send
                     for (uint256 j; j < step.assetsIn.length; j++) {
                         IConnector(step.connector).withdrawAsset(_strategyId, msg.sender, step.assetsIn[j]);
                     }
+
+                    ILiquidStrategy(_strategyModule).deleteUserPosition(_strategyId, msg.sender, step.assetsIn);
                 }
             } catch Error(string memory reason) {
-                revert(string(abi.encodePacked("Step ", i - 1, " failed: ", reason)));
+                revert ExecuteStepFailed(string(abi.encodePacked("Step ", i, " failed: ", reason)));
             }
         }
 
-        ILiquidStrategy.AssetBalance memory userAssetBalance =
-            ILiquidStrategy(_strategyModule).getUserAssetBalance(_strategyId, msg.sender, steps[0].assetsIn, 0);
-
-        // zero free for now
-        uint256 _fee = 0;
-
-        // update strategy stats
-        ILiquidStrategy(_strategyModule).updateStrategyStats(
-            _strategyId, steps[0].assetsIn, userAssetBalance.amounts, msg.sender, _fee, 1
-        );
-
         // update user strategy stats
         ILiquidStrategy(_strategyModule).updateUserStrategy(_strategyId, msg.sender, 1);
 
@@ -210,12 +257,4 @@ contract Engine is ERC4626 {
         // Emits Exit event
         emit Exit(_strategyId, msg.sender);
     }
-
-    function _verifyResult(uint256 _shareAmount, address _assetOut, address _strategyModule)
-        internal
-        view
-        returns (bool)
-    {
-        return ERC4626(_assetOut).balanceOf(_strategyModule) == _shareAmount;
-    }
 }

src/curators/interface/IEngine.sol

@@ -2,5 +2,7 @@
 pragma solidity ^0.8.13;
 
 interface IEngine {
+    function verifyAssetOut(address _assetOut) external returns (bool);
     function join(bytes32 _strategyId, address _strategyModule, uint256[] memory _amounts) external;
+    function exit(bytes32 _strategyId, address _strategyModule) external;
 }

src/curators/interface/IStrategy.sol

@@ -66,6 +66,8 @@ interface ILiquidStrategy is IConnector {
 
     function transferToken(address _token, uint256 _amount) external returns (bool);
 
+    function transferToken(address _user, address _token, uint256 _amount) external returns (bool);
+
     /**
      * @dev set value to true if a user joins a strategy, else, false
      * @param _strategyId  strategy identity
@@ -115,6 +117,32 @@ interface ILiquidStrategy is IConnector {
         uint256 _indicator
     ) external;
 
+    /**
+     * @dev Update user token balance
+     * @param _strategyId unique identifier of the strategy.
+     * @param _user address of the user whose balance is being updated.
+     * @param _token address of the token.
+     * @param _amount the amount of token to add or sub.
+     * @param _indicator determines the operation:
+     *                   0 to add,
+     *                   any other value to sub.
+     */
+    function updateUserTokenBalanceEngine(
+        bytes32 _strategyId,
+        address _user,
+        address _token,
+        uint256 _amount,
+        uint256 _indicator
+    ) external;
+
+    /**
+     * @dev Delete user strategy position
+     * @param _strategyId unique identifier of the strategy to update.
+     * @param _user address of the user whose strategy is being updated.
+     * @param _assets address of the user whose strategy is being updated.
+     */
+    function deleteUserPosition(bytes32 _strategyId, address _user, address[] memory _assets) external;
+
     /**
      * @dev Get strategy by strategy id
      * @param _strategyId strategy identity

src/curators/oracle.sol

@@ -1,12 +1,50 @@
 // SPDX-License-Identifier: GNU
 pragma solidity ^0.8.20;
 
+import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+
 import "./interface/IOracle.sol";
+import "../protocols/common/constant.sol";
+
+contract Oracle is Constants {
+    /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+    /*                       CUSTOM ERRORS                        */
+    /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
 
-contract Oracle {
+    /// @dev The Chainlink sequencer is down.
     error SequencerDown();
+
+    /// @dev The grace period has NOT passed after the sequencer is back up.
     error GracePeriodNotOver();
 
+    /// @dev The datafeed is address zero.
+    error InvalidDataFeed();
+
+    /// @dev The price of a token in USD is zero.
+    error InvalidPriceInUsd();
+
+    /**
+     * @notice function to get the price of token in USD
+     * @param _amount token amount to calculate in USD
+     * @param _token the addreas of the token
+     */
+    function getPriceInUSD(uint256 _amount, address _token) public view returns (uint256) {
+        address _dataFeed;
+
+        if (_token == CBBTC) _dataFeed = CBBTC_USD;
+        if (_token == ETH) _dataFeed = ETH_USD;
+        if (_token == USDC) _dataFeed = USDC_USD;
+        if (_token == DAI) _dataFeed = DAI_USD;
+
+        if (_dataFeed == address(0)) revert InvalidDataFeed();
+
+        int256 priceInUSD = getLatestAnswer(SEQUENCER_UPTIME_FEED, _dataFeed);
+
+        if (priceInUSD <= int256(0)) revert InvalidPriceInUsd();
+
+        return (uint256(priceInUSD) * _amount) / 10 ** ERC20(_token).decimals(); // returns usd value scaled to 8 decimal
+    }
+
     /**
      * @notice function to get the price of token in USD
      * @param _sequencerUptimeFeed uptime feed address on base network