Access lock (Uniswap#404)

.forge-snapshots/before swap hook, already cached dynamic fee.snap

@@ -1 +1 @@
-190294
+191802

.forge-snapshots/cached dynamic fee, no hooks.snap

@@ -1 +1 @@
-145654
+147163

.forge-snapshots/donate gas with 1 token.snap

@@ -1 +1 @@
-137259
+138487

.forge-snapshots/donate gas with 2 tokens.snap

@@ -1 +1 @@
-186980
+185645

.forge-snapshots/erc20 collect protocol fees.snap

@@ -1 +1 @@
-27033
+26991

.forge-snapshots/gas overhead of no-op lock.snap

@@ -1 +1 @@
-14918
+14940

.forge-snapshots/initialize.snap

@@ -1 +1 @@
-75025
+77663

.forge-snapshots/mint with empty hook.snap

@@ -1 +1 @@
-319543
+318152

.forge-snapshots/mint with native token.snap

@@ -1 +1 @@
-199129
+200911

.forge-snapshots/mint.snap

@@ -1 +1 @@
-202244
+200853

.forge-snapshots/modify position with noop.snap

@@ -1 +1 @@
-52336
+56127

.forge-snapshots/native collect protocol fees.snap

@@ -1 +1 @@
-38699
+38657

.forge-snapshots/poolManager bytecode size.snap

@@ -1 +1 @@
-25078
+25255

.forge-snapshots/simple swap with native.snap

@@ -1 +1 @@
-191062
+196049

.forge-snapshots/simple swap.snap

@@ -1 +1 @@
-202765
+204617

.forge-snapshots/swap against liquidity with native token.snap

@@ -1 +1 @@
-121658
+126645

.forge-snapshots/swap against liquidity.snap

@@ -1 +1 @@
-109483
+114135

.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 dynamic fee.snap

@@ -1 +1 @@
-189534
+191043

.forge-snapshots/swap with hooks.snap

@@ -1 +1 @@
-109462
+114114

.forge-snapshots/swap with noop.snap

@@ -1 +1 @@
-45202
+48947

.forge-snapshots/update dynamic fee in before swap.snap

@@ -1 +1 @@
-196118
+197626

src/PoolManager.sol

@@ -90,6 +90,18 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         return Lockers.getLocker(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
@@ -105,6 +117,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)
             {
@@ -127,6 +141,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);
     }
@@ -174,26 +191,27 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         if (pools[id].isNotInitialized()) revert PoolNotInitialized();
     }
 
-    modifier onlyByLocker() {
-        address locker = Lockers.getCurrentLocker();
-        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;
@@ -233,6 +251,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);
     }
 
@@ -244,14 +265,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;
@@ -285,6 +313,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
         );
@@ -298,14 +329,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);
@@ -317,6 +355,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
@@ -400,6 +441,10 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         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

@@ -23,7 +23,8 @@ interface IPoolManager is IFees, IClaims {
 
     /// @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 The ERC1155 being deposited is not the Uniswap ERC1155
     error NotPoolManagerToken();
@@ -123,6 +124,9 @@ interface IPoolManager is IFees, IClaims {
     /// @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;
     }