chg: add extra check on the withdrawal amount from the assetManager

unioncredit · Jul 24, 2024 · 1e6c540 · 1e6c540
1 parent 0ce7036
commit 1e6c540
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 7 deletions.
diff --git a/contracts/asset/AssetManager.sol b/contracts/asset/AssetManager.sol
@@ -330,7 +330,7 @@ contract AssetManager is Controller, ReentrancyGuardUpgradeable, IAssetManager {
      *  @param token ERC20 token address
      *  @param account User address
      *  @param amount ERC20 token address
-     *  @return Withdraw amount
+     *  @return The difference between the amount withdrawn and the amount transferred to the caller
      */
     function withdraw(
         address token,

diff --git a/contracts/market/UToken.sol b/contracts/market/UToken.sol
@@ -633,14 +633,17 @@ contract UToken is IUToken, Controller, ERC20PermitUpgradeable, ReentrancyGuardU
         if (remaining > amount) revert WithdrawFailed();
         actualAmount -= decimalScaling(remaining, underlyingDecimal);
 
+        // Ensure the actual withdrawal amount is not less than the minimum borrow amount
+        if (actualAmount < _minBorrow) revert AmountLessMinBorrow();
+
         fee = calculatingFee(actualAmount);
         uint256 accountBorrowsNew = borrowedAmount + actualAmount + fee;
         uint256 totalBorrowsNew = _totalBorrows + actualAmount + fee;
         if (totalBorrowsNew > _debtCeiling) revert AmountExceedGlobalMax();
 
         // Update internal balances
-        accountBorrows[msg.sender].principal += actualAmount + fee;
-        uint256 newPrincipal = _getBorrowed(msg.sender);
+        uint256 newPrincipal = actualAmount + fee;
+        accountBorrows[msg.sender].principal += newPrincipal;
         accountBorrows[msg.sender].interest = accountBorrowsNew - newPrincipal;
         accountBorrows[msg.sender].interestIndex = borrowIndex;
         _totalBorrows = totalBorrowsNew;

diff --git a/slither.db.json b/slither.db.json
diff --git a/test/foundry/uToken/TestBorrowRepay.t.sol b/test/foundry/uToken/TestBorrowRepay.t.sol
@@ -46,9 +46,7 @@ contract TestBorrowRepay is TestUTokenBase {
 
     function testBorrowWhenNotEnough(uint256 borrowAmount) public {
         vm.assume(
-            borrowAmount >= MIN_BORROW &&
-                borrowAmount > UNIT &&
-                borrowAmount < MAX_BORROW - (MAX_BORROW * ORIGINATION_FEE) / 1e18
+            borrowAmount >= MIN_BORROW + UNIT && borrowAmount < MAX_BORROW - (MAX_BORROW * ORIGINATION_FEE) / 1e18
         );
 
         vm.startPrank(ALICE);
@@ -62,8 +60,10 @@ contract TestBorrowRepay is TestUTokenBase {
 
         uint256 borrowed = uToken.getBorrowed(ALICE);
         uint256 realBorrowAmount = borrowAmount - UNIT;
+
         // borrowed amount should only include origination fee
         uint256 fees = (ORIGINATION_FEE * realBorrowAmount) / 1e18;
+
         assertEq(borrowed, realBorrowAmount + fees);
     }
 
@@ -205,4 +205,19 @@ contract TestBorrowRepay is TestUTokenBase {
 
         vm.stopPrank();
     }
+
+    function testBorrowReturnLessThanMinBorrow() public {
+        uint256 borrowAmount = MIN_BORROW;
+
+        vm.startPrank(ALICE);
+        vm.mockCall(
+            address(assetManagerMock),
+            abi.encodeWithSelector(AssetManager.withdraw.selector, erc20Mock, ALICE, borrowAmount),
+            abi.encode(MIN_BORROW - 1)
+        );
+        vm.expectRevert(UToken.AmountLessMinBorrow.selector);
+        uToken.borrow(ALICE, borrowAmount);
+
+        vm.stopPrank();
+    }
 }