enforce implementation address mined

.forge-snapshots/NormalSwap.snap

@@ -0,0 +1 @@
+79041

.forge-snapshots/ProtectedSwap.snap

@@ -0,0 +1 @@
+205221

.forge-snapshots/UnprotectedSwap.snap

@@ -0,0 +1 @@
+130276

contracts/middleware/BaseMiddleware.sol

@@ -14,7 +14,10 @@ contract BaseMiddleware is Proxy {
     IPoolManager public immutable manager;
     address public immutable implementation;
 
+    error FlagsMismatch();
+
     constructor(IPoolManager _manager, address _impl) {
+        _ensureValidFlags(_impl);
         manager = _manager;
         implementation = _impl;
     }
@@ -27,4 +30,10 @@ contract BaseMiddleware is Proxy {
     receive() external payable {
         _delegate(_implementation());
     }
+
+    function _ensureValidFlags(address _impl) internal view virtual {
+        if (uint160(address(this)) & Hooks.ALL_HOOK_MASK != uint160(_impl) & Hooks.ALL_HOOK_MASK) {
+            revert FlagsMismatch();
+        }
+    }
 }

contracts/middleware/BaseMiddlewareFactory.sol

@@ -4,9 +4,6 @@ pragma solidity ^0.8.19;
 import {IMiddlewareFactory} from "../interfaces/IMiddlewareFactory.sol";
 import {BaseMiddleware} from "./BaseMiddleware.sol";
 import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
-import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
-import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol";
-import {IBaseHook} from "../interfaces/IBaseHook.sol";
 
 contract BaseMiddlewareFactory is IMiddlewareFactory {
     mapping(address => address) private _implementations;
@@ -23,12 +20,11 @@ contract BaseMiddlewareFactory is IMiddlewareFactory {
 
     function createMiddleware(address implementation, bytes32 salt) external override returns (address middleware) {
         middleware = _deployMiddleware(implementation, salt);
-        Hooks.validateHookPermissions(IHooks(middleware), IBaseHook(implementation).getHookPermissions());
         _implementations[middleware] = implementation;
         emit MiddlewareCreated(implementation, middleware);
     }
 
     function _deployMiddleware(address implementation, bytes32 salt) internal virtual returns (address middleware) {
-        return address(new BaseMiddleware{salt: salt}(manager, implementation));
+        middleware = address(new BaseMiddleware{salt: salt}(manager, implementation));
     }
 }

contracts/middleware/MiddlewareProtect.sol

@@ -19,6 +19,7 @@ import {BeforeSwapDelta, BeforeSwapDeltaLibrary} from "@uniswap/v4-core/src/type
 import {console} from "./../../lib/forge-gas-snapshot/lib/forge-std/src/console.sol";
 import {LPFeeLibrary} from "@uniswap/v4-core/src/libraries/LPFeeLibrary.sol";
 import {TickMath} from "@uniswap/v4-core/src/libraries/TickMath.sol";
+import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
 
 contract MiddlewareProtect is MiddlewareRemove {
     using CustomRevert for bytes4;
@@ -29,25 +30,17 @@ contract MiddlewareProtect is MiddlewareRemove {
     using LPFeeLibrary for uint24;
     using BalanceDeltaLibrary for BalanceDelta;
 
-    error ForbiddenDynamicFee();
     error HookModifiedOutput();
+    error ForbiddenDynamicFee();
+    error MustEnableAfterSwapFlag();
 
     // todo: use tstore
     BalanceDelta private quote;
 
     uint160 public constant MIN_PRICE_LIMIT = TickMath.MIN_SQRT_PRICE + 1;
     uint160 public constant MAX_PRICE_LIMIT = TickMath.MAX_SQRT_PRICE - 1;
 
-    constructor(IPoolManager _manager, address _impl) MiddlewareRemove(_manager, _impl) {
-        IHooks middleware = IHooks(address(this));
-        if (
-            middleware.hasPermission(Hooks.BEFORE_SWAP_RETURNS_DELTA_FLAG)
-                || middleware.hasPermission(Hooks.AFTER_SWAP_RETURNS_DELTA_FLAG)
-                || middleware.hasPermission(Hooks.AFTER_ADD_LIQUIDITY_RETURNS_DELTA_FLAG)
-        ) {
-            HookPermissionForbidden.selector.revertWith(address(this));
-        }
-    }
+    constructor(IPoolManager _manager, address _impl) MiddlewareRemove(_manager, _impl) {}
 
     function beforeInitialize(address sender, PoolKey calldata key, uint160 sqrtPriceX96, bytes calldata hookData)
         external
@@ -92,7 +85,13 @@ contract MiddlewareProtect is MiddlewareRemove {
         external
         returns (bytes4, int128)
     {
-        if (delta != quote) revert HookModifiedOutput();
+        IHooks implementation = IHooks(address(implementation));
+        if (implementation.hasPermission(Hooks.BEFORE_SWAP_FLAG)) {
+            if (delta != quote) revert HookModifiedOutput();
+            if (!implementation.hasPermission(Hooks.AFTER_SWAP_FLAG)) {
+                return (BaseHook.afterSwap.selector, 0);
+            }
+        }
         (bool success, bytes memory returnData) = address(implementation).delegatecall(msg.data);
         if (!success) {
             _handleRevert(returnData);
@@ -120,4 +119,31 @@ contract MiddlewareProtect is MiddlewareRemove {
             revert(add(32, returnData), mload(returnData))
         }
     }
+
+    function _ensureValidFlags(address _impl) internal view virtual override {
+        IHooks This = IHooks(address(this));
+        if (
+            This.hasPermission(Hooks.BEFORE_SWAP_RETURNS_DELTA_FLAG)
+                || This.hasPermission(Hooks.AFTER_SWAP_RETURNS_DELTA_FLAG)
+                || This.hasPermission(Hooks.AFTER_ADD_LIQUIDITY_RETURNS_DELTA_FLAG)
+                || This.hasPermission(Hooks.AFTER_REMOVE_LIQUIDITY_RETURNS_DELTA_FLAG)
+        ) {
+            HookPermissionForbidden.selector.revertWith(address(this));
+        }
+        if (This.hasPermission(Hooks.BEFORE_SWAP_FLAG)) {
+            if (
+                uint160(address(this)) & Hooks.ALL_HOOK_MASK
+                    != uint160(_impl) & Hooks.ALL_HOOK_MASK | Hooks.AFTER_SWAP_FLAG
+            ) {
+                if (This.hasPermission(Hooks.AFTER_SWAP_FLAG)) {
+                    revert FlagsMismatch();
+                } else {
+                    // both flags match, but dev must enable AFTER_SWAP_FLAG
+                    revert MustEnableAfterSwapFlag();
+                }
+            }
+        } else if (uint160(address(this)) & Hooks.ALL_HOOK_MASK != uint160(_impl) & Hooks.ALL_HOOK_MASK) {
+            revert FlagsMismatch();
+        }
+    }
 }

contracts/middleware/MiddlewareProtectFactory.sol

@@ -9,6 +9,6 @@ contract MiddlewareProtectFactory is BaseMiddlewareFactory {
     constructor(IPoolManager _manager) BaseMiddlewareFactory(_manager) {}
 
     function _deployMiddleware(address implementation, bytes32 salt) internal override returns (address middleware) {
-        return address(new MiddlewareProtect{salt: salt}(manager, implementation));
+        middleware = address(new MiddlewareProtect{salt: salt}(manager, implementation));
     }
 }

contracts/middleware/MiddlewareRemove.sol

@@ -30,11 +30,7 @@ contract MiddlewareRemove is BaseMiddleware {
     bytes internal constant ZERO_BYTES = bytes("");
     uint256 public constant GAS_LIMIT = 10_000_000;
 
-    constructor(IPoolManager _manager, address _impl) BaseMiddleware(_manager, _impl) {
-        if (IHooks(address(this)).hasPermission(Hooks.AFTER_REMOVE_LIQUIDITY_RETURNS_DELTA_FLAG)) {
-            HookPermissionForbidden.selector.revertWith(address(this));
-        }
-    }
+    constructor(IPoolManager _manager, address _impl) BaseMiddleware(_manager, _impl) {}
 
     function beforeRemoveLiquidity(
         address sender,
@@ -87,4 +83,13 @@ contract MiddlewareRemove is BaseMiddleware {
             revert HookModifiedDeltas();
         }
     }
+
+    function _ensureValidFlags(address _impl) internal view virtual override {
+        if (uint160(address(this)) & Hooks.ALL_HOOK_MASK != uint160(_impl) & Hooks.ALL_HOOK_MASK) {
+            revert FlagsMismatch();
+        }
+        if (IHooks(address(this)).hasPermission(Hooks.AFTER_REMOVE_LIQUIDITY_RETURNS_DELTA_FLAG)) {
+            HookPermissionForbidden.selector.revertWith(address(this));
+        }
+    }
 }

contracts/middleware/MiddlewareRemoveFactory.sol

@@ -9,6 +9,6 @@ contract MiddlewareRemoveFactory is BaseMiddlewareFactory {
     constructor(IPoolManager _manager) BaseMiddlewareFactory(_manager) {}
 
     function _deployMiddleware(address implementation, bytes32 salt) internal override returns (address middleware) {
-        return address(new MiddlewareRemove{salt: salt}(manager, implementation));
+        middleware = address(new MiddlewareRemove{salt: salt}(manager, implementation));
     }
 }

test/BaseMiddlewareFactory.t.sol

@@ -27,6 +27,12 @@ contract BaseMiddlewareFactoryTest is Test, Deployers {
     Counter counter;
     address middleware;
 
+    uint160 COUNTER_FLAGS = uint160(
+        Hooks.BEFORE_INITIALIZE_FLAG | Hooks.AFTER_INITIALIZE_FLAG | Hooks.BEFORE_SWAP_FLAG | Hooks.AFTER_SWAP_FLAG
+            | Hooks.BEFORE_ADD_LIQUIDITY_FLAG | Hooks.AFTER_ADD_LIQUIDITY_FLAG | Hooks.BEFORE_REMOVE_LIQUIDITY_FLAG
+            | Hooks.AFTER_REMOVE_LIQUIDITY_FLAG | Hooks.BEFORE_DONATE_FLAG | Hooks.AFTER_DONATE_FLAG
+    );
+
     function setUp() public {
         deployFreshManagerAndRouters();
         (currency0, currency1) = deployMintAndApprove2Currencies();
@@ -36,55 +42,56 @@ contract BaseMiddlewareFactoryTest is Test, Deployers {
         token1 = TestERC20(Currency.unwrap(currency1));
 
         factory = new BaseMiddlewareFactory(manager);
-        counter = new Counter(manager);
+        counter = Counter(address(COUNTER_FLAGS));
+        vm.etch(address(counter), address(new Counter(manager)).code);
 
         token0.approve(address(router), type(uint256).max);
         token1.approve(address(router), type(uint256).max);
 
-        uint160 flags = uint160(
-            Hooks.BEFORE_INITIALIZE_FLAG | Hooks.AFTER_INITIALIZE_FLAG | Hooks.BEFORE_SWAP_FLAG | Hooks.AFTER_SWAP_FLAG
-                | Hooks.BEFORE_ADD_LIQUIDITY_FLAG | Hooks.AFTER_ADD_LIQUIDITY_FLAG | Hooks.BEFORE_REMOVE_LIQUIDITY_FLAG
-                | Hooks.AFTER_REMOVE_LIQUIDITY_FLAG | Hooks.BEFORE_DONATE_FLAG | Hooks.AFTER_DONATE_FLAG
-        );
-
         (address hookAddress, bytes32 salt) = HookMiner.find(
-            address(factory), flags, type(BaseMiddleware).creationCode, abi.encode(address(manager), address(counter))
+            address(factory),
+            COUNTER_FLAGS,
+            type(BaseMiddleware).creationCode,
+            abi.encode(address(manager), address(counter))
         );
         middleware = factory.createMiddleware(address(counter), salt);
         assertEq(hookAddress, middleware);
     }
 
     function testRevertOnSameDeployment() public {
-        uint160 flags = uint160(
-            Hooks.BEFORE_INITIALIZE_FLAG | Hooks.AFTER_INITIALIZE_FLAG | Hooks.BEFORE_SWAP_FLAG | Hooks.AFTER_SWAP_FLAG
-                | Hooks.BEFORE_ADD_LIQUIDITY_FLAG | Hooks.AFTER_ADD_LIQUIDITY_FLAG | Hooks.BEFORE_REMOVE_LIQUIDITY_FLAG
-                | Hooks.AFTER_REMOVE_LIQUIDITY_FLAG | Hooks.BEFORE_DONATE_FLAG | Hooks.AFTER_DONATE_FLAG
-        );
         (address hookAddress, bytes32 salt) = HookMiner.find(
-            address(factory), flags, type(BaseMiddleware).creationCode, abi.encode(address(manager), address(counter))
+            address(factory),
+            COUNTER_FLAGS,
+            type(BaseMiddleware).creationCode,
+            abi.encode(address(manager), address(counter))
         );
         factory.createMiddleware(address(counter), salt);
         // second deployment should revert
-        vm.expectRevert(bytes(""));
+        vm.expectRevert(ZERO_BYTES);
         factory.createMiddleware(address(counter), salt);
     }
 
     function testRevertOnIncorrectFlags() public {
-        Counter counter2 = new Counter(manager);
-        uint160 flags = uint160(Hooks.BEFORE_INITIALIZE_FLAG);
+        Counter counter2 = Counter(address(COUNTER_FLAGS));
+        vm.etch(address(counter), address(new Counter(manager)).code);
+        uint160 incorrectFlags = uint160(Hooks.BEFORE_INITIALIZE_FLAG);
 
         (address hookAddress, bytes32 salt) = HookMiner.find(
-            address(factory), flags, type(BaseMiddleware).creationCode, abi.encode(address(manager), address(counter2))
+            address(factory),
+            incorrectFlags,
+            type(BaseMiddleware).creationCode,
+            abi.encode(address(manager), address(counter2))
         );
         address implementation = address(counter2);
-        vm.expectRevert(abi.encodePacked(bytes16(Hooks.HookAddressNotValid.selector), hookAddress));
+        vm.expectRevert(BaseMiddleware.FlagsMismatch.selector);
         factory.createMiddleware(implementation, salt);
     }
 
     function testRevertOnIncorrectFlagsMined() public {
-        Counter counter2 = new Counter(manager);
+        Counter counter2 = Counter(address(COUNTER_FLAGS));
+        vm.etch(address(counter), address(new Counter(manager)).code);
         address implementation = address(counter2);
-        vm.expectRevert(); // HookAddressNotValid
+        vm.expectRevert(BaseMiddleware.FlagsMismatch.selector);
         factory.createMiddleware(implementation, bytes32("who needs to mine a salt?"));
     }
 
@@ -116,7 +123,7 @@ contract BaseMiddlewareFactoryTest is Test, Deployers {
         assertEq(counterProxy.afterSwapCount(id), 1);
 
         // counter does not store data itself
-        assertEq(counter.lastHookData(), bytes(""));
+        assertEq(counter.lastHookData(), ZERO_BYTES);
         assertEq(counter.beforeSwapCount(id), 0);
         assertEq(counter.afterSwapCount(id), 0);
     }