feat: Max withdraw on amount exceeding supplied (#727)

Co-authored-by: Alexandru Niculae <[email protected]>
Co-authored-by: YBM <[email protected]>
Co-authored-by: Francesc Armengol Carrasco <[email protected]>
Co-authored-by: DhairyaSethi <[email protected]>

Author: 5 people

snapshots/Spoke.Operations.json

@@ -22,7 +22,7 @@
   "usingAsCollateral: 0 borrows, enable": "54024",
   "usingAsCollateral: 1 borrow, disable": "153486",
   "usingAsCollateral: 1 borrow, enable": "24786",
-  "withdraw: 0 borrows, full": "133631",
-  "withdraw: 0 borrows, partial": "135734",
-  "withdraw: 1 borrow, partial": "252502"
+  "withdraw: 0 borrows, full": "133493",
+  "withdraw: 0 borrows, partial": "135596",
+  "withdraw: 1 borrow, partial": "252364"
 }

src/contracts/Spoke.sol

@@ -218,14 +218,11 @@ contract Spoke is ISpoke, Multicall, AccessManaged, EIP712 {
   ) external onlyPositionManager(onBehalfOf) {
     DataTypes.Reserve storage reserve = _reserves[reserveId];
     DataTypes.UserPosition storage userPosition = _userPositions[onBehalfOf][reserveId];
+    _validateWithdraw(reserve);
     uint256 assetId = reserve.assetId;
     IHubBase hub = reserve.hub;
 
-    // If uint256.max is passed, withdraw all user's supplied assets
-    if (amount == type(uint256).max) {
-      amount = hub.previewRemoveByShares(assetId, userPosition.suppliedShares);
-    }
-    _validateWithdraw(reserve, userPosition, amount);
+    amount = MathUtils.min(amount, hub.previewRemoveByShares(assetId, userPosition.suppliedShares));
 
     uint256 withdrawnShares = hub.remove(assetId, amount, msg.sender);
 
@@ -648,18 +645,9 @@ contract Spoke is ISpoke, Multicall, AccessManaged, EIP712 {
     require(!reserve.frozen, ReserveFrozen());
   }
 
-  function _validateWithdraw(
-    DataTypes.Reserve storage reserve,
-    DataTypes.UserPosition storage userPosition,
-    uint256 amount
-  ) internal view {
+  function _validateWithdraw(DataTypes.Reserve storage reserve) internal view {
     require(address(reserve.hub) != address(0), ReserveNotListed());
     require(!reserve.paused, ReservePaused());
-    uint256 suppliedAmount = reserve.hub.previewRemoveByShares(
-      reserve.assetId,
-      userPosition.suppliedShares
-    );
-    require(amount <= suppliedAmount, InsufficientSupply(suppliedAmount));
   }
 
   function _validateBorrow(DataTypes.Reserve storage reserve) internal view {

src/contracts/TreasurySpoke.sol

@@ -5,7 +5,8 @@ pragma solidity ^0.8.0;
 import {Ownable} from 'src/dependencies/openzeppelin/Ownable.sol';
 import {SafeERC20} from 'src/dependencies/openzeppelin/SafeERC20.sol';
 import {IERC20} from 'src/dependencies/openzeppelin/IERC20.sol';
-import {IHub} from 'src/interfaces/IHub.sol';
+import {MathUtils} from 'src/libraries/math/MathUtils.sol';
+import {IHubBase} from 'src/interfaces/IHubBase.sol';
 import {ITreasurySpoke, ISpokeBase} from 'src/interfaces/ITreasurySpoke.sol';
 
 /**
@@ -19,7 +20,7 @@ contract TreasurySpoke is ITreasurySpoke, Ownable {
   using SafeERC20 for IERC20;
 
   /// @inheritdoc ITreasurySpoke
-  IHub public immutable HUB;
+  IHubBase public immutable HUB;
 
   /**
    * @dev Constructor
@@ -29,7 +30,7 @@ contract TreasurySpoke is ITreasurySpoke, Ownable {
   constructor(address owner_, address hub_) Ownable(owner_) {
     require(hub_ != address(0), InvalidAddress());
 
-    HUB = IHub(hub_);
+    HUB = IHubBase(hub_);
   }
 
   /// @inheritdoc ITreasurySpoke
@@ -39,11 +40,8 @@ contract TreasurySpoke is ITreasurySpoke, Ownable {
 
   /// @inheritdoc ITreasurySpoke
   function withdraw(uint256 reserveId, uint256 amount, address) external onlyOwner {
-    // If uint256.max is passed, withdraw all supplied assets
-    if (amount == type(uint256).max) {
-      amount = HUB.getSpokeAddedAssets(reserveId, address(this));
-    }
-
+    // If amount to withdraw is greater than total supplied, withdraw all supplied assets
+    amount = MathUtils.min(amount, HUB.getSpokeAddedAssets(reserveId, address(this)));
     HUB.remove(reserveId, amount, msg.sender);
   }
 

src/interfaces/ISpoke.sol

@@ -105,7 +105,6 @@ interface ISpoke is ISpokeBase, IMulticall, IAccessManaged {
   error AssetNotListed();
   error ReserveExists();
   error ReserveNotListed();
-  error InsufficientSupply(uint256 supply);
   error ReserveNotBorrowable();
   error ReservePaused();
   error ReserveFrozen();

src/interfaces/ITreasurySpoke.sol

@@ -2,7 +2,7 @@
 // Copyright (c) 2025 Aave Labs
 pragma solidity ^0.8.0;
 
-import {IHub} from 'src/interfaces/IHub.sol';
+import {IHubBase} from 'src/interfaces/IHubBase.sol';
 import {ISpokeBase} from 'src/interfaces/ISpokeBase.sol';
 
 /**
@@ -61,5 +61,5 @@ interface ITreasurySpoke is ISpokeBase {
    * @notice Returns the address of the associated Hub.
    * @return The address of the Hub.
    */
-  function HUB() external view returns (IHub);
+  function HUB() external view returns (IHubBase);
 }

tests/Base.t.sol

@@ -1254,7 +1254,7 @@ abstract contract Base is Test {
     return IERC20(underlying);
   }
 
-  function getWithdrawalLimit(
+  function getTotalWithdrawable(
     ISpoke spoke,
     uint256 reserveId,
     address user

tests/unit/Spoke/Spoke.Multicall.t.sol

@@ -233,11 +233,11 @@ contract SpokeMulticall is SpokeBase {
   function test_multicall_forwards_first_revert() public {
     bytes[] memory calls = new bytes[](3);
     calls[0] = abi.encodeCall(ISpokeBase.supply, (_daiReserveId(spoke1), 120e18, alice));
-    calls[1] = abi.encodeCall(ISpokeBase.withdraw, (_daiReserveId(spoke1), 121e18, alice));
+    calls[1] = abi.encodeCall(ISpokeBase.withdraw, (_daiReserveId(spoke1), 0, alice));
     calls[2] = abi.encodeCall(ISpoke.setUsingAsCollateral, (_daiReserveId(spoke1), true, alice));
 
     vm.prank(alice);
-    vm.expectRevert(abi.encodeWithSelector(ISpoke.InsufficientSupply.selector, 120e18));
+    vm.expectRevert(IHub.InvalidAmount.selector);
     spoke1.multicall(calls);
   }
 }

tests/unit/Spoke/Spoke.Withdraw.Validation.t.sol

@@ -26,13 +26,15 @@ contract SpokeWithdrawValidationTest is SpokeBase {
     spoke1.withdraw(reserveId, amount, bob);
   }
 
-  function test_withdraw_revertsWith_InsufficientSupply_zero_supplied() public {
+  /// @dev Test passes 1 as amount with no supplied assets.
+  /// @dev The spoke contract changes the calling amount to the total user supplied, but since it's zero, it reverts.
+  function test_withdraw_revertsWith_InvalidAmount_zero_supplied() public {
     uint256 reserveId = _daiReserveId(spoke1);
     uint256 amount = 1;
 
     assertEq(spoke1.getUserSuppliedAssets(reserveId, alice), 0);
 
-    vm.expectRevert(abi.encodeWithSelector(ISpoke.InsufficientSupply.selector, 0));
+    vm.expectRevert(IHub.InvalidAmount.selector);
     vm.prank(alice);
     spoke1.withdraw(reserveId, amount, alice);
   }
@@ -43,42 +45,11 @@ contract SpokeWithdrawValidationTest is SpokeBase {
 
     assertEq(spoke1.getUserSuppliedAssets(reserveId, alice), 0);
 
-    vm.expectRevert(abi.encodeWithSelector(ISpoke.InsufficientSupply.selector, 0));
+    vm.expectRevert(IHub.InvalidAmount.selector);
     vm.prank(alice);
     spoke1.withdraw(reserveId, amount, alice);
   }
 
-  // Withdraw reverts when there is not enough avaulable liquidity
-  function test_withdraw_revertsWith_InsufficientSupply_with_supply() public {
-    uint256 amount = 100e18;
-    uint256 reserveId = _daiReserveId(spoke1);
-
-    // User spoke supply
-    Utils.supply({
-      spoke: spoke1,
-      reserveId: reserveId,
-      caller: alice,
-      amount: amount,
-      onBehalfOf: alice
-    });
-
-    uint256 withdrawalLimit = getWithdrawalLimit(spoke1, reserveId, alice);
-    assertGt(withdrawalLimit, 0);
-
-    vm.expectRevert(abi.encodeWithSelector(ISpoke.InsufficientSupply.selector, withdrawalLimit));
-    vm.prank(alice);
-    spoke1.withdraw(reserveId, withdrawalLimit + 1, alice);
-
-    // skip time but no index increase with no borrow
-    skip(365 days);
-    // withdrawal limit remains constant
-    assertEq(withdrawalLimit, getWithdrawalLimit(spoke1, reserveId, alice));
-
-    vm.expectRevert(abi.encodeWithSelector(ISpoke.InsufficientSupply.selector, withdrawalLimit));
-    vm.prank(alice);
-    spoke1.withdraw(reserveId, withdrawalLimit + 1, alice);
-  }
-
   // Withdrawal limit increases due to debt interest, but still cannot withdraw more than available liquidity
   function test_withdraw_revertsWith_InsufficientSupply_with_debt() public {
     uint256 supplyAmount = 100e18;
@@ -103,18 +74,23 @@ contract SpokeWithdrawValidationTest is SpokeBase {
       onBehalfOf: alice
     });
 
-    vm.expectRevert(abi.encodeWithSelector(ISpoke.InsufficientSupply.selector, supplyAmount));
+    vm.expectRevert(
+      abi.encodeWithSelector(IHub.InsufficientLiquidity.selector, supplyAmount - borrowAmount)
+    );
     vm.prank(alice);
     spoke1.withdraw({reserveId: reserveId, amount: supplyAmount + 1, onBehalfOf: alice});
 
     // accrue interest
     skip(365 days);
 
-    uint256 newWithdrawalLimit = getWithdrawalLimit(spoke1, reserveId, alice);
+    uint256 newWithdrawalLimit = getTotalWithdrawable(spoke1, reserveId, alice);
     // newWithdrawalLimit with accrued interest should be greater than supplyAmount
     assertGt(newWithdrawalLimit, supplyAmount);
 
-    vm.expectRevert(abi.encodeWithSelector(ISpoke.InsufficientSupply.selector, newWithdrawalLimit));
+    // Interest added on both sides, so can ignore
+    vm.expectRevert(
+      abi.encodeWithSelector(IHub.InsufficientLiquidity.selector, supplyAmount - borrowAmount)
+    );
     vm.prank(alice);
     spoke1.withdraw({reserveId: reserveId, amount: newWithdrawalLimit + 1, onBehalfOf: alice});
   }
@@ -152,18 +128,23 @@ contract SpokeWithdrawValidationTest is SpokeBase {
       onBehalfOf: alice
     });
 
-    vm.expectRevert(abi.encodeWithSelector(ISpoke.InsufficientSupply.selector, supplyAmount));
+    vm.expectRevert(
+      abi.encodeWithSelector(IHub.InsufficientLiquidity.selector, supplyAmount - borrowAmount)
+    );
     vm.prank(alice);
     spoke1.withdraw({reserveId: reserveId, amount: supplyAmount + 1, onBehalfOf: alice});
 
     // debt accrues
     skip(skipTime);
 
-    uint256 newWithdrawalLimit = getWithdrawalLimit(spoke1, reserveId, alice);
+    uint256 newWithdrawalLimit = getTotalWithdrawable(spoke1, reserveId, alice);
     // newWithdrawalLimit with accrued interest should be greater than supplyAmount
     vm.assume(newWithdrawalLimit > supplyAmount);
 
-    vm.expectRevert(abi.encodeWithSelector(ISpoke.InsufficientSupply.selector, newWithdrawalLimit));
+    // Interest added on both sides, so can ignore
+    vm.expectRevert(
+      abi.encodeWithSelector(IHub.InsufficientLiquidity.selector, supplyAmount - borrowAmount)
+    );
     vm.prank(alice);
     spoke1.withdraw({reserveId: reserveId, amount: newWithdrawalLimit + 1, onBehalfOf: alice});
   }

tests/unit/Spoke/Spoke.Withdraw.t.sol

@@ -168,6 +168,35 @@ contract SpokeWithdrawTest is SpokeBase {
     _checkSupplyRateIncreasing(addExRate, getAddExRate(daiAssetId), 'after withdraw');
   }
 
+  function test_withdraw_fuzz_all_greater_than_supplied(uint256 supplyAmount) public {
+    supplyAmount = bound(supplyAmount, 1, MAX_SUPPLY_AMOUNT);
+    Utils.supply({
+      spoke: spoke1,
+      reserveId: _daiReserveId(spoke1),
+      caller: bob,
+      amount: supplyAmount,
+      onBehalfOf: bob
+    });
+
+    _checkSuppliedAmounts(
+      daiAssetId,
+      _daiReserveId(spoke1),
+      spoke1,
+      bob,
+      supplyAmount,
+      'after supply'
+    );
+
+    uint256 addExRate = getAddExRate(daiAssetId);
+
+    // Withdraw all supplied assets
+    vm.prank(bob);
+    spoke1.withdraw(_daiReserveId(spoke1), supplyAmount + 1, bob);
+
+    _checkSuppliedAmounts(daiAssetId, _daiReserveId(spoke1), spoke1, bob, 0, 'after withdraw');
+    _checkSupplyRateIncreasing(addExRate, getAddExRate(daiAssetId), 'after withdraw');
+  }
+
   function test_withdraw_fuzz_all_with_interest(uint256 supplyAmount, uint256 borrowAmount) public {
     supplyAmount = bound(supplyAmount, 2, MAX_SUPPLY_AMOUNT);
     borrowAmount = bound(borrowAmount, 1, supplyAmount / 2);
@@ -838,4 +867,38 @@ contract SpokeWithdrawTest is SpokeBase {
     assertGe(hub1.convertToAddedAssets(daiAssetId, MAX_SUPPLY_AMOUNT), supplyExchangeRatio);
     assertGe(hub1.convertToDrawnAssets(daiAssetId, MAX_SUPPLY_AMOUNT), debtExchangeRatio);
   }
+
+  /// @dev Withdraw exceeding supplied amount withdraws everything
+  function test_withdraw_max_greater_than_supplied() public {
+    uint256 amount = 100e18;
+    uint256 reserveId = _daiReserveId(spoke1);
+
+    // User spoke supply
+    Utils.supply({
+      spoke: spoke1,
+      reserveId: reserveId,
+      caller: alice,
+      amount: amount,
+      onBehalfOf: alice
+    });
+
+    uint256 withdrawable = getTotalWithdrawable(spoke1, reserveId, alice);
+    assertGt(withdrawable, 0);
+
+    uint256 addExRateBefore = getAddExRate(daiAssetId);
+
+    // skip time but no index increase with no borrow
+    skip(365 days);
+    // withdrawable remains constant
+    assertEq(withdrawable, getTotalWithdrawable(spoke1, reserveId, alice));
+
+    vm.prank(alice);
+    spoke1.withdraw(reserveId, withdrawable + 1, alice);
+
+    assertEq(getTotalWithdrawable(spoke1, reserveId, alice), 0);
+    _checkSuppliedAmounts(daiAssetId, reserveId, spoke1, alice, 0, 'after withdraw');
+
+    // Check supply rate monotonically increasing after withdraw
+    _checkSupplyRateIncreasing(addExRateBefore, getAddExRate(daiAssetId), 'after withdraw');
+  }
 }