fix(dex): solve token left in dex contract issue

foundry.toml

@@ -4,7 +4,7 @@ optimizer = false
 fs_permissions = [{ access = "read-write", path = "./"}]
 
 [profile.optimized]
-optimizer = false
+optimizer = true
 optimizer_runs = 10000
 fs_permissions = [{ access = "read-write", path = "./"}]
 

src/core/UniversalLiquidatorRegistry.sol

@@ -28,7 +28,7 @@ contract UniversalLiquidatorRegistry is Ownable, IUniversalLiquidatorRegistry, U
                     && paths[_intermediateTokens[idx]][_buyToken].dex != bytes32(0)
             ) {
                 // found the intermediateToken and intermediateDex
-                DataTypes.SwapInfo[] memory retPaths = new DataTypes.SwapInfo[](
+                DataTypes.SwapInfo[] memory retPaths = new DataTypes.SwapInfo[](
                     2
                 );
                 retPaths[0] = DataTypes.SwapInfo(
@@ -46,7 +46,7 @@ contract UniversalLiquidatorRegistry is Ownable, IUniversalLiquidatorRegistry, U
         revert Errors.PathsNotExist();
     }
 
-    function setPath(bytes32 _dex, address[] memory _paths) external override onlyOwner {
+    function setPath(bytes32 _dex, address[] calldata _paths) external override onlyOwner {
         // dex should exist
         if (!_dexExists(_dex)) revert Errors.DexDoesNotExist();
         // path could also be an empty array
@@ -56,7 +56,7 @@ contract UniversalLiquidatorRegistry is Ownable, IUniversalLiquidatorRegistry, U
         paths[_paths[0]][_paths[_paths.length - 1]] = DataTypes.PathInfo(_dex, _paths);
     }
 
-    function setIntermediateToken(address[] memory _token) public override onlyOwner {
+    function setIntermediateToken(address[] calldata _token) public override onlyOwner {
         _intermediateTokens = _token;
     }
 

src/core/dexes/BalancerDex.sol

@@ -5,6 +5,7 @@ pragma solidity 0.8.17;
 import "openzeppelin/access/Ownable.sol";
 import "openzeppelin/token/ERC20/utils/SafeERC20.sol";
 import "openzeppelin/token/ERC20/IERC20.sol";
+import "./BasicDex.sol";
 
 // interfaces
 import "../../interfaces/ILiquidityDex.sol";
@@ -16,12 +17,13 @@ import "../../libraries/Addresses.sol";
 // constants and types
 import {BalancerDexStorage} from "../storage/BalancerDex.sol";
 
-contract BalancerDex is Ownable, ILiquidityDex, BalancerDexStorage {
+contract BalancerDex is Ownable, BasicDex, ILiquidityDex, BalancerDexStorage {
     using SafeERC20 for IERC20;
 
-    function doSwap(uint256 _sellAmount, uint256 _minBuyAmount, address _receiver, address[] memory _path)
+    function doSwap(uint256 _sellAmount, uint256 _minBuyAmount, address _receiver, address[] calldata _path)
         external
         override
+        afterSwapCheck(_path[0], _path[_path.length - 1])
         returns (uint256)
     {
         address sellToken = _path[0];
@@ -75,5 +77,7 @@ contract BalancerDex is Ownable, ILiquidityDex, BalancerDexStorage {
         return _poolIds[_token0][_token1];
     }
 
-    receive() external payable {}
+    function tokenWithdraw(address _token, uint256 _amount) public override onlyOwner {
+        super.tokenWithdraw(_token, _amount);
+    }
 }

src/core/dexes/BancorV2Dex.sol

@@ -5,6 +5,7 @@ pragma solidity 0.8.17;
 import "openzeppelin/access/Ownable.sol";
 import "openzeppelin/token/ERC20/utils/SafeERC20.sol";
 import "openzeppelin/token/ERC20/IERC20.sol";
+import "./BasicDex.sol";
 
 // interfaces
 import "../../interfaces/ILiquidityDex.sol";
@@ -18,39 +19,44 @@ import "../../libraries/Addresses.sol";
 // constants and types
 import {BancorV2DexStorage} from "../storage/BancorV2Dex.sol";
 
-contract BancorV2Dex is Ownable, ILiquidityDex, BancorV2DexStorage {
+contract BancorV2Dex is Ownable, BasicDex, ILiquidityDex, BancorV2DexStorage {
     using SafeERC20 for IERC20;
 
     // BancorDex's doSwap doesn't expect any address in the path to be bancorETH
     // they have to be regular tokens.
-    function doSwap(uint256 _sellAmount, uint256 _minBuyAmount, address _receiver, address[] memory _path)
-        public
+    function doSwap(uint256 _sellAmount, uint256 _minBuyAmount, address _receiver, address[] calldata _path)
+        external
         override
-        returns (uint256)
+        afterSwapCheck(_path[0], _path[_path.length - 1])
+        returns (uint256 outTokenReturned)
     {
+        uint256 sellAmount = _sellAmount;
+        uint256 minBuyAmount = _minBuyAmount;
+        address receiver = _receiver;
+
         address sellToken = _path[0];
         address buyToken = _path[_path.length - 1];
         address finalReceiver = _receiver;
         address network = getBancorNetworkContract();
 
         if (sellToken == Addresses._WETH) {
-            WETH9(Addresses._WETH).withdraw(_sellAmount);
+            WETH9(Addresses._WETH).withdraw(sellAmount);
             sellToken = Addresses._BANCOR_ETH;
         } else {
-            IERC20(sellToken).safeIncreaseAllowance(network, _sellAmount);
+            IERC20(sellToken).safeIncreaseAllowance(network, sellAmount);
         }
 
         if (buyToken == Addresses._WETH) {
             // we will be receiving eth here, and wrap it back to WETH
             buyToken = Addresses._BANCOR_ETH;
-            _receiver = address(this);
+            receiver = address(this);
         }
 
         address[] memory actualPath = IBancorNetwork(network).conversionPath(sellToken, buyToken);
 
-        uint256 outTokenReturned = IBancorNetwork(network).convertByPath{
-            value: sellToken == Addresses._BANCOR_ETH ? _sellAmount : 0
-        }(actualPath, _sellAmount, _minBuyAmount, payable(_receiver), _affiliateAccount, _affiliateFee);
+        outTokenReturned = IBancorNetwork(network).convertByPath{value: sellToken == Addresses._BANCOR_ETH ? sellAmount : 0}(
+            actualPath, sellAmount, minBuyAmount, payable(receiver), _affiliateAccount, _affiliateFee
+        );
 
         // If buyToken is bancorEth, then this contract has received ETH after the swap.
         // ETH should be wrapped back to WETH
@@ -60,7 +66,7 @@ contract BancorV2Dex is Ownable, ILiquidityDex, BancorV2DexStorage {
             IERC20(Addresses._WETH).safeTransfer(finalReceiver, outTokenReturned);
         }
 
-        return outTokenReturned;
+        //return outTokenReturned;
     }
 
     function configure(address _newAffiliateAccount, uint256 _newAffiliateFee) external onlyOwner {
@@ -80,5 +86,7 @@ contract BancorV2Dex is Ownable, ILiquidityDex, BancorV2DexStorage {
         return _affiliateFee;
     }
 
-    receive() external payable {}
+    function tokenWithdraw(address _token, uint256 _amount) public override onlyOwner {
+        super.tokenWithdraw(_token, _amount);
+    }
 }

src/core/dexes/BasicDex.sol

@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.17;
+
+// imported contracts and libraries
+import "openzeppelin/token/ERC20/IERC20.sol";
+
+// libraries
+import "../../libraries/Errors.sol";
+
+abstract contract BasicDex {
+    modifier afterSwapCheck(address _sellToken, address _buyToken) {
+        _;
+
+        if (IERC20(_sellToken).balanceOf(address(this)) > 0) {
+            revert Errors.InvalidBalance(_sellToken);
+        }
+        if (IERC20(_buyToken).balanceOf(address(this)) > 0) {
+            revert Errors.InvalidBalance(_buyToken);
+        }
+    }
+
+    function tokenWithdraw(address _token, uint256 _amount) public virtual {
+        (bool success,) = _token.call(abi.encodeWithSignature("transfer(address,uint256)", msg.sender, _amount));
+        require(success, "BasicDex: transfer failed");
+    }
+
+    receive() external payable {}
+}

src/core/dexes/CurveDex.sol

@@ -5,6 +5,7 @@ pragma solidity 0.8.17;
 import "openzeppelin/access/Ownable.sol";
 import "openzeppelin/token/ERC20/utils/SafeERC20.sol";
 import "openzeppelin/token/ERC20/IERC20.sol";
+import "./BasicDex.sol";
 
 // interfaces
 import "../../interfaces/ILiquidityDex.sol";
@@ -16,12 +17,13 @@ import "../../libraries/Addresses.sol";
 // constants and types
 import {CurveDexStorage} from "../storage/CurveDex.sol";
 
-contract CurveDex is Ownable, ILiquidityDex, CurveDexStorage {
+contract CurveDex is Ownable, BasicDex, ILiquidityDex, CurveDexStorage {
     using SafeERC20 for IERC20;
 
-    function doSwap(uint256 _sellAmount, uint256 _minBuyAmount, address _receiver, address[] memory _path)
+    function doSwap(uint256 _sellAmount, uint256 _minBuyAmount, address _receiver, address[] calldata _path)
         external
         override
+        afterSwapCheck(_path[0], _path[_path.length - 1])
         returns (uint256 receiveAmt)
     {
         uint256 sellAmount = _sellAmount;
@@ -61,5 +63,7 @@ contract CurveDex is Ownable, ILiquidityDex, CurveDexStorage {
         return _pool[_token0][_token1];
     }
 
-    receive() external payable {}
+    function tokenWithdraw(address _token, uint256 _amount) public override onlyOwner {
+        super.tokenWithdraw(_token, _amount);
+    }
 }

src/core/dexes/PancakeV3Dex.sol

@@ -5,6 +5,7 @@ pragma solidity 0.8.17;
 import "openzeppelin/access/Ownable.sol";
 import "openzeppelin/token/ERC20/utils/SafeERC20.sol";
 import "openzeppelin/token/ERC20/IERC20.sol";
+import "./BasicDex.sol";
 
 // interfaces
 import "../../interfaces/ILiquidityDex.sol";
@@ -16,12 +17,13 @@ import "../../libraries/Addresses.sol";
 // constants and types
 import {PancakeV3DexStorage} from "../storage/PancakeV3Dex.sol";
 
-contract PancakeV3Dex is Ownable, ILiquidityDex, PancakeV3DexStorage {
+contract PancakeV3Dex is Ownable, BasicDex, ILiquidityDex, PancakeV3DexStorage {
     using SafeERC20 for IERC20;
 
-    function doSwap(uint256 _sellAmount, uint256 _minBuyAmount, address _receiver, address[] memory _path)
+    function doSwap(uint256 _sellAmount, uint256 _minBuyAmount, address _receiver, address[] calldata _path)
         external
         override
+        afterSwapCheck(_path[0], _path[_path.length - 1])
         returns (uint256)
     {
         address sellToken = _path[0];
@@ -60,5 +62,7 @@ contract PancakeV3Dex is Ownable, ILiquidityDex, PancakeV3DexStorage {
         _pairFee[_token1][_token0] = _fee;
     }
 
-    receive() external payable {}
+    function tokenWithdraw(address _token, uint256 _amount) public override onlyOwner {
+        super.tokenWithdraw(_token, _amount);
+    }
 }

src/core/dexes/UniBasedDex.sol

@@ -5,6 +5,7 @@ pragma solidity 0.8.17;
 import "openzeppelin/access/Ownable.sol";
 import "openzeppelin/token/ERC20/utils/SafeERC20.sol";
 import "openzeppelin/token/ERC20/IERC20.sol";
+import "./BasicDex.sol";
 
 // interfaces
 import "../../interfaces/ILiquidityDex.sol";
@@ -14,16 +15,17 @@ import "../../interfaces/uniswap/v2/IUniswapV2Factory.sol";
 // constants and types
 import {BaseDexStorage} from "../storage/BaseDex.sol";
 
-contract UniBasedDex is Ownable, ILiquidityDex, BaseDexStorage {
+contract UniBasedDex is Ownable, BasicDex, ILiquidityDex, BaseDexStorage {
     using SafeERC20 for IERC20;
 
     constructor(address _initRouter) {
         _router = _initRouter;
     }
 
-    function doSwap(uint256 _sellAmount, uint256 _minBuyAmount, address _receiver, address[] memory _path)
+    function doSwap(uint256 _sellAmount, uint256 _minBuyAmount, address _receiver, address[] calldata _path)
         external
         override
+        afterSwapCheck(_path[0], _path[_path.length - 1])
         returns (uint256)
     {
         address sellToken = _path[0];
@@ -40,5 +42,7 @@ contract UniBasedDex is Ownable, ILiquidityDex, BaseDexStorage {
         return _router;
     }
 
-    receive() external payable {}
+    function tokenWithdraw(address _token, uint256 _amount) public override onlyOwner {
+        super.tokenWithdraw(_token, _amount);
+    }
 }

src/core/dexes/UniV3Dex.sol

@@ -5,6 +5,7 @@ pragma solidity 0.8.17;
 import "openzeppelin/access/Ownable.sol";
 import "openzeppelin/token/ERC20/utils/SafeERC20.sol";
 import "openzeppelin/token/ERC20/IERC20.sol";
+import "./BasicDex.sol";
 
 // interfaces
 import "../../interfaces/ILiquidityDex.sol";
@@ -16,12 +17,13 @@ import "../../libraries/Addresses.sol";
 // constants and types
 import {UniswapV3DexStorage} from "../storage/UniswapV3Dex.sol";
 
-contract UniV3Dex is Ownable, ILiquidityDex, UniswapV3DexStorage {
+contract UniV3Dex is Ownable, BasicDex, ILiquidityDex, UniswapV3DexStorage {
     using SafeERC20 for IERC20;
 
-    function doSwap(uint256 _sellAmount, uint256 _minBuyAmount, address _receiver, address[] memory _path)
+    function doSwap(uint256 _sellAmount, uint256 _minBuyAmount, address _receiver, address[] calldata _path)
         external
         override
+        afterSwapCheck(_path[0], _path[_path.length - 1])
         returns (uint256)
     {
         address sellToken = _path[0];
@@ -60,5 +62,7 @@ contract UniV3Dex is Ownable, ILiquidityDex, UniswapV3DexStorage {
         _pairFee[_token1][_token0] = _fee;
     }
 
-    receive() external payable {}
+    function tokenWithdraw(address _token, uint256 _amount) public override onlyOwner {
+        super.tokenWithdraw(_token, _amount);
+    }
 }

src/interfaces/ILiquidityDex.sol

@@ -2,7 +2,7 @@
 pragma solidity 0.8.17;
 
 interface ILiquidityDex {
-    function doSwap(uint256 _sellAmount, uint256 _minBuyAmount, address _receiver, address[] memory _path)
+    function doSwap(uint256 _sellAmount, uint256 _minBuyAmount, address _receiver, address[] calldata _path)
         external
         returns (uint256);
 }

src/interfaces/IUniversalLiquidatorRegistry.sol

@@ -7,9 +7,9 @@ import "../libraries/DataTypes.sol";
 interface IUniversalLiquidatorRegistry {
     function getPath(address _sellToken, address _buyToken) external view returns (DataTypes.SwapInfo[] memory);
 
-    function setPath(bytes32 _dex, address[] memory _paths) external;
+    function setPath(bytes32 _dex, address[] calldata _paths) external;
 
-    function setIntermediateToken(address[] memory _token) external;
+    function setIntermediateToken(address[] calldata _token) external;
 
     function addDex(bytes32 _name, address _address) external;
 

src/libraries/Errors.sol

@@ -9,4 +9,6 @@ library Errors {
     error PathsNotExist();
     // UniversalLiquidator errors
     error InvalidAddress();
+    // Dex errors
+    error InvalidBalance(address token);
 }

test/UniversalLiquidator.t.sol

@@ -8,7 +8,7 @@ import "openzeppelin/token/ERC20/IERC20.sol";
 import "../src/libraries/Errors.sol";
 
 // import test base and helpers.
-import {console2, AdvancedFixture} from "./AdvancedFixture.t.sol";
+import {console2, AdvancedFixture} from "./AdvancedFixture.sol";
 
 contract UniversalLiquidatorTest is AdvancedFixture {
     address _farmer;

test/UniversalLiquidatorRegistry.t.sol

@@ -9,7 +9,7 @@ import "../src/core/dexes/UniV3Dex.sol";
 import "../src/libraries/Errors.sol";
 
 // import test base and helpers.
-import {AdvancedFixture} from "./AdvancedFixture.t.sol";
+import {AdvancedFixture} from "./AdvancedFixture.sol";
 
 contract UniversalLiquidatorRegistryTest is AdvancedFixture {
     function testSetPath() public {

test/dexes/BalancerDex.t.sol

@@ -5,7 +5,7 @@ pragma solidity 0.8.17;
 import "../../src/core/dexes/BalancerDex.sol";
 
 // import test base and helpers.
-import {AdvancedFixture} from "../AdvancedFixture.t.sol";
+import {AdvancedFixture} from "../AdvancedFixture.sol";
 
 contract BalancerDexTest is AdvancedFixture {
     function testSetPoolId() public {

test/dexes/BancorV2Dex.t.sol

@@ -5,7 +5,7 @@ pragma solidity 0.8.17;
 import "../../src/core/dexes/BancorV2Dex.sol";
 
 // import test base and helpers.
-import {AdvancedFixture} from "../AdvancedFixture.t.sol";
+import {AdvancedFixture} from "../AdvancedFixture.sol";
 
 contract BancorV2DexTest is AdvancedFixture {
     function testSetPoolFee() public {

test/dexes/CurveDex.t.sol

@@ -5,7 +5,7 @@ pragma solidity 0.8.17;
 import "../../src/core/dexes/CurveDex.sol";
 
 // import test base and helpers.
-import {AdvancedFixture} from "../AdvancedFixture.t.sol";
+import {AdvancedFixture} from "../AdvancedFixture.sol";
 
 contract CurveDexTest is AdvancedFixture {
     function testSetPoolId() public {