[yAudit-03] Fix swap fee accrual on realized gain

src/contracts/AbstractARM.sol

@@ -428,7 +428,7 @@ abstract contract AbstractARM is OwnableOperable, ERC20Upgradeable, ReentrancyGu
             amountOut = convertedAmountIn * config.buyPrice / PRICE_SCALE;
         }
 
-        _validateAndConsumeSwapLiquidity(config, isBuySide, outToken, amountOut);
+        _validateAndConsumeSwapLiquidity(config, isBuySide, outToken, amountIn, amountOut);
 
         // Transfer the input tokens from the caller to this ARM contract
         inToken.transferFrom(msg.sender, address(this), amountIn);
@@ -468,7 +468,7 @@ abstract contract AbstractARM is OwnableOperable, ERC20Upgradeable, ReentrancyGu
             amountIn = convertedAmountOut * PRICE_SCALE / config.buyPrice + 3;
         }
 
-        _validateAndConsumeSwapLiquidity(config, isBuySide, outToken, amountOut);
+        _validateAndConsumeSwapLiquidity(config, isBuySide, outToken, amountIn, amountOut);
 
         // Transfer the input tokens from the caller to this ARM contract
         inToken.transferFrom(msg.sender, address(this), amountIn);
@@ -518,16 +518,18 @@ abstract contract AbstractARM is OwnableOperable, ERC20Upgradeable, ReentrancyGu
     /// @dev Validate swap reserves and consume the per-base liquidity limit.
     /// @param isBuySide True when the ARM buys base asset and pays out liquidity asset.
     /// @param outToken Swap output token address.
+    /// @param amountIn Swap input token amount.
     /// @param amountOut Swap output token amount.
     function _validateAndConsumeSwapLiquidity(
         BaseAssetConfig storage config,
         bool isBuySide,
         IERC20 outToken,
+        uint256 amountIn,
         uint256 amountOut
     ) private {
         uint256 remaining;
         if (isBuySide) {
-            _accrueSwapFee(config.buyPrice, config.crossPrice, amountOut);
+            _accrueSwapFee(config, amountIn, amountOut);
             remaining = config.buyLiquidityRemaining;
             if (amountOut > remaining) revert InsufficientLiquidity();
             unchecked {
@@ -581,13 +583,14 @@ abstract contract AbstractARM is OwnableOperable, ERC20Upgradeable, ReentrancyGu
         return baseDecimals > liquidityAssetDecimals ? amount * 1e12 : amount / 1e12;
     }
 
-    /// @dev Accrue fees on discounted buy-side swaps using the recognized NAV gain.
-    /// @param buyPrice Price the ARM paid for the base asset.
-    /// @param crossPrice Price used to value the base asset in totalAssets().
+    /// @dev Accrue fees on buy-side swaps using the gain realized from the settled input and output amounts.
+    /// @param config Base asset configuration used to value the input received by the ARM.
+    /// @param amountIn Base asset amount received by the ARM.
     /// @param amountOut Liquidity asset amount paid out by the ARM.
-    function _accrueSwapFee(uint256 buyPrice, uint256 crossPrice, uint256 amountOut) internal {
-        uint256 feeMultiplier = (crossPrice - buyPrice) * uint256(fee) * PRICE_SCALE / (buyPrice * FEE_SCALE);
-        feesAccrued = SafeCast.toUint128(feesAccrued + amountOut * feeMultiplier / PRICE_SCALE);
+    function _accrueSwapFee(BaseAssetConfig storage config, uint256 amountIn, uint256 amountOut) internal {
+        uint256 realizedAssets = _convertToAssets(config, amountIn) * config.crossPrice / PRICE_SCALE;
+        uint256 gain = realizedAssets > amountOut ? realizedAssets - amountOut : 0;
+        feesAccrued = SafeCast.toUint128(feesAccrued + gain * uint256(fee) / FEE_SCALE);
     }
 
     ////////////////////////////////////////////////////

test/fork/EthenaARM/SwapExactTokensForTokens.t.sol

@@ -90,8 +90,7 @@ contract Fork_Concrete_EthenaARM_swapExactTokensForTokens_Test_ is Fork_Shared_T
         // Precompute expected amount out
         uint256 traderate = _buyPrice();
         uint256 expectedAmountOut = (susde.convertToAssets(AMOUNT_IN) * traderate) / 1e36;
-        uint256 expectedFee =
-            expectedAmountOut * _swapFeeMultiplier(_buyPrice(), _crossPrice(), ethenaARM.fee()) / PRICE_SCALE;
+        uint256 expectedFee = _expectedBuySideFee(AMOUNT_IN, expectedAmountOut);
 
         // Expected events
         vm.expectEmit({emitter: address(susde)});
@@ -112,9 +111,7 @@ contract Fork_Concrete_EthenaARM_swapExactTokensForTokens_Test_ is Fork_Shared_T
         assertEq(obtained[1], expectedAmountOut, "Obtained USDe amount should match expected output");
         assertEq(usdeBalanceAfter, usdeBalanceBefore + expectedAmountOut, "USDe balance should have increased");
         assertEq(susdeBalanceBefore, susdeBalanceAfter + AMOUNT_IN, "SUSDe balance should have decreased");
-        assertEq(
-            ethenaARM.feesAccrued() - feesAccruedBefore, expectedFee, "Fees accrued should match output multiplier"
-        );
+        assertEq(ethenaARM.feesAccrued() - feesAccruedBefore, expectedFee, "Fees accrued should match realized gain");
     }
 
     function test_swapExactTokensForTokens_SUSDE_To_USDE_WithOutstandingWithdrawals_Sig1() public {

test/fork/EthenaARM/shared/Shared.sol

@@ -120,10 +120,10 @@ abstract contract Fork_Shared_Test is Base_Test_ {
         crossPrice = crossPriceMem;
     }
 
-    function _swapFeeMultiplier(uint256 buyPrice, uint256 crossPrice, uint256 fee) internal view returns (uint256) {
-        uint256 priceScale = PRICE_SCALE;
-        if (buyPrice == 0 || fee == 0) return 0;
-        return (crossPrice - buyPrice) * fee * priceScale / (buyPrice * FEE_SCALE);
+    function _expectedBuySideFee(uint256 amountIn, uint256 amountOut) internal view returns (uint256) {
+        uint256 realizedAssets = susde.convertToAssets(amountIn) * _crossPrice() / PRICE_SCALE;
+        uint256 gain = realizedAssets > amountOut ? realizedAssets - amountOut : 0;
+        return gain * ethenaARM.fee() / FEE_SCALE;
     }
 
     function _ignite() internal virtual {

test/fork/MultiAssetARM/SwapExactTokensForTokens.t.sol

@@ -162,8 +162,7 @@ contract Fork_Concrete_MultiAssetARM_swapExactTokensForTokens_Test_ is Fork_Shar
     function _swapBuy(IERC20 token) internal {
         uint256 buyPrice = _buyPrice(token);
         uint256 expectedAmountOut = _convertToAssets(token, AMOUNT_IN) * buyPrice / PRICE_SCALE;
-        uint256 expectedFee =
-            expectedAmountOut * _swapFeeMultiplier(buyPrice, _crossPrice(token), arm.fee()) / PRICE_SCALE;
+        uint256 expectedFee = _expectedBuySideFee(token, AMOUNT_IN, expectedAmountOut);
 
         uint256 wethBefore = weth.balanceOf(address(this));
         uint256 baseBefore = token.balanceOf(address(this));

test/fork/MultiAssetARM/SwapTokensForExactTokens.t.sol

@@ -153,7 +153,7 @@ contract Fork_Concrete_MultiAssetARM_swapTokensForExactTokens_Test_ is Fork_Shar
     function _swapBuyExact(IERC20 token) internal {
         uint256 buyPrice = _buyPrice(token);
         uint256 expectedAmountIn = _convertToShares(token, AMOUNT_OUT) * PRICE_SCALE / buyPrice + 3;
-        uint256 expectedFee = AMOUNT_OUT * _swapFeeMultiplier(buyPrice, _crossPrice(token), arm.fee()) / PRICE_SCALE;
+        uint256 expectedFee = _expectedBuySideFee(token, expectedAmountIn, AMOUNT_OUT);
 
         uint256 wethBefore = weth.balanceOf(address(this));
         uint256 baseBefore = token.balanceOf(address(this));

test/fork/MultiAssetARM/shared/Shared.sol

@@ -313,9 +313,10 @@ abstract contract Fork_Shared_Test is Base_Test_ {
         return _adapter(token).convertToShares(assets);
     }
 
-    function _swapFeeMultiplier(uint256 buyPrice, uint256 crossPrice, uint256 fee) internal pure returns (uint256) {
-        if (buyPrice == 0 || fee == 0) return 0;
-        return (crossPrice - buyPrice) * fee * PRICE_SCALE / (buyPrice * FEE_SCALE);
+    function _expectedBuySideFee(IERC20 token, uint256 amountIn, uint256 amountOut) internal view returns (uint256) {
+        uint256 realizedAssets = _convertToAssets(token, amountIn) * _crossPrice(token) / PRICE_SCALE;
+        uint256 gain = realizedAssets > amountOut ? realizedAssets - amountOut : 0;
+        return gain * arm.fee() / FEE_SCALE;
     }
 
     //////////////////////////////////////////////////////

test/fork/PaxosARM/SwapExactTokensForTokens.t.sol

@@ -69,8 +69,7 @@ contract Fork_Concrete_PaxosARM_swapExactTokensForTokens_Test_ is Fork_Shared_Te
     function _swapBuy(IERC20 token) internal {
         uint256 buyPrice = _buyPrice(token);
         uint256 expectedAmountOut = AMOUNT_IN * buyPrice / PRICE_SCALE;
-        uint256 expectedFee =
-            expectedAmountOut * _swapFeeMultiplier(buyPrice, _crossPrice(token), arm.fee()) / PRICE_SCALE;
+        uint256 expectedFee = _expectedBuySideFee(token, AMOUNT_IN, expectedAmountOut);
 
         uint256 usdcBefore = usdc.balanceOf(address(this));
         uint256 baseBefore = token.balanceOf(address(this));

test/fork/PaxosARM/SwapTokensForExactTokens.t.sol

@@ -22,7 +22,7 @@ contract Fork_Concrete_PaxosARM_swapTokensForExactTokens_Test_ is Fork_Shared_Te
     function test_swapTokensForExactTokens_Pyusd_To_Usdc() public {
         uint256 buyPrice = _buyPrice(pyusd);
         uint256 expectedAmountIn = AMOUNT_OUT * PRICE_SCALE / buyPrice + 3;
-        uint256 expectedFee = AMOUNT_OUT * _swapFeeMultiplier(buyPrice, _crossPrice(pyusd), arm.fee()) / PRICE_SCALE;
+        uint256 expectedFee = _expectedBuySideFee(pyusd, expectedAmountIn, AMOUNT_OUT);
 
         uint256 usdcBefore = usdc.balanceOf(address(this));
         uint256 baseBefore = pyusd.balanceOf(address(this));

test/fork/PaxosARM/shared/Shared.sol

@@ -223,8 +223,9 @@ abstract contract Fork_Shared_Test is Base_Test_ {
         return PaxosAssetAdapter(adapter);
     }
 
-    function _swapFeeMultiplier(uint256 buyPrice, uint256 crossPrice, uint256 fee) internal pure returns (uint256) {
-        if (buyPrice == 0 || fee == 0) return 0;
-        return (crossPrice - buyPrice) * fee * PRICE_SCALE / (buyPrice * FEE_SCALE);
+    function _expectedBuySideFee(IERC20 token, uint256 amountIn, uint256 amountOut) internal view returns (uint256) {
+        uint256 realizedAssets = amountIn * _crossPrice(token) / PRICE_SCALE;
+        uint256 gain = realizedAssets > amountOut ? realizedAssets - amountOut : 0;
+        return gain * arm.fee() / FEE_SCALE;
     }
 }