add safecallback (#123)

* add safecallback * use manager
Uniswap · Jun 18, 2024 · 1ca56ab · 1ca56ab
1 parent 5082779
commit 1ca56ab
Show file tree

Hide file tree

Showing 17 changed files with 136 additions and 119 deletions.
diff --git a/.forge-snapshots/FullRangeAddInitialLiquidity.snap b/.forge-snapshots/FullRangeAddInitialLiquidity.snap
@@ -1 +1 @@
-311073
+311181
diff --git a/.forge-snapshots/FullRangeAddLiquidity.snap b/.forge-snapshots/FullRangeAddLiquidity.snap
@@ -1 +1 @@
-122882
+122990
diff --git a/.forge-snapshots/FullRangeFirstSwap.snap b/.forge-snapshots/FullRangeFirstSwap.snap
@@ -1 +1 @@
-80283
+80220
diff --git a/.forge-snapshots/FullRangeInitialize.snap b/.forge-snapshots/FullRangeInitialize.snap
@@ -1 +1 @@
-1015169
+1015181
diff --git a/.forge-snapshots/FullRangeRemoveLiquidity.snap b/.forge-snapshots/FullRangeRemoveLiquidity.snap
@@ -1 +1 @@
-110476
+110566
diff --git a/.forge-snapshots/FullRangeRemoveLiquidityAndRebalance.snap b/.forge-snapshots/FullRangeRemoveLiquidityAndRebalance.snap
@@ -1 +1 @@
-239954
+240044
diff --git a/.forge-snapshots/FullRangeSecondSwap.snap b/.forge-snapshots/FullRangeSecondSwap.snap
@@ -1 +1 @@
-45993
+45930
diff --git a/.forge-snapshots/FullRangeSwap.snap b/.forge-snapshots/FullRangeSwap.snap
@@ -1 +1 @@
-79414
+79351
diff --git a/.forge-snapshots/TWAMMSubmitOrder.snap b/.forge-snapshots/TWAMMSubmitOrder.snap
@@ -1 +1 @@
-122355
+122336
diff --git a/contracts/BaseHook.sol b/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();

diff --git a/contracts/base/ImmutableState.sol b/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;
+    }
+}
diff --git a/contracts/base/SafeCallback.sol b/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);
+}
diff --git a/contracts/hooks/examples/FullRange.sol b/contracts/hooks/examples/FullRange.sol
@@ -26,7 +26,7 @@ import {BeforeSwapDelta, BeforeSwapDeltaLibrary} from "@uniswap/v4-core/src/type
 
 import "../../libraries/LiquidityAmounts.sol";
 
-contract FullRange is BaseHook, IUnlockCallback {
+contract FullRange is BaseHook {
     using CurrencyLibrary for Currency;
     using CurrencySettler for Currency;
     using PoolIdLibrary for PoolKey;
@@ -85,7 +85,7 @@ contract FullRange is BaseHook, IUnlockCallback {
 
     mapping(PoolId => PoolInfo) public poolInfo;
 
-    constructor(IPoolManager _poolManager) BaseHook(_poolManager) {}
+    constructor(IPoolManager _manager) BaseHook(_manager) {}
 
     modifier ensure(uint256 deadline) {
         if (deadline < block.timestamp) revert ExpiredPastDeadline();
@@ -126,13 +126,13 @@ contract FullRange is BaseHook, IUnlockCallback {
 
         PoolId poolId = key.toId();
 
-        (uint160 sqrtPriceX96,,,) = poolManager.getSlot0(poolId);
+        (uint160 sqrtPriceX96,,,) = manager.getSlot0(poolId);
 
         if (sqrtPriceX96 == 0) revert PoolNotInitialized();
 
         PoolInfo storage pool = poolInfo[poolId];
 
-        uint128 poolLiquidity = poolManager.getLiquidity(poolId);
+        uint128 poolLiquidity = manager.getLiquidity(poolId);
 
         liquidity = LiquidityAmounts.getLiquidityForAmounts(
             sqrtPriceX96,
@@ -184,7 +184,7 @@ contract FullRange is BaseHook, IUnlockCallback {
 
         PoolId poolId = key.toId();
 
-        (uint160 sqrtPriceX96,,,) = poolManager.getSlot0(poolId);
+        (uint160 sqrtPriceX96,,,) = manager.getSlot0(poolId);
 
         if (sqrtPriceX96 == 0) revert PoolNotInitialized();
 
@@ -260,17 +260,17 @@ contract FullRange is BaseHook, IUnlockCallback {
         internal
         returns (BalanceDelta delta)
     {
-        delta = abi.decode(poolManager.unlock(abi.encode(CallbackData(msg.sender, key, params))), (BalanceDelta));
+        delta = abi.decode(manager.unlock(abi.encode(CallbackData(msg.sender, key, params))), (BalanceDelta));
     }
 
     function _settleDeltas(address sender, PoolKey memory key, BalanceDelta delta) internal {
-        key.currency0.settle(poolManager, sender, uint256(int256(-delta.amount0())), false);
-        key.currency1.settle(poolManager, sender, uint256(int256(-delta.amount1())), false);
+        key.currency0.settle(manager, sender, uint256(int256(-delta.amount0())), false);
+        key.currency1.settle(manager, sender, uint256(int256(-delta.amount1())), false);
     }
 
     function _takeDeltas(address sender, PoolKey memory key, BalanceDelta delta) internal {
-        poolManager.take(key.currency0, sender, uint256(uint128(delta.amount0())));
-        poolManager.take(key.currency1, sender, uint256(uint128(delta.amount1())));
+        manager.take(key.currency0, sender, uint256(uint128(delta.amount0())));
+        manager.take(key.currency1, sender, uint256(uint128(delta.amount1())));
     }
 
     function _removeLiquidity(PoolKey memory key, IPoolManager.ModifyLiquidityParams memory params)
@@ -286,42 +286,37 @@ contract FullRange is BaseHook, IUnlockCallback {
 
         uint256 liquidityToRemove = FullMath.mulDiv(
             uint256(-params.liquidityDelta),
-            poolManager.getLiquidity(poolId),
+            manager.getLiquidity(poolId),
             UniswapV4ERC20(pool.liquidityToken).totalSupply()
         );
 
         params.liquidityDelta = -(liquidityToRemove.toInt256());
-        (delta,) = poolManager.modifyLiquidity(key, params, ZERO_BYTES);
+        (delta,) = manager.modifyLiquidity(key, params, ZERO_BYTES);
         pool.hasAccruedFees = false;
     }
 
-    function unlockCallback(bytes calldata rawData)
-        external
-        override(IUnlockCallback, BaseHook)
-        poolManagerOnly
-        returns (bytes memory)
-    {
+    function _unlockCallback(bytes calldata rawData) internal override returns (bytes memory) {
         CallbackData memory data = abi.decode(rawData, (CallbackData));
         BalanceDelta delta;
 
         if (data.params.liquidityDelta < 0) {
             delta = _removeLiquidity(data.key, data.params);
             _takeDeltas(data.sender, data.key, delta);
         } else {
-            (delta,) = poolManager.modifyLiquidity(data.key, data.params, ZERO_BYTES);
+            (delta,) = manager.modifyLiquidity(data.key, data.params, ZERO_BYTES);
             _settleDeltas(data.sender, data.key, delta);
         }
         return abi.encode(delta);
     }
 
     function _rebalance(PoolKey memory key) public {
         PoolId poolId = key.toId();
-        (BalanceDelta balanceDelta,) = poolManager.modifyLiquidity(
+        (BalanceDelta balanceDelta,) = manager.modifyLiquidity(
             key,
             IPoolManager.ModifyLiquidityParams({
                 tickLower: MIN_TICK,
                 tickUpper: MAX_TICK,
-                liquidityDelta: -(poolManager.getLiquidity(poolId).toInt256()),
+                liquidityDelta: -(manager.getLiquidity(poolId).toInt256()),
                 salt: 0
             }),
             ZERO_BYTES
@@ -333,9 +328,9 @@ contract FullRange is BaseHook, IUnlockCallback {
             ) * FixedPointMathLib.sqrt(FixedPoint96.Q96)
         ).toUint160();
 
-        (uint160 sqrtPriceX96,,,) = poolManager.getSlot0(poolId);
+        (uint160 sqrtPriceX96,,,) = manager.getSlot0(poolId);
 
-        poolManager.swap(
+        manager.swap(
             key,
             IPoolManager.SwapParams({
                 zeroForOne: newSqrtPriceX96 < sqrtPriceX96,
@@ -353,7 +348,7 @@ contract FullRange is BaseHook, IUnlockCallback {
             uint256(uint128(balanceDelta.amount1()))
         );
 
-        (BalanceDelta balanceDeltaAfter,) = poolManager.modifyLiquidity(
+        (BalanceDelta balanceDeltaAfter,) = manager.modifyLiquidity(
             key,
             IPoolManager.ModifyLiquidityParams({
                 tickLower: MIN_TICK,
@@ -368,6 +363,6 @@ contract FullRange is BaseHook, IUnlockCallback {
         uint128 donateAmount0 = uint128(balanceDelta.amount0() + balanceDeltaAfter.amount0());
         uint128 donateAmount1 = uint128(balanceDelta.amount1() + balanceDeltaAfter.amount1());
 
-        poolManager.donate(key, donateAmount0, donateAmount1, ZERO_BYTES);
+        manager.donate(key, donateAmount0, donateAmount1, ZERO_BYTES);
     }
 }
diff --git a/contracts/hooks/examples/GeomeanOracle.sol b/contracts/hooks/examples/GeomeanOracle.sol
@@ -61,7 +61,7 @@ contract GeomeanOracle is BaseHook {
         return uint32(block.timestamp);
     }
 
-    constructor(IPoolManager _poolManager) BaseHook(_poolManager) {}
+    constructor(IPoolManager _manager) BaseHook(_manager) {}
 
     function getHookPermissions() public pure override returns (Hooks.Permissions memory) {
         return Hooks.Permissions({
@@ -86,20 +86,20 @@ contract GeomeanOracle is BaseHook {
         external
         view
         override
-        poolManagerOnly
+        onlyByManager
         returns (bytes4)
     {
         // This is to limit the fragmentation of pools using this oracle hook. In other words,
         // there may only be one pool per pair of tokens that use this hook. The tick spacing is set to the maximum
         // because we only allow max range liquidity in this pool.
-        if (key.fee != 0 || key.tickSpacing != poolManager.MAX_TICK_SPACING()) revert OnlyOneOraclePoolAllowed();
+        if (key.fee != 0 || key.tickSpacing != manager.MAX_TICK_SPACING()) revert OnlyOneOraclePoolAllowed();
         return GeomeanOracle.beforeInitialize.selector;
     }
 
     function afterInitialize(address, PoolKey calldata key, uint160, int24, bytes calldata)
         external
         override
-        poolManagerOnly
+        onlyByManager
         returns (bytes4)
     {
         PoolId id = key.toId();
@@ -110,9 +110,9 @@ contract GeomeanOracle is BaseHook {
     /// @dev Called before any action that potentially modifies pool price or liquidity, such as swap or modify position
     function _updatePool(PoolKey calldata key) private {
         PoolId id = key.toId();
-        (, int24 tick,,) = poolManager.getSlot0(id);
+        (, int24 tick,,) = manager.getSlot0(id);
 
-        uint128 liquidity = poolManager.getLiquidity(id);
+        uint128 liquidity = manager.getLiquidity(id);
 
         (states[id].index, states[id].cardinality) = observations[id].write(
             states[id].index, _blockTimestamp(), tick, liquidity, states[id].cardinality, states[id].cardinalityNext
@@ -124,8 +124,8 @@ contract GeomeanOracle is BaseHook {
         PoolKey calldata key,
         IPoolManager.ModifyLiquidityParams calldata params,
         bytes calldata
-    ) external override poolManagerOnly returns (bytes4) {
-        int24 maxTickSpacing = poolManager.MAX_TICK_SPACING();
+    ) external override onlyByManager returns (bytes4) {
+        int24 maxTickSpacing = manager.MAX_TICK_SPACING();
         if (
             params.tickLower != TickMath.minUsableTick(maxTickSpacing)
                 || params.tickUpper != TickMath.maxUsableTick(maxTickSpacing)
@@ -139,14 +139,14 @@ contract GeomeanOracle is BaseHook {
         PoolKey calldata,
         IPoolManager.ModifyLiquidityParams calldata,
         bytes calldata
-    ) external view override poolManagerOnly returns (bytes4) {
+    ) external view override onlyByManager returns (bytes4) {
         revert OraclePoolMustLockLiquidity();
     }
 
     function beforeSwap(address, PoolKey calldata key, IPoolManager.SwapParams calldata, bytes calldata)
         external
         override
-        poolManagerOnly
+        onlyByManager
         returns (bytes4, BeforeSwapDelta, uint24)
     {
         _updatePool(key);
@@ -163,9 +163,9 @@ contract GeomeanOracle is BaseHook {
 
         ObservationState memory state = states[id];
 
-        (, int24 tick,,) = poolManager.getSlot0(id);
+        (, int24 tick,,) = manager.getSlot0(id);
 
-        uint128 liquidity = poolManager.getLiquidity(id);
+        uint128 liquidity = manager.getLiquidity(id);
 
         return observations[id].observe(_blockTimestamp(), secondsAgos, tick, state.index, liquidity, state.cardinality);
     }