fix: #13

reservoir-labs · Dec 30, 2024 · af0178e · af0178e
1 parent e880dfe
commit af0178e
Show file tree

Hide file tree

Showing 7 changed files with 34 additions and 24 deletions.
diff --git a/src/ReservoirPair.sol b/src/ReservoirPair.sol
@@ -292,7 +292,7 @@ abstract contract ReservoirPair is IAssetManagedPair, ReservoirERC20, RGT {
             uint256 lReserveOut = lIsToken0 ? aReserve0 : aReserve1;
 
             if (lReserveOut - lTokenOutManaged < aAmount) {
-                assetManager.returnAsset(lIsToken0, aAmount - (lReserveOut - lTokenOutManaged));
+                assetManager.returnAsset(lIsToken0 ? aAmount - (lReserveOut - lTokenOutManaged) : 0, lIsToken0 ? 0 : aAmount - (lReserveOut - lTokenOutManaged));
                 require(_safeTransfer(aToken, aDestination, aAmount), TransferFailed());
             } else {
                 revert TransferFailed();
@@ -458,7 +458,7 @@ abstract contract ReservoirPair is IAssetManagedPair, ReservoirERC20, RGT {
 
             rAmtSkimmed = lTokenAmtManaged - type(uint104).max;
 
-            assetManager.returnAsset(aToken == token0, rAmtSkimmed);
+            assetManager.returnAsset(aToken == token0 ? rAmtSkimmed : 0, aToken == token1 ? 0 :rAmtSkimmed);
             address(aToken).safeTransfer(lRecoverer, rAmtSkimmed);
         }
     }

diff --git a/src/asset-management/EulerV2Manager.sol b/src/asset-management/EulerV2Manager.sol
@@ -26,6 +26,7 @@ contract EulerV2Manager is IAssetManager, Owned(msg.sender), RGT {
     event Divestment(IAssetManagedPair pair, IERC20 token, uint256 shares);
 
     error OutstandingSharesForVault();
+    error ReturnAssetZeroAmount();
 
     /// @dev Mapping from an ERC20 token to an Euler V2 vault.
     /// This implies that for a given asset, there can only be one vault at any one time.
@@ -235,12 +236,9 @@ contract EulerV2Manager is IAssetManager, Owned(msg.sender), RGT {
         _adjustManagement(lPair, lAmount0Change, lAmount1Change);
     }
 
-    function returnAsset(bool aToken0, uint256 aAmount) external {
-        require(aAmount > 0, "AM: ZERO_AMOUNT_REQUESTED");
-        IAssetManagedPair lPair = IAssetManagedPair(msg.sender);
-        int256 lAmount0Change = aToken0 ? -aAmount.toInt256() : int256(0);
-        int256 lAmount1Change = aToken0 ? int256(0) : -aAmount.toInt256();
-        _adjustManagement(lPair, lAmount0Change, lAmount1Change);
+    function returnAsset(uint256 aToken0Amt, uint256 aToken1Amt) external {
+        require(aToken0Amt > 0 || aToken1Amt > 0, ReturnAssetZeroAmount());
+        _adjustManagement(IAssetManagedPair(msg.sender), -aToken0Amt.toInt256(), -aToken1Amt.toInt256());
     }
 
     function _calculateChangeAmount(uint256 aReserve, uint256 aManaged) internal view returns (int256 rAmountChange) {

diff --git a/src/interfaces/IAssetManager.sol b/src/interfaces/IAssetManager.sol
@@ -14,5 +14,5 @@ interface IAssetManager {
 
     /// @notice called by the pair when it requires assets managed by the manager to be returned to the pair
     /// in order to fulfill swap requests or burn requests
-    function returnAsset(bool aToken0, uint256 aAmount) external;
+    function returnAsset(uint256 aToken0Amt, uint256 aToken1Amt) external;
 }
diff --git a/test/__mocks/AssetManager.sol b/test/__mocks/AssetManager.sol
@@ -1,10 +1,14 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.0;
 
+import { SafeCast } from "@openzeppelin/utils/math/SafeCast.sol";
+
 import { IAssetManager, IERC20 } from "src/interfaces/IAssetManager.sol";
 import { IAssetManagedPair } from "src/interfaces/IAssetManagedPair.sol";
 
 contract AssetManager is IAssetManager {
+    using SafeCast for uint256;
+
     mapping(IAssetManagedPair => mapping(IERC20 => uint256)) public getBalance;
 
     function adjustManagement(IAssetManagedPair aPair, int256 aToken0Amount, int256 aToken1Amount) public {
@@ -44,11 +48,10 @@ contract AssetManager is IAssetManager {
     // solhint-disable-next-line no-empty-blocks
     function afterLiquidityEvent() external { }
 
-    function returnAsset(bool aToken0, uint256 aAmount) external {
+    function returnAsset(uint256 aToken0Amt, uint256 aToken1Amt) external {
         IAssetManagedPair lPair = IAssetManagedPair(msg.sender);
-        int256 lAmount0Change = -int256(aToken0 ? aAmount : 0);
-        int256 lAmount1Change = -int256(aToken0 ? 0 : aAmount);
-        (aToken0 ? lPair.token0() : lPair.token1()).approve(address(msg.sender), aAmount);
-        adjustManagement(lPair, lAmount0Change, lAmount1Change);
+        lPair.token0().approve(address(msg.sender), aToken0Amt);
+        lPair.token1().approve(address(msg.sender), aToken1Amt);
+        adjustManagement(lPair, -aToken0Amt.toInt256(), -aToken1Amt.toInt256());
     }
 }
diff --git a/test/__mocks/AssetManagerReenter.sol b/test/__mocks/AssetManagerReenter.sol
@@ -1,11 +1,15 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.0;
 
+import { SafeCast } from "@openzeppelin/utils/math/SafeCast.sol";
+
 import { IAssetManager, IERC20 } from "src/interfaces/IAssetManager.sol";
 import { IAssetManagedPair } from "src/interfaces/IAssetManagedPair.sol";
 import { ReservoirPair } from "src/ReservoirPair.sol";
 
 contract AssetManagerReenter is IAssetManager {
+    using SafeCast for uint256;
+
     mapping(IAssetManagedPair => mapping(IERC20 => uint256)) public _getBalance;
 
     // this is solely to test reentrancy for ReservoirPair::mint/burn when the pair syncs
@@ -49,12 +53,11 @@ contract AssetManagerReenter is IAssetManager {
     // solhint-disable-next-line no-empty-blocks
     function afterLiquidityEvent() external { }
 
-    function returnAsset(bool aToken0, uint256 aAmount) external {
-        (aToken0 ? IAssetManagedPair(msg.sender).token0() : IAssetManagedPair(msg.sender).token1()).approve(
-            address(msg.sender), aAmount
-        );
-        IAssetManagedPair(msg.sender).adjustManagement(
-            aToken0 ? -int256(aAmount) : int256(0), aToken0 ? int256(0) : -int256(aAmount)
-        );
+    function returnAsset(uint256 aToken0Amt, uint256 aToken1Amt) external {
+        IAssetManagedPair(msg.sender).token0().approve(msg.sender, aToken0Amt);
+        IAssetManagedPair(msg.sender).token1().approve(msg.sender, aToken1Amt);
+
+        IAssetManagedPair(msg.sender).adjustManagement(-aToken0Amt.toInt256(), -aToken1Amt.toInt256());
     }
 }
+
diff --git a/test/__mocks/ReturnAssetExploit.sol b/test/__mocks/ReturnAssetExploit.sol
@@ -22,6 +22,6 @@ contract ReturnAssetExploit {
     }
 
     function attack(IAssetManager aAssetManager) external {
-        aAssetManager.returnAsset(false, 123_555);
+        aAssetManager.returnAsset(0, 123_555);
     }
 }
diff --git a/test/integration/Euler.t.sol b/test/integration/Euler.t.sol
@@ -603,7 +603,10 @@ contract EulerIntegrationTest is BaseTest {
         (int256 lExpectedToken0Calldata, int256 lExpectedToken1Calldata) =
             _pair.token0() == USDC ? (int256(-10), int256(0)) : (int256(0), int256(-10));
         _tokenA.mint(address(_pair), lReserveTokenA * 2);
-        vm.expectCall(address(_manager), abi.encodeCall(_manager.returnAsset, (_pair.token0() == USDC, 10)));
+        vm.expectCall(
+            address(_manager),
+            abi.encodeCall(_manager.returnAsset, (_pair.token0() == USDC ? 10 : 0, _pair.token1() == USDC ? 10 : 0))
+        );
         vm.expectCall(
             address(_pair), abi.encodeCall(_pair.adjustManagement, (lExpectedToken0Calldata, lExpectedToken1Calldata))
         );
@@ -636,7 +639,10 @@ contract EulerIntegrationTest is BaseTest {
         (int256 lExpectedToken0Calldata, int256 lExpectedToken1Calldata) =
             _pair.token0() == USDC ? (int256(-10), int256(0)) : (int256(0), int256(-10));
         _tokenA.mint(address(_pair), lReserveTokenA * 2);
-        vm.expectCall(address(_manager), abi.encodeCall(_manager.returnAsset, (_pair.token0() == USDC, 10)));
+        vm.expectCall(
+            address(_manager),
+            abi.encodeCall(_manager.returnAsset, (_pair.token0() == USDC ? 10 : 0, _pair.token1() == USDC ? 10 : 0))
+        );
         vm.expectCall(
             address(_pair), abi.encodeCall(_pair.adjustManagement, (lExpectedToken0Calldata, lExpectedToken1Calldata))
         );