Access lock (#404)

.forge-snapshots/erc20 collect protocol fees.snap

@@ -1 +1 @@
-26979
+26979

.forge-snapshots/modify position with noop.snap

@@ -1 +1 @@
-52267
+52267

.forge-snapshots/native collect protocol fees.snap

@@ -1 +1 @@
-38645
+38645

.forge-snapshots/swap burn claim for input.snap

@@ -1 +1 @@
-128222
+133204

.forge-snapshots/swap mint output as claim.snap

@@ -1 +1 @@
-212934
+216787

.forge-snapshots/swap with noop.snap

@@ -1 +1 @@
-45214
+45214

src/PoolManager.sol

@@ -94,6 +94,18 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         return LockDataLibrary.getLock(i);
     }
 
+    /// @notice This will revert if a function is called by any address other than the current locker OR the most recently called, pre-permissioned hook.
+    modifier onlyByLocker() {
+        _checkLocker(msg.sender, Lockers.getCurrentLocker(), Lockers.getCurrentHook());
+        _;
+    }
+
+    function _checkLocker(address caller, address locker, IHooks hook) internal pure {
+        if (caller == locker) return;
+        if (caller == address(hook) && hook.hasPermissionToAccessLock()) return;
+        revert LockedBy(locker, address(hook));
+    }
+
     /// @inheritdoc IPoolManager
     function initialize(PoolKey memory key, uint160 sqrtPriceX96, bytes calldata hookData)
         external
@@ -109,6 +121,8 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         if (key.currency0 >= key.currency1) revert CurrenciesOutOfOrderOrEqual();
         if (!key.hooks.isValidHookAddress(key.fee)) revert Hooks.HookAddressNotValid(address(key.hooks));
 
+        (bool set) = Lockers.setCurrentHook(key.hooks);
+
         if (key.hooks.shouldCallBeforeInitialize()) {
             if (key.hooks.beforeInitialize(msg.sender, key, sqrtPriceX96, hookData) != IHooks.beforeInitialize.selector)
             {
@@ -131,6 +145,9 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
             }
         }
 
+        // We only want to clear the current hook if it was set in setCurrentHook in this execution frame.
+        if (set) Lockers.clearCurrentHook();
+
         // On intitalize we emit the key's fee, which tells us all fee settings a pool can have: either a static swap fee or dynamic swap fee and if the hook has enabled swap or withdraw fees.
         emit Initialize(id, key.currency0, key.currency1, key.fee, key.tickSpacing, key.hooks);
     }
@@ -178,26 +195,27 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         if (pools[id].isNotInitialized()) revert PoolNotInitialized();
     }
 
-    modifier onlyByLocker() {
-        address locker = lockData.getActiveLock();
-        if (msg.sender != locker) revert LockedBy(locker);
-        _;
-    }
-
     /// @inheritdoc IPoolManager
     function modifyPosition(
         PoolKey memory key,
         IPoolManager.ModifyPositionParams memory params,
         bytes calldata hookData
     ) external override noDelegateCall onlyByLocker returns (BalanceDelta delta) {
+        (bool set) = Lockers.setCurrentHook(key.hooks);
+
         PoolId id = key.toId();
         _checkPoolInitialized(id);
 
         if (key.hooks.shouldCallBeforeModifyPosition()) {
             bytes4 selector = key.hooks.beforeModifyPosition(msg.sender, key, params, hookData);
             // Sentinel return value used to signify that a NoOp occurred.
-            if (key.hooks.isValidNoOpCall(selector)) return BalanceDeltaLibrary.MAXIMUM_DELTA;
-            else if (selector != IHooks.beforeModifyPosition.selector) revert Hooks.InvalidHookResponse();
+            if (key.hooks.isValidNoOpCall(selector)) {
+                // We only want to clear the current hook if it was set in setCurrentHook in this execution frame.
+                if (set) Lockers.clearCurrentHook();
+                return BalanceDeltaLibrary.MAXIMUM_DELTA;
+            } else if (selector != IHooks.beforeModifyPosition.selector) {
+                revert Hooks.InvalidHookResponse();
+            }
         }
 
         Pool.FeeAmounts memory feeAmounts;
@@ -237,6 +255,9 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
             }
         }
 
+        // We only want to clear the current hook if it was set in setCurrentHook in this execution frame.
+        if (set) Lockers.clearCurrentHook();
+
         emit ModifyPosition(id, msg.sender, params.tickLower, params.tickUpper, params.liquidityDelta);
     }
 
@@ -248,14 +269,21 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         onlyByLocker
         returns (BalanceDelta delta)
     {
+        (bool set) = Lockers.setCurrentHook(key.hooks);
+
         PoolId id = key.toId();
         _checkPoolInitialized(id);
 
         if (key.hooks.shouldCallBeforeSwap()) {
             bytes4 selector = key.hooks.beforeSwap(msg.sender, key, params, hookData);
             // Sentinel return value used to signify that a NoOp occurred.
-            if (key.hooks.isValidNoOpCall(selector)) return BalanceDeltaLibrary.MAXIMUM_DELTA;
-            else if (selector != IHooks.beforeSwap.selector) revert Hooks.InvalidHookResponse();
+            if (key.hooks.isValidNoOpCall(selector)) {
+                // We only want to clear the current hook if it was set in setCurrentHook in this execution frame.
+                if (set) Lockers.clearCurrentHook();
+                return BalanceDeltaLibrary.MAXIMUM_DELTA;
+            } else if (selector != IHooks.beforeSwap.selector) {
+                revert Hooks.InvalidHookResponse();
+            }
         }
 
         uint256 feeForProtocol;
@@ -289,6 +317,9 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
             }
         }
 
+        // We only want to clear the current hook if it was set in setCurrentHook in this execution frame.
+        if (set) Lockers.clearCurrentHook();
+
         emit Swap(
             id, msg.sender, delta.amount0(), delta.amount1(), state.sqrtPriceX96, state.liquidity, state.tick, swapFee
         );
@@ -302,14 +333,21 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         onlyByLocker
         returns (BalanceDelta delta)
     {
+        (bool set) = Lockers.setCurrentHook(key.hooks);
+
         PoolId id = key.toId();
         _checkPoolInitialized(id);
 
         if (key.hooks.shouldCallBeforeDonate()) {
             bytes4 selector = key.hooks.beforeDonate(msg.sender, key, amount0, amount1, hookData);
             // Sentinel return value used to signify that a NoOp occurred.
-            if (key.hooks.isValidNoOpCall(selector)) return BalanceDeltaLibrary.MAXIMUM_DELTA;
-            else if (selector != IHooks.beforeDonate.selector) revert Hooks.InvalidHookResponse();
+            if (key.hooks.isValidNoOpCall(selector)) {
+                // We only want to clear the current hook if it was set in setCurrentHook in this execution frame.
+                if (set) Lockers.clearCurrentHook();
+                return BalanceDeltaLibrary.MAXIMUM_DELTA;
+            } else if (selector != IHooks.beforeDonate.selector) {
+                revert Hooks.InvalidHookResponse();
+            }
         }
 
         delta = pools[id].donate(amount0, amount1);
@@ -321,6 +359,9 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
                 revert Hooks.InvalidHookResponse();
             }
         }
+
+        // We only want to clear the current hook if it was set in setCurrentHook in this execution frame.
+        if (set) Lockers.clearCurrentHook();
     }
 
     /// @inheritdoc IPoolManager
@@ -404,6 +445,18 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         return value;
     }
 
+    function getLockLength() external view returns (uint256 _length) {
+        return Lockers.length();
+    }
+
+    function getLockNonzeroDeltaCount() external view returns (uint256 _nonzeroDeltaCount) {
+        return Lockers.nonzeroDeltaCount();
+    }
+
+    function getCurrentHook() external view returns (IHooks) {
+        return Lockers.getCurrentHook();
+    }
+
     /// @notice receive native tokens for native pools
     receive() external payable {}
 }

src/interfaces/IPoolManager.sol

@@ -22,7 +22,8 @@ interface IPoolManager is IFees {
 
     /// @notice Thrown when a function is called by an address that is not the current locker
     /// @param locker The current locker
-    error LockedBy(address locker);
+    /// @param currentHook The most recently called hook
+    error LockedBy(address locker, address currentHook);
 
     /// @notice Pools are limited to type(int16).max tickSpacing in #initialize, to prevent overflow
     error TickSpacingTooLarge();
@@ -119,6 +120,9 @@ interface IPoolManager is IFees {
     /// @notice Returns the length of the lockers array, which is the number of locks open on the PoolManager.
     function getLockLength() external view returns (uint256 _length);
 
+    /// @notice Returns the most recently called hook.
+    function getCurrentHook() external view returns (IHooks _currentHook);
+
     /// @notice Returns the number of nonzero deltas open on the PoolManager that must be zerod by the close of the initial lock.
     function getLockNonzeroDeltaCount() external view returns (uint256 _nonzeroDeltaCount);
 

src/libraries/Hooks.sol

@@ -20,10 +20,11 @@ library Hooks {
     uint256 internal constant BEFORE_DONATE_FLAG = 1 << 153;
     uint256 internal constant AFTER_DONATE_FLAG = 1 << 152;
     uint256 internal constant NO_OP_FLAG = 1 << 151;
+    uint256 internal constant ACCESS_LOCK_FLAG = 1 << 150;
 
     bytes4 public constant NO_OP_SELECTOR = bytes4(keccak256(abi.encodePacked("NoOp")));
 
-    struct Calls {
+    struct Permissions {
         bool beforeInitialize;
         bool afterInitialize;
         bool beforeModifyPosition;
@@ -33,6 +34,7 @@ library Hooks {
         bool beforeDonate;
         bool afterDonate;
         bool noOp;
+        bool accessLock;
     }
 
     /// @notice Thrown if the address will not lead to the specified hook calls being called
@@ -44,17 +46,19 @@ library Hooks {
 
     /// @notice Utility function intended to be used in hook constructors to ensure
     /// the deployed hooks address causes the intended hooks to be called
-    /// @param calls The hooks that are intended to be called
-    /// @dev calls param is memory as the function will be called from constructors
-    function validateHookAddress(IHooks self, Calls memory calls) internal pure {
+    /// @param permissions The hooks that are intended to be called
+    /// @dev permissions param is memory as the function will be called from constructors
+    function validateHookPermissions(IHooks self, Permissions memory permissions) internal pure {
         if (
-            calls.beforeInitialize != shouldCallBeforeInitialize(self)
-                || calls.afterInitialize != shouldCallAfterInitialize(self)
-                || calls.beforeModifyPosition != shouldCallBeforeModifyPosition(self)
-                || calls.afterModifyPosition != shouldCallAfterModifyPosition(self)
-                || calls.beforeSwap != shouldCallBeforeSwap(self) || calls.afterSwap != shouldCallAfterSwap(self)
-                || calls.beforeDonate != shouldCallBeforeDonate(self) || calls.afterDonate != shouldCallAfterDonate(self)
-                || calls.noOp != hasPermissionToNoOp(self)
+            permissions.beforeInitialize != shouldCallBeforeInitialize(self)
+                || permissions.afterInitialize != shouldCallAfterInitialize(self)
+                || permissions.beforeModifyPosition != shouldCallBeforeModifyPosition(self)
+                || permissions.afterModifyPosition != shouldCallAfterModifyPosition(self)
+                || permissions.beforeSwap != shouldCallBeforeSwap(self)
+                || permissions.afterSwap != shouldCallAfterSwap(self)
+                || permissions.beforeDonate != shouldCallBeforeDonate(self)
+                || permissions.afterDonate != shouldCallAfterDonate(self) || permissions.noOp != hasPermissionToNoOp(self)
+                || permissions.accessLock != hasPermissionToAccessLock(self)
         ) {
             revert HookAddressNotValid(address(self));
         }
@@ -74,7 +78,7 @@ library Hooks {
         return address(hook) == address(0)
             ? !fee.isDynamicFee() && !fee.hasHookSwapFee() && !fee.hasHookWithdrawFee()
             : (
-                uint160(address(hook)) >= AFTER_DONATE_FLAG || fee.isDynamicFee() || fee.hasHookSwapFee()
+                uint160(address(hook)) >= ACCESS_LOCK_FLAG || fee.isDynamicFee() || fee.hasHookSwapFee()
                     || fee.hasHookWithdrawFee()
             );
     }
@@ -111,6 +115,10 @@ library Hooks {
         return uint256(uint160(address(self))) & AFTER_DONATE_FLAG != 0;
     }
 
+    function hasPermissionToAccessLock(IHooks self) internal pure returns (bool) {
+        return uint256(uint160(address(self))) & ACCESS_LOCK_FLAG != 0;
+    }
+
     function hasPermissionToNoOp(IHooks self) internal pure returns (bool) {
         return uint256(uint160(address(self))) & NO_OP_FLAG != 0;
     }

src/libraries/Lockers.sol

@@ -1,13 +1,18 @@
 // SPDX-License-Identifier: BUSL-1.1
 pragma solidity ^0.8.20;
 
+import {IHooks} from "../interfaces/IHooks.sol";
+
 /// @notice This is a temporary library that allows us to use transient storage (tstore/tload)
 /// for the lockers array and nonzero delta count.
 /// TODO: This library can be deleted when we have the transient keyword support in solidity.
 library Lockers {
     // The starting slot for an array of lockers, stored transiently.
     uint256 constant LOCKERS_SLOT = uint256(keccak256("Lockers")) - 1;
 
+    // The starting slot for an array of hook addresses per locker, stored transiently.
+    uint256 constant HOOK_ADDRESS_SLOT = uint256(keccak256("HookAddress")) - 1;
+
     // The slot holding the number of nonzero deltas.
     uint256 constant NONZERO_DELTA_COUNT = uint256(keccak256("NonzeroDeltaCount")) - 1;
 
@@ -94,4 +99,34 @@ library Lockers {
             tstore(slot, count)
         }
     }
+
+    function getCurrentHook() internal view returns (IHooks currentHook) {
+        return IHooks(getHook(length()));
+    }
+
+    function getHook(uint256 i) internal view returns (address hook) {
+        uint256 slot = HOOK_ADDRESS_SLOT + i;
+        assembly {
+            hook := tload(slot)
+        }
+    }
+
+    function setCurrentHook(IHooks currentHook) internal returns (bool set) {
+        // Set the hook address for the current locker if the address is 0.
+        // If the address is nonzero, a hook has already been set for this lock, and is not allowed to be updated or cleared at the end of the call.
+        if (address(getCurrentHook()) == address(0)) {
+            uint256 slot = HOOK_ADDRESS_SLOT + length();
+            assembly {
+                tstore(slot, currentHook)
+            }
+            return true;
+        }
+    }
+
+    function clearCurrentHook() internal {
+        uint256 slot = HOOK_ADDRESS_SLOT + length();
+        assembly {
+            tstore(slot, 0)
+        }
+    }
 }