Move msg.sender check to modifier (#564)

* Move msg.sender check to modifier * Removing old check * remove extra line * Update SkipCallsTestHook.t.sol * snap rename
Uniswap · Apr 5, 2024 · 11aa1f5 · 11aa1f5
1 parent d6df729
commit 11aa1f5
Show file tree

Hide file tree

Showing 24 changed files with 53 additions and 25 deletions.
diff --git a/.forge-snapshots/addLiquidity with empty hook.snap b/.forge-snapshots/addLiquidity with empty hook.snap
@@ -1 +1 @@
-265373
+265381
diff --git a/.forge-snapshots/addLiquidity with native token.snap b/.forge-snapshots/addLiquidity with native token.snap
@@ -1 +1 @@
-140136
+140224
diff --git a/.forge-snapshots/addLiquidity.snap b/.forge-snapshots/addLiquidity.snap
@@ -1 +1 @@
-145453
+145541
diff --git a/.forge-snapshots/donate gas with 1 token.snap b/.forge-snapshots/donate gas with 1 token.snap
@@ -1 +1 @@
-101049
+101137
diff --git a/.forge-snapshots/donate gas with 2 tokens.snap b/.forge-snapshots/donate gas with 2 tokens.snap
@@ -1 +1 @@
-132039
+132127
diff --git a/.forge-snapshots/initialize.snap b/.forge-snapshots/initialize.snap
@@ -1 +1 @@
-51180
+51268
diff --git a/.forge-snapshots/poolManager bytecode size.snap b/.forge-snapshots/poolManager bytecode size.snap
@@ -1 +1 @@
-22965
+23083
diff --git a/.forge-snapshots/removeLiquidity with empty hook.snap b/.forge-snapshots/removeLiquidity with empty hook.snap
@@ -1 +1 @@
-55984
+56072
diff --git a/.forge-snapshots/removeLiquidity with native token.snap b/.forge-snapshots/removeLiquidity with native token.snap
@@ -1 +1 @@
-148170
+148258
diff --git a/.forge-snapshots/removeLiquidity.snap b/.forge-snapshots/removeLiquidity.snap
@@ -1 +1 @@
-149634
+149722
diff --git a/.forge-snapshots/simple swap with native.snap b/.forge-snapshots/simple swap with native.snap
@@ -1 +1 @@
-132814
+132902
diff --git a/.forge-snapshots/simple swap.snap b/.forge-snapshots/simple swap.snap
@@ -1 +1 @@
-146690
+146778
diff --git a/.forge-snapshots/swap against liquidity with native token.snap b/.forge-snapshots/swap against liquidity with native token.snap
@@ -1 +1 @@
-72273
+72361
diff --git a/.forge-snapshots/swap against liquidity.snap b/.forge-snapshots/swap against liquidity.snap
@@ -1 +1 @@
-60276
+60364
diff --git a/.forge-snapshots/swap burn 6909 for input.snap b/.forge-snapshots/swap burn 6909 for input.snap
@@ -1 +1 @@
-80306
+80394
diff --git a/.forge-snapshots/swap burn native 6909 for input.snap b/.forge-snapshots/swap burn native 6909 for input.snap
@@ -1 +1 @@
-76303
+76391
diff --git a/.forge-snapshots/swap mint native output as 6909.snap b/.forge-snapshots/swap mint native output as 6909.snap
@@ -1 +1 @@
-138654
+138742
diff --git a/.forge-snapshots/swap mint output as 6909.snap b/.forge-snapshots/swap mint output as 6909.snap
@@ -1 +1 @@
-155505
+155593
diff --git a/.forge-snapshots/swap skips hook call if hook is caller.snap b/.forge-snapshots/swap skips hook call if hook is caller.snap
@@ -0,0 +1 @@
+156254
diff --git a/.forge-snapshots/swap with dynamic fee.snap b/.forge-snapshots/swap with dynamic fee.snap
@@ -1 +1 @@
-89580
+89668
diff --git a/.forge-snapshots/swap with hooks.snap b/.forge-snapshots/swap with hooks.snap
@@ -1 +1 @@
-60254
+60342
diff --git a/.forge-snapshots/update dynamic fee in before swap.snap b/.forge-snapshots/update dynamic fee in before swap.snap
@@ -1 +1 @@
-140246
+140294
diff --git a/src/libraries/Hooks.sol b/src/libraries/Hooks.sol
@@ -96,17 +96,24 @@ library Hooks {
 
     /// @notice performs a hook call using the given calldata on the given hook
     function callHook(IHooks self, bytes memory data) internal {
-        if (msg.sender == address(self)) return;
         (bytes4 expectedSelector, bytes4 selector) = _callHook(self, data);
 
         if (selector != expectedSelector) {
             revert InvalidHookResponse();
         }
     }
 
+    /// @notice modifier to prevent calling a hook if they initiated the action
+    modifier noSelfCall(IHooks self) {
+        if (msg.sender != address(self)) {
+            _;
+        }
+    }
+
     /// @notice calls beforeInitialize hook if permissioned and validates return value
     function beforeInitialize(IHooks self, PoolKey memory key, uint160 sqrtPriceX96, bytes calldata hookData)
         internal
+        noSelfCall(self)
     {
         if (self.hasPermission(BEFORE_INITIALIZE_FLAG)) {
             self.callHook(
@@ -118,6 +125,7 @@ library Hooks {
     /// @notice calls afterInitialize hook if permissioned and validates return value
     function afterInitialize(IHooks self, PoolKey memory key, uint160 sqrtPriceX96, int24 tick, bytes calldata hookData)
         internal
+        noSelfCall(self)
     {
         if (self.hasPermission(AFTER_INITIALIZE_FLAG)) {
             self.callHook(
@@ -132,7 +140,7 @@ library Hooks {
         PoolKey memory key,
         IPoolManager.ModifyLiquidityParams memory params,
         bytes calldata hookData
-    ) internal {
+    ) internal noSelfCall(self) {
         if (params.liquidityDelta > 0 && self.hasPermission(BEFORE_ADD_LIQUIDITY_FLAG)) {
             self.callHook(abi.encodeWithSelector(IHooks.beforeAddLiquidity.selector, msg.sender, key, params, hookData));
         } else if (params.liquidityDelta <= 0 && self.hasPermission(BEFORE_REMOVE_LIQUIDITY_FLAG)) {
@@ -149,7 +157,7 @@ library Hooks {
         IPoolManager.ModifyLiquidityParams memory params,
         BalanceDelta delta,
         bytes calldata hookData
-    ) internal {
+    ) internal noSelfCall(self) {
         if (params.liquidityDelta > 0 && self.hasPermission(AFTER_ADD_LIQUIDITY_FLAG)) {
             self.callHook(
                 abi.encodeWithSelector(IHooks.afterAddLiquidity.selector, msg.sender, key, params, delta, hookData)
@@ -164,6 +172,7 @@ library Hooks {
     /// @notice calls beforeSwap hook if permissioned and validates return value
     function beforeSwap(IHooks self, PoolKey memory key, IPoolManager.SwapParams memory params, bytes calldata hookData)
         internal
+        noSelfCall(self)
     {
         if (self.hasPermission(BEFORE_SWAP_FLAG)) {
             self.callHook(abi.encodeWithSelector(IHooks.beforeSwap.selector, msg.sender, key, params, hookData));
@@ -177,7 +186,7 @@ library Hooks {
         IPoolManager.SwapParams memory params,
         BalanceDelta delta,
         bytes calldata hookData
-    ) internal {
+    ) internal noSelfCall(self) {
         if (self.hasPermission(AFTER_SWAP_FLAG)) {
             self.callHook(abi.encodeWithSelector(IHooks.afterSwap.selector, msg.sender, key, params, delta, hookData));
         }
@@ -186,6 +195,7 @@ library Hooks {
     /// @notice calls beforeDonate hook if permissioned and validates return value
     function beforeDonate(IHooks self, PoolKey memory key, uint256 amount0, uint256 amount1, bytes calldata hookData)
         internal
+        noSelfCall(self)
     {
         if (self.hasPermission(BEFORE_DONATE_FLAG)) {
             self.callHook(
@@ -197,6 +207,7 @@ library Hooks {
     /// @notice calls afterDonate hook if permissioned and validates return value
     function afterDonate(IHooks self, PoolKey memory key, uint256 amount0, uint256 amount1, bytes calldata hookData)
         internal
+        noSelfCall(self)
     {
         if (self.hasPermission(AFTER_DONATE_FLAG)) {
             self.callHook(

diff --git a/test/SkipCallsTestHook.t.sol b/test/SkipCallsTestHook.t.sol
@@ -160,6 +160,22 @@ contract SkipCallsTest is Test, Deployers, GasSnapshot {
         assertEq(skipCallsTestHook.counter(), 2);
     }
 
+    function test_gas_beforeSwap_skipIfCalledByHook() public {
+        SkipCallsTestHook skipCallsTestHook = SkipCallsTestHook(
+            address(uint160(type(uint160).max & clearAllHookPermisssionsMask | Hooks.BEFORE_SWAP_FLAG))
+        );
+
+        deploy(skipCallsTestHook);
+        approveAndAddLiquidity(skipCallsTestHook);
+        assertEq(skipCallsTestHook.counter(), 0);
+
+        // swaps and increments counter
+        snapStart("swap skips hook call if hook is caller");
+        swapRouter.swap(key, swapParams, testSettings, abi.encode(address(this)));
+        snapEnd();
+        assertEq(skipCallsTestHook.counter(), 1);
+    }
+
     function test_afterSwap_skipIfCalledByHook() public {
         SkipCallsTestHook skipCallsTestHook = SkipCallsTestHook(
             address(uint160(type(uint160).max & clearAllHookPermisssionsMask | Hooks.AFTER_SWAP_FLAG))