[audit] Staked ETH Rounding - Disable exact output (#460)

marktoda · ccashwell · gretzke · web-flow · commit eeb3eff28dd5 · 2025-05-02T12:01:36.000-04:00
* fix: steth rounding * feat: add fork tests * feat: fork tests * fix(audit): revert exactOutput swaps for wstETH<>stETH Due to an unavoidable property of the share mechanics used internally within the stETH implementation, it is not currently possible for v4 to support exact outputs in either direction when interacting with stETH. Rather than violate the spirit of the exactOut request by either overshooting the conversion target or ignoring the unexpected delta, this implementation opts to disallow exactOut by reverting in the `beforeSwap` hook. This resolves H-01. * feat: add wsteth fork tests since it has a very specific rounding implementation we test it with the real thing * Use infura for fork testing * add infura api key to environment * bump deprecated ubuntu version * Only run fork tests if API key is present, else skip (#462) * fix: comments * revert: semgrep * fix: merge * fix: snapshots * fix: snapshots --------- Co-authored-by: Chris Cashwell <chris.cashwell@uniswap.org> Co-authored-by: gretzke <daniel@gretzke.de>
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -26,3 +26,4 @@ jobs:
         env:
           FOUNDRY_PROFILE: ci
           FORGE_SNAPSHOT_CHECK: true
+          INFURA_API_KEY: ${{ secrets.INFURA_API_KEY }}
diff --git a/foundry.toml b/foundry.toml
@@ -38,5 +38,4 @@ sepolia = "https://rpc.sepolia.org"
 unichain_sepolia = "https://sepolia.unichain.org"
 base_sepolia = "https://sepolia.base.org"
 arbitrum_sepolia = "https://sepolia-rollup.arbitrum.io/rpc"
-
-# See more config options https://github.com/foundry-rs/foundry/tree/master/config
+mainnet = "https://mainnet.infura.io/v3/${INFURA_API_KEY}"
diff --git a/snapshots/StateViewTest.json b/snapshots/StateViewTest.json
@@ -1,12 +1,12 @@
 {
-  "StateView_extsload_getFeeGrowthGlobals": "2203",
-  "StateView_extsload_getFeeGrowthInside": "7875",
-  "StateView_extsload_getLiquidity": "1439",
-  "StateView_extsload_getPositionInfo": "2761",
-  "StateView_extsload_getPositionLiquidity": "1691",
-  "StateView_extsload_getSlot0": "1486",
-  "StateView_extsload_getTickBitmap": "1432",
-  "StateView_extsload_getTickFeeGrowthOutside": "2490",
-  "StateView_extsload_getTickInfo": "2693",
-  "StateView_extsload_getTickLiquidity": "1686"
+  "StateView_extsload_getFeeGrowthGlobals": "8703",
+  "StateView_extsload_getFeeGrowthInside": "24375",
+  "StateView_extsload_getLiquidity": "5939",
+  "StateView_extsload_getPositionInfo": "11261",
+  "StateView_extsload_getPositionLiquidity": "6191",
+  "StateView_extsload_getSlot0": "5986",
+  "StateView_extsload_getTickBitmap": "5932",
+  "StateView_extsload_getTickFeeGrowthOutside": "8990",
+  "StateView_extsload_getTickInfo": "11193",
+  "StateView_extsload_getTickLiquidity": "6186"
 }
diff --git a/src/base/hooks/BaseTokenWrapperHook.sol b/src/base/hooks/BaseTokenWrapperHook.sol
@@ -36,6 +36,12 @@ abstract contract BaseTokenWrapperHook is BaseHook, DeltaResolver {
     /// @dev Fee must be 0 as wrapper pools don't charge fees
     error InvalidPoolFee();
 
+    /// @notice Thrown when exact input swaps are not supported
+    error ExactInputNotSupported();
+
+    /// @notice Thrown when exact output swaps are not supported
+    error ExactOutputNotSupported();
+
     /// @notice The wrapped token currency (e.g., WETH)
     Currency public immutable wrapperCurrency;
 
@@ -118,27 +124,25 @@ abstract contract BaseTokenWrapperHook is BaseHook, DeltaResolver {
         returns (bytes4, BeforeSwapDelta swapDelta, uint24)
     {
         bool isExactInput = params.amountSpecified < 0;
+        if (isExactInput && !_supportsExactInput()) revert ExactInputNotSupported();
+        if (!isExactInput && !_supportsExactOutput()) revert ExactOutputNotSupported();
 
         if (wrapZeroForOne == params.zeroForOne) {
             // we are wrapping
             uint256 inputAmount =
                 isExactInput ? uint256(-params.amountSpecified) : _getWrapInputRequired(uint256(params.amountSpecified));
-            _take(underlyingCurrency, address(this), inputAmount);
-            uint256 wrappedAmount = _deposit(inputAmount);
-            _settle(wrapperCurrency, address(this), wrappedAmount);
+            (uint256 actualUnderlyingAmount, uint256 wrappedAmount) = _deposit(inputAmount);
             int128 amountUnspecified =
-                isExactInput ? -wrappedAmount.toInt256().toInt128() : inputAmount.toInt256().toInt128();
+                isExactInput ? -wrappedAmount.toInt256().toInt128() : actualUnderlyingAmount.toInt256().toInt128();
             swapDelta = toBeforeSwapDelta(-params.amountSpecified.toInt128(), amountUnspecified);
         } else {
             // we are unwrapping
             uint256 inputAmount = isExactInput
                 ? uint256(-params.amountSpecified)
                 : _getUnwrapInputRequired(uint256(params.amountSpecified));
-            _take(wrapperCurrency, address(this), inputAmount);
-            uint256 unwrappedAmount = _withdraw(inputAmount);
-            _settle(underlyingCurrency, address(this), unwrappedAmount);
+            (uint256 actualWrappedAmount, uint256 unwrappedAmount) = _withdraw(inputAmount);
             int128 amountUnspecified =
-                isExactInput ? -unwrappedAmount.toInt256().toInt128() : inputAmount.toInt256().toInt128();
+                isExactInput ? -unwrappedAmount.toInt256().toInt128() : actualWrappedAmount.toInt256().toInt128();
             swapDelta = toBeforeSwapDelta(-params.amountSpecified.toInt128(), amountUnspecified);
         }
 
@@ -155,17 +159,29 @@ abstract contract BaseTokenWrapperHook is BaseHook, DeltaResolver {
 
     /// @notice Deposits underlying tokens to receive wrapper tokens
     /// @param underlyingAmount The amount of underlying tokens to deposit
+    /// @return actualUnderlyingAmount the actual number of underlying tokens used, i.e. to account for rebasing rounding errors
     /// @return wrappedAmount The amount of wrapper tokens received
-    /// @dev Implementing contracts should handle the wrapping operation
-    ///      The base contract will handle settling tokens with the pool manager
-    function _deposit(uint256 underlyingAmount) internal virtual returns (uint256 wrappedAmount);
+    /// @dev Implementing contracts should handle:
+    //    - taking tokens from PoolManager
+    //    - performing the wrapping operation
+    //    - settling tokens on PoolManager
+    function _deposit(uint256 underlyingAmount)
+        internal
+        virtual
+        returns (uint256 actualUnderlyingAmount, uint256 wrappedAmount);
 
     /// @notice Withdraws wrapper tokens to receive underlying tokens
     /// @param wrappedAmount The amount of wrapper tokens to withdraw
+    /// @return actualWrappedAmount the actual number of wrapped tokens used, i.e. to account for rebasing rounding errors
     /// @return underlyingAmount The amount of underlying tokens received
-    /// @dev Implementing contracts should handle the unwrapping operation
-    ///      The base contract will handle settling tokens with the pool manager
-    function _withdraw(uint256 wrappedAmount) internal virtual returns (uint256 underlyingAmount);
+    /// @dev Implementing contracts should handle:
+    //    - taking tokens from PoolManager
+    //    - performing the unwrapping operation
+    //    - settling tokens on PoolManager
+    function _withdraw(uint256 wrappedAmount)
+        internal
+        virtual
+        returns (uint256 actualWrappedAmount, uint256 underlyingAmount);
 
     /// @notice Calculates underlying tokens needed to receive desired wrapper tokens
     /// @param wrappedAmount The desired amount of wrapper tokens
@@ -184,4 +200,18 @@ abstract contract BaseTokenWrapperHook is BaseHook, DeltaResolver {
     function _getUnwrapInputRequired(uint256 underlyingAmount) internal view virtual returns (uint256) {
         return underlyingAmount;
     }
+
+    /// @notice Indicates whether the hook supports exact output swaps
+    /// @dev Default implementation returns true
+    /// @dev Override for wrappers that cannot support exact output swaps
+    function _supportsExactOutput() internal view virtual returns (bool) {
+        return true;
+    }
+
+    /// @notice Indicates whether the hook supports exact input swaps
+    /// @dev Default implementation returns true
+    /// @dev Override for wrappers that cannot support exact input swaps
+    function _supportsExactInput() internal view virtual returns (bool) {
+        return true;
+    }
 }
diff --git a/src/hooks/WETHHook.sol b/src/hooks/WETHHook.sol
@@ -27,15 +27,26 @@ contract WETHHook is BaseTokenWrapperHook {
     }
 
     /// @inheritdoc BaseTokenWrapperHook
-    function _deposit(uint256 underlyingAmount) internal override returns (uint256) {
-        weth.deposit{value: underlyingAmount}();
-        return underlyingAmount; // 1:1 ratio
+    /// @dev Note the WETH deposit relies on the WETH wrapper having a receive function that mints WETH to msg.sender
+    function _deposit(uint256 underlyingAmount) internal override returns (uint256, uint256) {
+        // Sync WETH on PoolManager
+        poolManager.sync(wrapperCurrency);
+        // take ETH from PoolManager and deposit directly into the WETH contract
+        // this will mint WETH to msg.sender (PoolManager in this case)
+        _take(underlyingCurrency, address(weth), underlyingAmount);
+        // Settle on PoolManager which will take into account the new weth
+        poolManager.settle();
+        return (underlyingAmount, underlyingAmount); // 1:1 ratio
     }
 
     /// @inheritdoc BaseTokenWrapperHook
-    function _withdraw(uint256 wrapperAmount) internal override returns (uint256) {
+    function _withdraw(uint256 wrapperAmount) internal override returns (uint256, uint256) {
+        // take WETH into this hook contract
+        _take(wrapperCurrency, address(this), wrapperAmount);
+        // Withdraw WETH - this returns ETH back to this hook contract
         weth.withdraw(wrapperAmount);
-        return wrapperAmount; // 1:1 ratio
+        _settle(underlyingCurrency, address(this), wrapperAmount);
+        return (wrapperAmount, wrapperAmount); // 1:1 ratio
     }
 
     /// @notice Required to receive ETH
diff --git a/src/hooks/WstETHHook.sol b/src/hooks/WstETHHook.sol
@@ -3,15 +3,20 @@ pragma solidity 0.8.26;
 
 import {ERC20} from "solmate/src/tokens/ERC20.sol";
 import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
-import {Currency, CurrencyLibrary} from "@uniswap/v4-core/src/types/Currency.sol";
+import {FixedPointMathLib} from "solmate/src/utils/FixedPointMathLib.sol";
+import {Currency} from "@uniswap/v4-core/src/types/Currency.sol";
 import {BaseTokenWrapperHook} from "../base/hooks/BaseTokenWrapperHook.sol";
-import {IWstETH} from "../interfaces/external/IWstETH.sol";
+import {IWstETH, IStETH} from "../interfaces/external/IWstETH.sol";
+import {SafeTransferLib} from "solmate/src/utils/SafeTransferLib.sol";
 
 /// @title Wrapped Staked ETH (wstETH) Hook
 /// @notice Hook for wrapping/unwrapping stETH/wstETH in Uniswap V4 pools
 /// @dev Implements dynamic exchange rate wrapping/unwrapping between stETH and wstETH
 /// @dev wstETH represents stETH with accrued staking rewards, maintaining a dynamic exchange rate
 contract WstETHHook is BaseTokenWrapperHook {
+    using FixedPointMathLib for uint256;
+    using SafeTransferLib for ERC20;
+
     /// @notice The wstETH contract used for wrapping/unwrapping operations
     IWstETH public immutable wstETH;
 
@@ -22,28 +27,39 @@ contract WstETHHook is BaseTokenWrapperHook {
     constructor(IPoolManager _manager, IWstETH _wsteth)
         BaseTokenWrapperHook(
             _manager,
-            Currency.wrap(address(_wsteth)), // wrapper token is wsteth
+            Currency.wrap(address(_wsteth)), // wrapper token is wstETH
             Currency.wrap(_wsteth.stETH()) // underlying token is stETH
         )
     {
         wstETH = _wsteth;
-        ERC20(Currency.unwrap(underlyingCurrency)).approve(address(wstETH), type(uint256).max);
+        ERC20(Currency.unwrap(underlyingCurrency)).safeApprove(address(wstETH), type(uint256).max);
     }
 
     /// @inheritdoc BaseTokenWrapperHook
-    /// @notice Wraps stETH to wstETH
-    /// @param underlyingAmount Amount of stETH to wrap
-    /// @return Amount of wstETH received
-    function _deposit(uint256 underlyingAmount) internal override returns (uint256) {
-        return wstETH.wrap(underlyingAmount);
+    function _deposit(uint256 underlyingAmount)
+        internal
+        override
+        returns (uint256 actualUnderlyingAmount, uint256 wrappedAmount)
+    {
+        _take(underlyingCurrency, address(this), underlyingAmount);
+        // For wrapping, the key is ensuring we wrap exactly what we got
+        actualUnderlyingAmount = IStETH(Currency.unwrap(underlyingCurrency)).balanceOf(address(this));
+
+        // Wrap exactly what we have (which might be 1-2 wei less than requested)
+        wrappedAmount = wstETH.wrap(actualUnderlyingAmount);
+        _settle(wrapperCurrency, address(this), wrappedAmount);
     }
 
     /// @inheritdoc BaseTokenWrapperHook
-    /// @notice Unwraps wstETH to stETH
-    /// @param wrapperAmount Amount of wstETH to unwrap
-    /// @return Amount of stETH received
-    function _withdraw(uint256 wrapperAmount) internal override returns (uint256) {
-        return wstETH.unwrap(wrapperAmount);
+    function _withdraw(uint256 wrapperAmount)
+        internal
+        override
+        returns (uint256 actualWrappedAmount, uint256 unwrappedAmount)
+    {
+        _take(wrapperCurrency, address(this), wrapperAmount);
+        actualWrappedAmount = wrapperAmount;
+        unwrappedAmount = wstETH.unwrap(actualWrappedAmount);
+        _settle(underlyingCurrency, address(this), unwrappedAmount);
     }
 
     /// @inheritdoc BaseTokenWrapperHook
@@ -52,7 +68,7 @@ contract WstETHHook is BaseTokenWrapperHook {
     /// @return Amount of stETH required
     /// @dev Uses current stETH/wstETH exchange rate for calculation
     function _getWrapInputRequired(uint256 wrappedAmount) internal view override returns (uint256) {
-        return wstETH.getStETHByWstETH(wrappedAmount);
+        return wrappedAmount.divWadUp(wstETH.tokensPerStEth());
     }
 
     /// @inheritdoc BaseTokenWrapperHook
@@ -63,4 +79,9 @@ contract WstETHHook is BaseTokenWrapperHook {
     function _getUnwrapInputRequired(uint256 underlyingAmount) internal view override returns (uint256) {
         return wstETH.getWstETHByStETH(underlyingAmount);
     }
+
+    /// @inheritdoc BaseTokenWrapperHook
+    function _supportsExactOutput() internal pure override returns (bool) {
+        return false;
+    }
 }
diff --git a/src/interfaces/external/IWstETH.sol b/src/interfaces/external/IWstETH.sol
@@ -11,5 +11,15 @@ interface IWstETH {
     function unwrap(uint256 _wstETHAmount) external returns (uint256);
     function getStETHByWstETH(uint256 _wstETHAmount) external view returns (uint256);
     function getWstETHByStETH(uint256 _stETHAmount) external view returns (uint256);
+    function tokensPerStEth() external view returns (uint256);
+    function stEthPerToken() external view returns (uint256);
     function stETH() external view returns (address);
 }
+
+interface IStETH {
+    function getSharesByPooledEth(uint256 stEthAmount) external view returns (uint256);
+    function getPooledEthByShares(uint256 _sharesAmount) external view returns (uint256);
+    function sharesOf(address _account) external view returns (uint256);
+    function transferShares(address recipient, uint256 shares) external;
+    function balanceOf(address _account) external view returns (uint256);
+}
diff --git a/test/hooks/WETHHook.t.sol b/test/hooks/WETHHook.t.sol
@@ -94,7 +94,7 @@ contract WETHHookTest is Test, Deployers {
 
         vm.startPrank(alice);
         vm.expectEmit(true, true, true, true);
-        emit Transfer(address(0), address(hook), wrapAmount);
+        emit Transfer(address(0), address(manager), wrapAmount);
         vm.expectEmit(true, true, true, true);
         emit Transfer(address(manager), address(alice), wrapAmount);
 
@@ -175,7 +175,7 @@ contract WETHHookTest is Test, Deployers {
 
         vm.startPrank(alice);
         vm.expectEmit(true, true, true, true);
-        emit Transfer(address(0), address(hook), wrapAmount);
+        emit Transfer(address(0), address(manager), wrapAmount);
         vm.expectEmit(true, true, true, true);
         emit Transfer(address(manager), address(alice), wrapAmount);
 
diff --git a/test/hooks/WstETHHook.fork.t.sol b/test/hooks/WstETHHook.fork.t.sol
diff --git a/test/hooks/WstETHHook.t.sol b/test/hooks/WstETHHook.t.sol
diff --git a/test/mocks/MockWstETH.sol b/test/mocks/MockWstETH.sol
diff --git a/test/shared/TestRouter.sol b/test/shared/TestRouter.sol
