From 16bf664f4224d1767b86261ae24958d133e4d76f Mon Sep 17 00:00:00 2001
From: 0xtekgrinder <0xtekgrinder@protonmail.com>
Date: Wed, 17 Jul 2024 11:49:05 +0200
Subject: [PATCH] feat: swap vests assets and check no outgoing strategy assets

---
 contracts/BaseStrategy.sol | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/contracts/BaseStrategy.sol b/contracts/BaseStrategy.sol
index 0a63b44..ff8ff3e 100644
--- a/contracts/BaseStrategy.sol
+++ b/contracts/BaseStrategy.sol
@@ -525,23 +525,24 @@ abstract contract BaseStrategy is ERC4626, AccessControl {
      * @custom:requires KEEPER_ROLE
      */
     function swap(address[] calldata tokens, bytes[] calldata callDatas) public onlyRole(KEEPER_ROLE) {
-        address localAsset = asset();
-        address localStrategyAsset = STRATEGY_ASSET;
-        uint256 assetBalance = IERC20(localAsset).balanceOf(address(this));
-        uint256 strategyAssetBalance = IERC20(localStrategyAsset).balanceOf(address(this));
+        address _asset = asset();
+        address _strategyAsset = STRATEGY_ASSET;
+
+        uint256 assetBalance = IERC20(_asset).balanceOf(address(this));
+        uint256 strategyAssetBalance = IERC20(_strategyAsset).balanceOf(address(this));
 
         _swap(tokens, callDatas);
 
-        uint256 newStrategyAssetBalance = IERC20(localStrategyAsset).balanceOf(address(this));
+        uint256 newAssetBalance = IERC20(_asset).balanceOf(address(this));
 
-        if (
-            IERC20(localAsset).balanceOf(address(this)) < assetBalance || newStrategyAssetBalance < strategyAssetBalance
-        ) {
+        _handleUserGain(newAssetBalance);
+        _afterDeposit(newAssetBalance);
+
+        uint256 newStrategyAssetBalance = IERC20(_strategyAsset).balanceOf(address(this));
+
+        if (newStrategyAssetBalance < strategyAssetBalance) {
             revert OutgoingAssets();
         }
-
-        _handleUserGain(newStrategyAssetBalance);
-        _afterDeposit(newStrategyAssetBalance);
     }
 
     /**