fix: close imbalance check (#606)

* fix: close imbalance check

* fix: close imbalance check

Author: beeb

src/UsdnProtocol/libraries/UsdnProtocolActionsUtilsLibrary.sol

@@ -202,17 +202,9 @@ library UsdnProtocolActionsUtilsLibrary {
         data_.tempPositionValue =
             _assetToRemove(balanceLong, data_.lastPrice, liqPriceWithoutPenalty, data_.totalExpoToClose);
 
-        uint128 priceAfterFees =
-            (data_.lastPrice - data_.lastPrice * s._positionFeeBps / Constants.BPS_DIVISOR).toUint128();
-
-        uint256 posValueAfterFees =
-            _assetToRemove(balanceLong, priceAfterFees, liqPriceWithoutPenalty, data_.totalExpoToClose);
-
-        // we perform the imbalance check with the position value after fees
-        // the position value after fees is smaller than the position value before fees so the subtraction is safe
-        _checkImbalanceLimitClose(
-            s, data_.totalExpoToClose, posValueAfterFees, data_.tempPositionValue - posValueAfterFees
-        );
+        // we perform the imbalance check with the full position value subtracted from the long side, which is
+        // representative of the state of the balances after this initiate action
+        _checkImbalanceLimitClose(s, data_.totalExpoToClose, data_.tempPositionValue);
     }
 
     /**
@@ -341,15 +333,12 @@ library UsdnProtocolActionsUtilsLibrary {
      * the close limit on the vault side, otherwise revert
      * @param s The storage of the protocol
      * @param posTotalExpoToClose The total expo to remove position
-     * @param posValueToCloseAfterFees The value to remove from the position after the fees are applied
-     * @param fees The fees applied to the position, going to the vault
+     * @param posValueToClose The value to remove from the position (and the long balance)
      */
-    function _checkImbalanceLimitClose(
-        Types.Storage storage s,
-        uint256 posTotalExpoToClose,
-        uint256 posValueToCloseAfterFees,
-        uint256 fees
-    ) internal view {
+    function _checkImbalanceLimitClose(Types.Storage storage s, uint256 posTotalExpoToClose, uint256 posValueToClose)
+        internal
+        view
+    {
         int256 closeExpoImbalanceLimitBps;
         if (msg.sender == address(s._rebalancer)) {
             closeExpoImbalanceLimitBps = s._rebalancerCloseExpoImbalanceLimitBps;
@@ -362,9 +351,9 @@ library UsdnProtocolActionsUtilsLibrary {
             return;
         }
 
-        int256 newLongBalance = s._balanceLong.toInt256().safeSub(posValueToCloseAfterFees.toInt256());
+        int256 newLongBalance = s._balanceLong.toInt256().safeSub(posValueToClose.toInt256());
         uint256 newTotalExpo = s._totalExpo - posTotalExpoToClose;
-        int256 currentVaultExpo = s._balanceVault.toInt256().safeAdd(s._pendingBalanceVault + fees.toInt256());
+        int256 currentVaultExpo = s._balanceVault.toInt256().safeAdd(s._pendingBalanceVault);
 
         int256 imbalanceBps = Utils._calcImbalanceCloseBps(currentVaultExpo, newLongBalance, newTotalExpo);
 

test/unit/UsdnProtocol/Actions/_ImbalanceLimitClose.t.sol

@@ -28,7 +28,7 @@ contract TestImbalanceLimitClose is UsdnProtocolBaseFixture {
      */
     function test_checkImbalanceLimitClose() public view {
         (, uint256 longAmount, uint256 totalExpoValueToLimit) = _getCloseLimitValues(false);
-        protocol.i_checkImbalanceLimitClose(totalExpoValueToLimit / 2, longAmount, 0);
+        protocol.i_checkImbalanceLimitClose(totalExpoValueToLimit / 2, longAmount);
     }
 
     /**
@@ -40,7 +40,7 @@ contract TestImbalanceLimitClose is UsdnProtocolBaseFixture {
      */
     function test_checkImbalanceLimitCloseOnLimit() public view {
         (, uint256 longAmount, uint256 totalExpoValueToLimit) = _getCloseLimitValues(false);
-        protocol.i_checkImbalanceLimitClose(totalExpoValueToLimit - 1, longAmount, 0);
+        protocol.i_checkImbalanceLimitClose(totalExpoValueToLimit - 1, longAmount);
     }
 
     /**
@@ -55,7 +55,7 @@ contract TestImbalanceLimitClose is UsdnProtocolBaseFixture {
         vm.expectRevert(
             abi.encodeWithSelector(IUsdnProtocolErrors.UsdnProtocolImbalanceLimitReached.selector, closeLimitBps)
         );
-        protocol.i_checkImbalanceLimitClose(totalExpoValueToLimit, longAmount, 0);
+        protocol.i_checkImbalanceLimitClose(totalExpoValueToLimit, longAmount);
     }
 
     /**
@@ -75,7 +75,7 @@ contract TestImbalanceLimitClose is UsdnProtocolBaseFixture {
             abi.encodeWithSelector(IUsdnProtocolErrors.UsdnProtocolImbalanceLimitReached.selector, closeLimitBps)
         );
 
-        protocol.i_checkImbalanceLimitClose(totalExpoValueToLimit, longAmount, 0);
+        protocol.i_checkImbalanceLimitClose(totalExpoValueToLimit, longAmount);
     }
 
     /**
@@ -90,7 +90,7 @@ contract TestImbalanceLimitClose is UsdnProtocolBaseFixture {
         // disable close limit
         protocol.setExpoImbalanceLimits(200, 200, 600, 0, 0, 0);
 
-        protocol.i_checkImbalanceLimitClose(totalExpoValueToLimit + 1, longAmount, 0);
+        protocol.i_checkImbalanceLimitClose(totalExpoValueToLimit + 1, longAmount);
     }
 
     /**
@@ -149,7 +149,7 @@ contract TestImbalanceLimitClose is UsdnProtocolBaseFixture {
         vm.expectRevert(
             abi.encodeWithSelector(IUsdnProtocolErrors.UsdnProtocolImbalanceLimitReached.selector, type(int256).max)
         );
-        protocol.i_checkImbalanceLimitClose(0, 0, 0);
+        protocol.i_checkImbalanceLimitClose(0, 0);
     }
 
     /**
@@ -174,7 +174,7 @@ contract TestImbalanceLimitClose is UsdnProtocolBaseFixture {
                 IUsdnProtocolErrors.UsdnProtocolImbalanceLimitReached.selector, uint256(expectedImbalance)
             )
         );
-        protocol.i_checkImbalanceLimitClose(totalExpoValueToLimit, longAmount, 0);
+        protocol.i_checkImbalanceLimitClose(totalExpoValueToLimit, longAmount);
     }
 
     /**

test/unit/UsdnProtocol/Actions/fuzzing/_ImbalanceLimitClose.fuzzing.t.sol

@@ -54,6 +54,6 @@ contract TestImbalanceLimitCloseFuzzing is UsdnProtocolBaseFixture {
             );
         }
 
-        protocol.i_checkImbalanceLimitClose(totalExpoToRemove, closeAmount, 0);
+        protocol.i_checkImbalanceLimitClose(totalExpoToRemove, closeAmount);
     }
 }

test/unit/UsdnProtocol/utils/Handler.sol

@@ -404,11 +404,8 @@ contract UsdnProtocolHandler is UsdnProtocolImpl, Test {
         Long._checkImbalanceLimitOpen(s, openTotalExpoValue, openCollatValue);
     }
 
-    function i_checkImbalanceLimitClose(uint256 posTotalExpoToClose, uint256 posValueToClose, uint256 fees)
-        external
-        view
-    {
-        ActionsUtils._checkImbalanceLimitClose(s, posTotalExpoToClose, posValueToClose, fees);
+    function i_checkImbalanceLimitClose(uint256 posTotalExpoToClose, uint256 posValueToClose) external view {
+        ActionsUtils._checkImbalanceLimitClose(s, posTotalExpoToClose, posValueToClose);
     }
 
     function i_getLeverage(uint128 price, uint128 liqPrice) external pure returns (uint256) {

test/utils/IUsdnProtocolHandler.sol

@@ -202,9 +202,7 @@ interface IUsdnProtocolHandler is IUsdnProtocol {
 
     function i_checkImbalanceLimitOpen(uint256 openTotalExpoValue, uint256 openCollatValue) external view;
 
-    function i_checkImbalanceLimitClose(uint256 posTotalExpoToClose, uint256 posValueToCloseAfterFees, uint256 fees)
-        external
-        view;
+    function i_checkImbalanceLimitClose(uint256 posTotalExpoToClose, uint256 posValueToClose) external view;
 
     function i_getLeverage(uint128 price, uint128 liqPrice) external pure returns (uint256);