Hook Middleware

.forge-snapshots/FullRangeAddInitialLiquidity.snap

@@ -1 +1 @@
-311073
+311181

.forge-snapshots/FullRangeAddLiquidity.snap

@@ -1 +1 @@
-122882
+122990

.forge-snapshots/FullRangeFirstSwap.snap

@@ -1 +1 @@
-80283
+80220

.forge-snapshots/FullRangeInitialize.snap

@@ -1 +1 @@
-1015169
+1015181

.forge-snapshots/FullRangeRemoveLiquidity.snap

@@ -1 +1 @@
-110476
+110566

.forge-snapshots/FullRangeRemoveLiquidityAndRebalance.snap

@@ -1 +1 @@
-239954
+240044

.forge-snapshots/FullRangeSecondSwap.snap

@@ -1 +1 @@
-45993
+45930

.forge-snapshots/FullRangeSwap.snap

@@ -1 +1 @@
-79414
+79351

.forge-snapshots/TWAMMSubmitOrder.snap

@@ -1 +1 @@
-122355
+122336

contracts/BaseHook.sol

@@ -7,28 +7,19 @@ import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
 import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
 import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
 import {BeforeSwapDelta} from "@uniswap/v4-core/src/types/BeforeSwapDelta.sol";
+import {SafeCallback} from "./base/SafeCallback.sol";
+import {ImmutableState} from "./base/ImmutableState.sol";
 
-abstract contract BaseHook is IHooks {
-    error NotPoolManager();
+abstract contract BaseHook is IHooks, SafeCallback {
     error NotSelf();
     error InvalidPool();
     error LockFailure();
     error HookNotImplemented();
 
-    /// @notice The address of the pool manager
-    IPoolManager public immutable poolManager;
-
-    constructor(IPoolManager _poolManager) {
-        poolManager = _poolManager;
+    constructor(IPoolManager _manager) ImmutableState(_manager) {
         validateHookAddress(this);
     }
 
-    /// @dev Only the pool manager may call this function
-    modifier poolManagerOnly() {
-        if (msg.sender != address(poolManager)) revert NotPoolManager();
-        _;
-    }
-
     /// @dev Only this address may call this function
     modifier selfOnly() {
         if (msg.sender != address(this)) revert NotSelf();
@@ -50,7 +41,7 @@ abstract contract BaseHook is IHooks {
         Hooks.validateHookPermissions(_this, getHookPermissions());
     }
 
-    function unlockCallback(bytes calldata data) external virtual poolManagerOnly returns (bytes memory) {
+    function _unlockCallback(bytes calldata data) internal virtual override returns (bytes memory) {
         (bool success, bytes memory returnData) = address(this).call(data);
         if (success) return returnData;
         if (returnData.length == 0) revert LockFailure();

contracts/base/ImmutableState.sol

@@ -0,0 +1,12 @@
+// SPDX-License-Identifier: GPL-2.0-or-later
+pragma solidity ^0.8.19;
+
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+
+contract ImmutableState {
+    IPoolManager public immutable manager;
+
+    constructor(IPoolManager _manager) {
+        manager = _manager;
+    }
+}

contracts/base/SafeCallback.sol

@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import {IUnlockCallback} from "@uniswap/v4-core/src/interfaces/callback/IUnlockCallback.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {ImmutableState} from "./ImmutableState.sol";
+
+abstract contract SafeCallback is ImmutableState, IUnlockCallback {
+    error NotManager();
+
+    modifier onlyByManager() {
+        if (msg.sender != address(manager)) revert NotManager();
+        _;
+    }
+
+    /// @dev We force the onlyByManager modifier by exposing a virtual function after the onlyByManager check.
+    function unlockCallback(bytes calldata data) external onlyByManager returns (bytes memory) {
+        return _unlockCallback(data);
+    }
+
+    function _unlockCallback(bytes calldata data) internal virtual returns (bytes memory);
+}

contracts/hooks/examples/FeeTaker.sol

@@ -0,0 +1,104 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import {BaseHook} from "../../BaseHook.sol";
+import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
+import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
+import {Currency, CurrencyLibrary} from "@uniswap/v4-core/src/types/Currency.sol";
+import {SafeCast} from "@uniswap/v4-core/src/libraries/SafeCast.sol";
+
+abstract contract FeeTaker is BaseHook {
+    using SafeCast for uint256;
+
+    bytes internal constant ZERO_BYTES = bytes("");
+
+    constructor(IPoolManager _poolManager) BaseHook(_poolManager) {}
+
+    /**
+     * @notice This hook takes a fee from the unspecified token after a swap.
+     * @dev This can be overridden if more permissions are needed.
+     */
+    function getHookPermissions() public pure virtual override returns (Hooks.Permissions memory) {
+        return Hooks.Permissions({
+            beforeInitialize: false,
+            afterInitialize: false,
+            beforeAddLiquidity: false,
+            afterAddLiquidity: false,
+            beforeRemoveLiquidity: false,
+            afterRemoveLiquidity: false,
+            beforeSwap: false,
+            afterSwap: true,
+            beforeDonate: false,
+            afterDonate: false,
+            beforeSwapReturnDelta: false,
+            afterSwapReturnDelta: true,
+            afterAddLiquidityReturnDelta: false,
+            afterRemoveLiquidityReturnDelta: false
+        });
+    }
+
+    function afterSwap(
+        address sender,
+        PoolKey calldata key,
+        IPoolManager.SwapParams calldata params,
+        BalanceDelta delta,
+        bytes calldata hookData
+    ) external override onlyByManager returns (bytes4, int128) {
+        //(Currency currencyUnspecified, amountUnspecified) = key.getUnspecified(params);
+
+        // fee will be in the unspecified token of the swap
+        bool currency0Specified = (params.amountSpecified < 0 == params.zeroForOne);
+        (Currency currencyUnspecified, int128 amountUnspecified) =
+            (currency0Specified) ? (key.currency1, delta.amount1()) : (key.currency0, delta.amount0());
+        // if exactOutput swap, get the absolute output amount
+        if (amountUnspecified < 0) amountUnspecified = -amountUnspecified;
+
+        uint256 feeAmount = _feeAmount(amountUnspecified);
+        // mint ERC6909 instead of take to avoid edge case where PM doesn't have enough balance
+        manager.mint(address(this), CurrencyLibrary.toId(currencyUnspecified), feeAmount);
+
+        (bytes4 selector, int128 amount) = _afterSwap(sender, key, params, delta, hookData);
+        return (selector, feeAmount.toInt128() + amount);
+    }
+
+    function withdraw(Currency[] calldata currencies) external {
+        manager.unlock(abi.encode(currencies));
+    }
+
+    function _unlockCallback(bytes calldata rawData) internal override returns (bytes memory) {
+        Currency[] memory currencies = abi.decode(rawData, (Currency[]));
+        uint256 length = currencies.length;
+        for (uint256 i = 0; i < length;) {
+            uint256 amount = manager.balanceOf(address(this), CurrencyLibrary.toId(currencies[i]));
+            manager.burn(address(this), CurrencyLibrary.toId(currencies[i]), amount);
+            manager.take(currencies[i], _recipient(), amount);
+            unchecked {
+                ++i;
+            }
+        }
+        return ZERO_BYTES;
+    }
+
+    /**
+     * @dev This is a virtual function that should be overridden so it returns the fee charged for a given amount.
+     */
+    function _feeAmount(int128 amountUnspecified) internal view virtual returns (uint256);
+
+    /**
+     * @dev This is a virtual function that should be overridden so it returns the address to receive the fee.
+     */
+    function _recipient() internal view virtual returns (address);
+
+    /**
+     * @dev This can be overridden to add logic after a swap.
+     */
+    function _afterSwap(address, PoolKey memory, IPoolManager.SwapParams memory, BalanceDelta, bytes calldata)
+        internal
+        virtual
+        returns (bytes4, int128)
+    {
+        return (BaseHook.afterSwap.selector, 0);
+    }
+}

contracts/hooks/examples/FeeTaking.sol

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {SafeCast} from "@uniswap/v4-core/src/libraries/SafeCast.sol";
+import {Owned} from "solmate/auth/Owned.sol";
+import {FeeTaker} from "./FeeTaker.sol";
+
+contract FeeTaking is FeeTaker, Owned {
+    using SafeCast for uint256;
+
+    uint128 private constant TOTAL_BIPS = 10000;
+    uint128 public immutable swapFeeBips;
+    address public treasury;
+
+    constructor(IPoolManager _poolManager, uint128 _swapFeeBips, address _owner, address _treasury)
+        FeeTaker(_poolManager)
+        Owned(_owner)
+    {
+        swapFeeBips = _swapFeeBips;
+        treasury = _treasury;
+    }
+
+    function setTreasury(address _treasury) external onlyOwner {
+        treasury = _treasury;
+    }
+
+    function _feeAmount(int128 amountUnspecified) internal view override returns (uint256) {
+        return uint128(amountUnspecified) * swapFeeBips / TOTAL_BIPS;
+    }
+
+    function _recipient() internal view override returns (address) {
+        return treasury;
+    }
+}