fix: manageDebt mask refactorred

Gearbox-protocol · Apr 28, 2023 · 269e78c · 269e78c
1 parent df38db4
commit 269e78c
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 16 deletions.
diff --git a/contracts/credit/CreditFacadeV3.sol b/contracts/credit/CreditFacadeV3.sol
@@ -748,10 +748,10 @@ contract CreditFacadeV3 is ICreditFacade, ACLNonReentrantTrait, BalanceHelperTra
     function _manageDebt(
         address creditAccount,
         bytes memory callData,
-        uint256 enabledTokensMask,
+        uint256 _enabledTokensMask,
         address borrower,
         ManageDebtAction action
-    ) internal returns (uint256) {
+    ) internal returns (uint256 enabledTokensMask) {
         // It is forbidden to take new debt if increaseDebtForbidden mode is enabled
         if (params.isIncreaseDebtForbidden) {
             revert IncreaseDebtForbiddenException();
@@ -762,21 +762,20 @@ contract CreditFacadeV3 is ICreditFacade, ACLNonReentrantTrait, BalanceHelperTra
         // Checks that the borrowed amount does not violate the per block limit
         _checkAndUpdateBorrowedBlockLimit(amount); // F:[FA-18A]
 
+        uint256 newBorrowedAmount;
         // Requests the Credit Manager to borrow additional funds from the pool
-        (uint256 newBorrowedAmount, bool underlyingBalanceIsZero) =
-            creditManager.manageDebt(creditAccount, amount, enabledTokensMask, action); // F:[FA-17]
+        (newBorrowedAmount, enabledTokensMask) =
+            creditManager.manageDebt(creditAccount, amount, _enabledTokensMask, action); // F:[FA-17]
 
         // Checks that the new total borrowed amount is within bounds
         _revertIfOutOfBorrowedLimits(newBorrowedAmount); // F:[FA-18B]
 
         // Emits event
         if (action == ManageDebtAction.INCREASE_DEBT) {
             emit IncreaseBorrowedAmount(borrower, amount); // F:[FA-17]
-            return enabledTokensMask | UNDERLYING_TOKEN_MASK;
+        } else {
+            emit DecreaseBorrowedAmount(borrower, amount); // F:[FA-19]
         }
-
-        emit DecreaseBorrowedAmount(borrower, amount); // F:[FA-19]
-        return underlyingBalanceIsZero ? enabledTokensMask & (~UNDERLYING_TOKEN_MASK) : enabledTokensMask;
     }
 
     function _addCollateral(address creditAccount, bytes memory callData, address borrower)

diff --git a/contracts/credit/CreditManagerV3.sol b/contracts/credit/CreditManagerV3.sol
@@ -10,7 +10,7 @@ import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol
 import {SafeCast} from "@openzeppelin/contracts/utils/math/SafeCast.sol";
 
 // LIBS & TRAITS
-import {BitMask} from "../libraries/BitMask.sol";
+import {UNDERLYING_TOKEN_MASK, BitMask} from "../libraries/BitMask.sol";
 import {ReentrancyGuard} from "@openzeppelin/contracts/security/ReentrancyGuard.sol";
 import {SanityCheckTrait} from "../traits/SanityCheckTrait.sol";
 import {BalanceHelperTrait} from "../traits/BalanceHelperTrait.sol";
@@ -395,11 +395,11 @@ contract CreditManagerV3 is ICreditManagerV3, SanityCheckTrait, ReentrancyGuard,
     /// @param amount Amount to increase / decrease the principal by
     /// @param action Increase/decrease bed debt
     /// @return newBorrowedAmount The new debt principal
-    function manageDebt(address creditAccount, uint256 amount, uint256 enableTokenMask, ManageDebtAction action)
+    function manageDebt(address creditAccount, uint256 amount, uint256 _enabledTokensMask, ManageDebtAction action)
         external
         nonReentrant
         creditFacadeOnly // F:[CM-2]
-        returns (uint256 newBorrowedAmount, bool underlyingBalanceIsZero)
+        returns (uint256 newBorrowedAmount, uint256 enabledTokensMask)
     {
         (uint256 borrowedAmount, uint256 cumulativeIndexAtOpen_RAY, uint256 cumulativeIndexNow_RAY) =
             _getCreditAccountParameters(creditAccount);
@@ -420,14 +420,15 @@ contract CreditManagerV3 is ICreditManagerV3, SanityCheckTrait, ReentrancyGuard,
 
             // Requests the pool to lend additional funds to the Credit Account
             IPoolService(pool).lendCreditAccount(amount, creditAccount); // F:[CM-20]
+            enabledTokensMask = _enabledTokensMask | UNDERLYING_TOKEN_MASK;
         } else {
             // Decrease
             uint256 amountRepaid = amount;
             uint256 amountProfit = 0;
 
             if (supportsQuotas) {
                 (amountRepaid, amountProfit) =
-                    _computeQuotasAmountDebtDecrease(creditAccount, amountRepaid, amountProfit, enableTokenMask);
+                    _computeQuotasAmountDebtDecrease(creditAccount, amountRepaid, amountProfit, _enabledTokensMask);
             }
 
             if (amountRepaid > 0) {
@@ -480,7 +481,10 @@ contract CreditManagerV3 is ICreditManagerV3, SanityCheckTrait, ReentrancyGuard,
 
             // Pays the amount back to the pool
             ICreditAccount(creditAccount).safeTransfer(underlying, pool, amount); // F:[CM-21]
-            underlyingBalanceIsZero = _balanceOf(creditAccount, underlying) <= 1;
+
+            enabledTokensMask = _balanceOf(underlying, creditAccount) <= 1
+                ? _enabledTokensMask & (~UNDERLYING_TOKEN_MASK)
+                : _enabledTokensMask;
 
             // TODO: delete after tests or write Invaraiant test
             require(borrowedAmount - newBorrowedAmount == amountRepaid, "Ooops, something was wring");

diff --git a/contracts/interfaces/ICreditManagerV3.sol b/contracts/interfaces/ICreditManagerV3.sol
@@ -118,11 +118,11 @@ interface ICreditManagerV3 is ICreditManagerV3Events, IVersion {
     ///
     /// @param creditAccount Address of the Credit Account to change debt for
     /// @param amount Amount to increase / decrease the principal by
-    /// @param  action Increase/decrease
+    /// @param action Increase/decrease
     /// @return newBorrowedAmount The new debt principal
-    function manageDebt(address creditAccount, uint256 amount, uint256 enabledTokensMask, ManageDebtAction action)
+    function manageDebt(address creditAccount, uint256 amount, uint256 _enabledTokensMask, ManageDebtAction action)
         external
-        returns (uint256 newBorrowedAmount, bool underlyingBalanceIsZero);
+        returns (uint256 newBorrowedAmount, uint256 enableTokenMask);
 
     /// @dev Adds collateral to borrower's credit account
     /// @param payer Address of the account which will be charged to provide additional collateral