Middleware remove

.github/workflows/lint.yml

@@ -21,5 +21,5 @@ jobs:
         with:
           version: nightly
 
-      - name: Run tests
+      - name: Check format
         run: forge fmt --check

contracts/interfaces/IMiddlewareFactory.sol

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: UNLICENSED
-pragma solidity ^0.8.19;
+pragma solidity ^0.8.24;
 
 interface IMiddlewareFactory {
     event MiddlewareCreated(address implementation, address middleware);

contracts/middleware/BaseMiddleware.sol

@@ -11,11 +11,14 @@ import {Proxy} from "@openzeppelin/contracts/proxy/Proxy.sol";
 
 contract BaseMiddleware is Proxy {
     /// @notice The address of the pool manager
-    IPoolManager public immutable poolManager;
+    IPoolManager public immutable manager;
     address public immutable implementation;
 
-    constructor(IPoolManager _poolManager, address _impl) {
-        poolManager = _poolManager;
+    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

@@ -1,20 +1,17 @@
 // SPDX-License-Identifier: UNLICENSED
-pragma solidity ^0.8.19;
+pragma solidity ^0.8.24;
 
 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;
 
-    IPoolManager public immutable poolManager;
+    IPoolManager public immutable manager;
 
-    constructor(IPoolManager _poolManager) {
-        poolManager = _poolManager;
+    constructor(IPoolManager _manager) {
+        manager = _manager;
     }
 
     function getImplementation(address middleware) external view override returns (address implementation) {
@@ -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}(poolManager, implementation));
+        middleware = address(new BaseMiddleware{salt: salt}(manager, implementation));
     }
 }

contracts/middleware/MiddlewareRemove.sol

@@ -0,0 +1,95 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.24;
+
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
+import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
+import {Proxy} from "@openzeppelin/contracts/proxy/Proxy.sol";
+import {BaseMiddleware} from "./BaseMiddleware.sol";
+import {BaseHook} from "../BaseHook.sol";
+import {BalanceDeltaLibrary} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
+import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol";
+import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
+import {CustomRevert} from "@uniswap/v4-core/src/libraries/CustomRevert.sol";
+import {NonZeroDeltaCount} from "@uniswap/v4-core/src/libraries/NonZeroDeltaCount.sol";
+import {IExttload} from "@uniswap/v4-core/src/interfaces/IExttload.sol";
+import {StateLibrary} from "@uniswap/v4-core/src/libraries/StateLibrary.sol";
+import {PoolIdLibrary} from "@uniswap/v4-core/src/types/PoolId.sol";
+
+contract MiddlewareRemove is BaseMiddleware {
+    using CustomRevert for bytes4;
+    using Hooks for IHooks;
+    using StateLibrary for IPoolManager;
+    using PoolIdLibrary for PoolKey;
+
+    error HookPermissionForbidden(address hooks);
+    error HookModifiedPrice();
+    error HookModifiedDeltas();
+    error FailedImplementationCall();
+
+    bytes internal constant ZERO_BYTES = bytes("");
+    uint256 public constant GAS_LIMIT = 10_000_000;
+
+    constructor(IPoolManager _manager, address _impl) BaseMiddleware(_manager, _impl) {}
+
+    function beforeRemoveLiquidity(
+        address sender,
+        PoolKey calldata key,
+        IPoolManager.ModifyLiquidityParams calldata params,
+        bytes calldata hookData
+    ) external returns (bytes4) {
+        address(this).delegatecall{gas: GAS_LIMIT}(abi.encodeWithSelector(this._callAndEnsurePrice.selector, msg.data));
+        return BaseHook.beforeRemoveLiquidity.selector;
+    }
+
+    function _callAndEnsurePrice(bytes calldata data) external {
+        (, PoolKey memory key,,) = abi.decode(data[4:], (address, PoolKey, IPoolManager.ModifyLiquidityParams, bytes));
+
+        (uint160 priceBefore,,,) = manager.getSlot0(key.toId());
+        (bool success,) = address(implementation).delegatecall(data);
+        if (!success) {
+            revert FailedImplementationCall();
+        }
+        (uint160 priceAfter,,,) = manager.getSlot0(key.toId());
+        if (priceAfter != priceBefore) {
+            // purpousely revert to cause the whole hook to reset
+            revert HookModifiedPrice();
+        }
+    }
+
+    function afterRemoveLiquidity(
+        address sender,
+        PoolKey calldata key,
+        IPoolManager.ModifyLiquidityParams calldata params,
+        BalanceDelta delta,
+        bytes calldata hookData
+    ) external returns (bytes4, BalanceDelta) {
+        address(this).delegatecall{gas: GAS_LIMIT}(
+            abi.encodeWithSelector(this._callAndEnsureZeroDeltas.selector, msg.data)
+        );
+        return (BaseHook.afterRemoveLiquidity.selector, BalanceDeltaLibrary.ZERO_DELTA);
+    }
+
+    function _callAndEnsureZeroDeltas(bytes calldata data) external {
+        bytes32 slot = bytes32(NonZeroDeltaCount.NONZERO_DELTA_COUNT_SLOT);
+        uint256 countBefore = uint256(IExttload(address(manager)).exttload(slot));
+        (bool success,) = address(implementation).delegatecall(data);
+        if (!success) {
+            revert FailedImplementationCall();
+        }
+        uint256 countAfter = uint256(IExttload(address(manager)).exttload(slot));
+        if (countAfter != countBefore) {
+            // purpousely revert to cause the whole hook to reset
+            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

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import {BaseMiddlewareFactory} from "./BaseMiddlewareFactory.sol";
+import {MiddlewareRemove} from "./MiddlewareRemove.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+
+contract MiddlewareRemoveFactory is BaseMiddlewareFactory {
+    constructor(IPoolManager _manager) BaseMiddlewareFactory(_manager) {}
+
+    function _deployMiddleware(address implementation, bytes32 salt) internal override returns (address middleware) {
+        middleware = address(new MiddlewareRemove{salt: salt}(manager, implementation));
+    }
+}

foundry.toml

@@ -6,6 +6,7 @@ optimizer_runs = 1000000
 ffi = true
 fs_permissions = [{ access = "read-write", path = ".forge-snapshots/"}]
 evm_version = "cancun"
+gas_limit = "3000000000"
 
 [profile.ci]
 fuzz_runs = 100000

test/BaseMiddlewareFactory.t.sol

@@ -25,9 +25,14 @@ contract BaseMiddlewareFactoryTest is Test, Deployers {
 
     BaseMiddlewareFactory factory;
     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();
@@ -37,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?"));
     }
 
@@ -117,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);
     }