test: returnAsset attack

src/asset-management/AaveManager.sol

@@ -131,12 +131,12 @@ contract AaveManager is IAssetManager, Owned(msg.sender), ReentrancyGuard {
     {
         rShares = aAmount.mulDivUp(totalShares[aAaveToken], aAaveToken.balanceOf(address(this)));
 
-        uint256 lCurrentShares = shares[aPair][aToken];
-//
-//        // this is to prevent underflow as we round up in the previous division operation
-//        if (rShares > lCurrentShares) {
-//            rShares = lCurrentShares;
-//        }
+        // this is the vuln, where even if the shares are 0 this will still succeed
+        // uint256 lCurrentShares = shares[aPair][aToken];
+        // this is to prevent underflow as we round up in the previous division operation
+        // if (rShares > lCurrentShares) {
+        //    rShares = lCurrentShares;
+        // }
 
         shares[aPair][aToken] -= rShares;
         totalShares[aAaveToken] -= rShares;

test/__mocks/ReturnAssetExploit.sol

@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.0;
+
+import "src/interfaces/IAssetManager.sol";
+import { ReservoirPair } from "src/ReservoirPair.sol";
+import { SafeTransferLib } from "solady/utils/SafeTransferLib.sol";
+
+contract ReturnAssetExploit {
+    using SafeTransferLib for address;
+
+    IERC20 public token0;
+    IERC20 public token1;
+
+    constructor(ReservoirPair aPair) {
+        token0 = aPair.token0();
+        token1 = aPair.token1();
+    }
+
+    function adjustManagement(int256 token0Change, int256 token1Change) external {
+        address(token0).safeTransferFrom(msg.sender, address(this), uint256(-token0Change));
+        address(token1).safeTransferFrom(msg.sender, address(this), uint256(-token1Change));
+    }
+
+    function attack(IAssetManager aAssetManager) external {
+        aAssetManager.returnAsset(false, 123555);
+    }
+}

test/integration/Aave.t.sol

@@ -16,6 +16,7 @@ import { MathUtils } from "src/libraries/MathUtils.sol";
 import { AaveManager, IAssetManager } from "src/asset-management/AaveManager.sol";
 import { GenericFactory, IERC20 } from "src/GenericFactory.sol";
 import { IUSDC } from "test/interfaces/IUSDC.sol";
+import { ReturnAssetExploit } from "../__mocks/ReturnAssetExploit.sol";
 
 struct Network {
     string rpcUrl;
@@ -491,6 +492,8 @@ contract AaveIntegrationTest is BaseTest {
 
         // assert
         assertEq(_manager.shares(_pair, USDC), uint256(lAmountToManagePair));
+        console.log(_manager.getBalance(_pair, USDC));
+        console.log(lAaveTokenAmt2);
         assertTrue(MathUtils.within1(_manager.getBalance(_pair, USDC), lAaveTokenAmt2));
 
         uint256 lExpectedShares =
@@ -1037,6 +1040,7 @@ contract AaveIntegrationTest is BaseTest {
         uint256 lBalAfterTimePair = _manager.getBalance(_pair, USDC);
         uint256 lBalAfterTimeOther = _manager.getBalance(lOtherPair, USDC);
         uint256 lClaimed = _manager.claimRewardForMarket(lUSDCMarket, lWavax);
+        // dummy amount of proceeds from selling the rewards
         uint256 lAmtUSDC = 9_019_238;
         _deal(address(USDC), address(this), lAmtUSDC);
         // supply the USDC for aaveUSDC
@@ -1135,4 +1139,15 @@ contract AaveIntegrationTest is BaseTest {
         assertEq(_manager.shares(lOtherPair, USDC), 0);
         assertEq(_manager.shares(lThirdPair, USDC), 0);
     }
+
+    function testReturnAsset_Attack() external allNetworks allPairs {
+
+        // arrange
+        _increaseManagementOneToken(11e6);
+        ReturnAssetExploit lExploit = new ReturnAssetExploit(_pair);
+
+        // act & assert - attack should fail with our fix, and should succeed without the fix
+        vm.expectRevert(stdError.arithmeticError);
+        lExploit.attack(_manager);
+    }
 }