create outline

contracts/middleware/MiddlewareProtect.sol

@@ -0,0 +1,56 @@
+// 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 {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";
+import {MiddlewareRemove} from "./MiddlewareRemove.sol";
+import {BeforeSwapDelta, BeforeSwapDeltaLibrary} from "@uniswap/v4-core/src/types/BeforeSwapDelta.sol";
+
+contract MiddlewareProtect is MiddlewareRemove {
+    using CustomRevert for bytes4;
+    using Hooks for IHooks;
+    using StateLibrary for IPoolManager;
+    using PoolIdLibrary for PoolKey;
+
+    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)
+                || middleware.hasPermission(Hooks.AFTER_REMOVE_LIQUIDITY_RETURNS_DELTA_FLAG)
+        ) {
+            HookPermissionForbidden.selector.revertWith(address(this));
+        }
+    }
+
+    // function beforeSwap(address, PoolKey calldata, IPoolManager.SwapParams calldata, bytes calldata)
+    //     external
+    //     returns (bytes4, BeforeSwapDelta, uint24)
+    // {
+    //     return (BaseHook.beforeSwap.selector, BeforeSwapDeltaLibrary.ZERO_DELTA, 0);
+    // }
+
+    function beforeAddLiquidity(
+        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, sender, key, params, hookData)
+        );
+        return BaseHook.beforeAddLiquidity.selector;
+    }
+}

contracts/middleware/MiddlewareProtectFactory.sol

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import {BaseMiddlewareFactory} from "./BaseMiddlewareFactory.sol";
+import {MiddlewareProtect} from "./MiddlewareProtect.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+
+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));
+    }
+}

test/MiddlewareProtectFactory.t.sol

@@ -0,0 +1,202 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import {Test} from "forge-std/Test.sol";
+import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol";
+import {HooksFrontrun} from "./middleware/HooksFrontrun.sol";
+import {MiddlewareRemove} from "../contracts/middleware/MiddlewareRemove.sol";
+import {MiddlewareProtect} from "../contracts/middleware/MiddlewareProtect.sol";
+import {Deployers} from "@uniswap/v4-core/test/utils/Deployers.sol";
+import {TestERC20} from "@uniswap/v4-core/src/test/TestERC20.sol";
+import {CurrencyLibrary, Currency} from "@uniswap/v4-core/src/types/Currency.sol";
+import {PoolId, PoolIdLibrary} from "@uniswap/v4-core/src/types/PoolId.sol";
+import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
+import {HookEnabledSwapRouter} from "./utils/HookEnabledSwapRouter.sol";
+import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
+import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
+import {console} from "../../../lib/forge-std/src/console.sol";
+import {HooksRevert} from "./middleware/HooksRevert.sol";
+import {HooksOutOfGas} from "./middleware/HooksOutOfGas.sol";
+import {MiddlewareProtectFactory} from "./../contracts/middleware/MiddlewareProtectFactory.sol";
+import {HookMiner} from "./utils/HookMiner.sol";
+import {HooksReturnDeltas} from "./middleware/HooksReturnDeltas.sol";
+import {Counter} from "./middleware/Counter.sol";
+import {SafeCallback} from "./../contracts/base/SafeCallback.sol";
+
+contract MiddlewareProtectFactoryTest is Test, Deployers {
+    HookEnabledSwapRouter router;
+    TestERC20 token0;
+    TestERC20 token1;
+
+    MiddlewareProtectFactory factory;
+    Counter counter;
+    address middleware;
+    HooksFrontrun hooksFrontrun;
+
+    function setUp() public {
+        deployFreshManagerAndRouters();
+        (currency0, currency1) = deployMintAndApprove2Currencies();
+
+        router = new HookEnabledSwapRouter(manager);
+        token0 = TestERC20(Currency.unwrap(currency0));
+        token1 = TestERC20(Currency.unwrap(currency1));
+
+        factory = new MiddlewareProtectFactory(manager);
+        counter = new Counter(manager);
+
+        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(MiddlewareProtect).creationCode,
+            abi.encode(address(manager), address(counter))
+        );
+        middleware = factory.createMiddleware(address(counter), salt);
+        assertEq(hookAddress, middleware);
+
+        flags = uint160(Hooks.BEFORE_SWAP_FLAG | Hooks.AFTER_SWAP_FLAG);
+        hooksFrontrun = HooksFrontrun(address(uint160(flags)));
+        HooksFrontrun impl = new HooksFrontrun(manager);
+        vm.etch(address(hooksFrontrun), address(impl).code);
+    }
+
+    function testRevertOnDeltas() public {
+        HooksReturnDeltas hooksReturnDeltas = new HooksReturnDeltas(manager);
+        uint160 flags = uint160(Hooks.AFTER_SWAP_FLAG | Hooks.AFTER_SWAP_RETURNS_DELTA_FLAG);
+
+        (address hookAddress, bytes32 salt) = HookMiner.find(
+            address(factory),
+            flags,
+            type(MiddlewareProtect).creationCode,
+            abi.encode(address(manager), address(hooksReturnDeltas))
+        );
+        address implementation = address(hooksReturnDeltas);
+        vm.expectRevert(abi.encodePacked(bytes16(MiddlewareRemove.HookPermissionForbidden.selector), hookAddress));
+        factory.createMiddleware(implementation, salt);
+    }
+
+    function testFrontrun() public {
+        return;
+        (PoolKey memory key,) =
+            initPoolAndAddLiquidity(currency0, currency1, IHooks(address(0)), 100, SQRT_PRICE_1_1, ZERO_BYTES);
+        BalanceDelta swapDelta = swap(key, true, 0.001 ether, ZERO_BYTES);
+
+        (key,) = initPoolAndAddLiquidity(
+            currency0, currency1, IHooks(address(hooksFrontrun)), 100, SQRT_PRICE_1_1, ZERO_BYTES
+        );
+        BalanceDelta swapDelta2 = swap(key, true, 0.001 ether, ZERO_BYTES);
+
+        // while both swaps are in the same pool, the second swap is more expensive
+        assertEq(swapDelta.amount1(), swapDelta2.amount1());
+        assertTrue(abs(swapDelta.amount0()) < abs(swapDelta2.amount0()));
+        assertTrue(manager.balanceOf(address(hooksFrontrun), CurrencyLibrary.toId(key.currency0)) > 0);
+    }
+
+    function testRevertOnFrontrun() public {
+        return;
+        uint160 flags = uint160(Hooks.BEFORE_SWAP_FLAG | Hooks.AFTER_SWAP_FLAG);
+
+        (address hookAddress, bytes32 salt) = HookMiner.find(
+            address(factory),
+            flags,
+            type(MiddlewareProtect).creationCode,
+            abi.encode(address(manager), address(hooksFrontrun))
+        );
+        address implementation = address(hooksFrontrun);
+        address hookAddressCreated = factory.createMiddleware(implementation, salt);
+        assertEq(hookAddressCreated, hookAddress);
+        MiddlewareProtect middlewareProtect = MiddlewareProtect(payable(hookAddress));
+
+        (key,) = initPoolAndAddLiquidity(
+            currency0, currency1, IHooks(address(middlewareProtect)), 100, SQRT_PRICE_1_1, ZERO_BYTES
+        );
+        //vm.expectRevert(MiddlewareProtect.ActionBetweenHook.selector);
+        swap(key, true, 0.001 ether, ZERO_BYTES);
+    }
+
+    function abs(int256 x) internal pure returns (uint256) {
+        return x >= 0 ? uint256(x) : uint256(-x);
+    }
+
+    // from BaseMiddlewareFactory.t.sol
+    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
+        );
+        (, bytes32 salt) = HookMiner.find(
+            address(factory),
+            flags,
+            type(MiddlewareProtect).creationCode,
+            abi.encode(address(manager), address(counter))
+        );
+        factory.createMiddleware(address(counter), salt);
+        // second deployment should revert
+        vm.expectRevert(bytes(""));
+        factory.createMiddleware(address(counter), salt);
+    }
+
+    function testRevertOnIncorrectFlags() public {
+        Counter counter2 = new Counter(manager);
+        uint160 flags = uint160(Hooks.BEFORE_INITIALIZE_FLAG);
+
+        (address hookAddress, bytes32 salt) = HookMiner.find(
+            address(factory),
+            flags,
+            type(MiddlewareProtect).creationCode,
+            abi.encode(address(manager), address(counter2))
+        );
+        address implementation = address(counter2);
+        vm.expectRevert(abi.encodePacked(bytes16(Hooks.HookAddressNotValid.selector), hookAddress));
+        factory.createMiddleware(implementation, salt);
+    }
+
+    function testRevertOnIncorrectFlagsMined() public {
+        Counter counter2 = new Counter(manager);
+        address implementation = address(counter2);
+        vm.expectRevert(); // HookAddressNotValid
+        factory.createMiddleware(implementation, bytes32("who needs to mine a salt?"));
+    }
+
+    function testRevertOnIncorrectCaller() public {
+        vm.expectRevert(SafeCallback.NotManager.selector);
+        counter.afterDonate(address(this), key, 0, 0, ZERO_BYTES);
+    }
+
+    function testCounters() public {
+        (PoolKey memory key, PoolId id) =
+            initPoolAndAddLiquidity(currency0, currency1, IHooks(middleware), 3000, SQRT_PRICE_1_1, ZERO_BYTES);
+
+        Counter counterProxy = Counter(middleware);
+        assertEq(counterProxy.beforeInitializeCount(id), 1);
+        assertEq(counterProxy.afterInitializeCount(id), 1);
+        assertEq(counterProxy.beforeSwapCount(id), 0);
+        assertEq(counterProxy.afterSwapCount(id), 0);
+        assertEq(counterProxy.beforeAddLiquidityCount(id), 1);
+        assertEq(counterProxy.afterAddLiquidityCount(id), 1);
+        assertEq(counterProxy.beforeRemoveLiquidityCount(id), 0);
+        assertEq(counterProxy.afterRemoveLiquidityCount(id), 0);
+        assertEq(counterProxy.beforeDonateCount(id), 0);
+        assertEq(counterProxy.afterDonateCount(id), 0);
+
+        assertEq(counterProxy.lastHookData(), ZERO_BYTES);
+        swap(key, true, 1, bytes("hi"));
+        assertEq(counterProxy.lastHookData(), bytes("hi"));
+        assertEq(counterProxy.beforeSwapCount(id), 1);
+        assertEq(counterProxy.afterSwapCount(id), 1);
+
+        // counter does not store data itself
+        assertEq(counter.lastHookData(), bytes(""));
+        assertEq(counter.beforeSwapCount(id), 0);
+        assertEq(counter.afterSwapCount(id), 0);
+    }
+}

test/MiddlewareRemoveFactory.t.sol

@@ -16,8 +16,8 @@ import {StateLibrary} from "@uniswap/v4-core/src/libraries/StateLibrary.sol";
 import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
 import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
 import {console} from "../../../lib/forge-std/src/console.sol";
-import {RemoveReverts} from "./middleware/RemoveReverts.sol";
-import {RemoveOutOfGas} from "./middleware/RemoveOutOfGas.sol";
+import {HooksRevert} from "./middleware/HooksRevert.sol";
+import {HooksOutOfGas} from "./middleware/HooksOutOfGas.sol";
 import {MiddlewareRemoveFactory} from "./../contracts/middleware/MiddlewareRemoveFactory.sol";
 import {HookMiner} from "./utils/HookMiner.sol";
 import {SafeCallback} from "./../contracts/base/SafeCallback.sol";
@@ -156,24 +156,27 @@ contract MiddlewareRemoveFactoryTest is Test, Deployers {
         );
         testOn(address(counter), salt);
 
-        RemoveReverts removeReverts = new RemoveReverts(manager);
-        flags = uint160(Hooks.BEFORE_REMOVE_LIQUIDITY_FLAG | Hooks.AFTER_REMOVE_LIQUIDITY_FLAG);
+        HooksRevert hooksRevert = new HooksRevert(manager);
+        flags = uint160(
+            Hooks.BEFORE_SWAP_FLAG | Hooks.AFTER_SWAP_FLAG | Hooks.BEFORE_REMOVE_LIQUIDITY_FLAG
+                | Hooks.AFTER_REMOVE_LIQUIDITY_FLAG
+        );
         (hookAddress, salt) = HookMiner.find(
             address(factory),
             flags,
             type(MiddlewareRemove).creationCode,
-            abi.encode(address(manager), address(removeReverts))
+            abi.encode(address(manager), address(hooksRevert))
         );
-        testOn(address(removeReverts), salt);
+        testOn(address(hooksRevert), salt);
 
-        RemoveOutOfGas removeOutOfGas = new RemoveOutOfGas(manager);
+        HooksOutOfGas hooksOutOfGas = new HooksOutOfGas(manager);
         (hookAddress, salt) = HookMiner.find(
             address(factory),
             flags,
             type(MiddlewareRemove).creationCode,
-            abi.encode(address(manager), address(removeOutOfGas))
+            abi.encode(address(manager), address(hooksOutOfGas))
         );
-        testOn(address(removeOutOfGas), salt);
+        testOn(address(hooksOutOfGas), salt);
     }
 
     // creates a middleware on an implementation
@@ -265,5 +268,9 @@ contract MiddlewareRemoveFactoryTest is Test, Deployers {
         assertEq(counter.lastHookData(), bytes(""));
         assertEq(counter.beforeSwapCount(id), 0);
         assertEq(counter.afterSwapCount(id), 0);
+
+        removeLiquidity(currency0, currency1, IHooks(middleware), 3000, SQRT_PRICE_1_1, ZERO_BYTES);
+        assertEq(counterProxy.beforeRemoveLiquidityCount(id), 1);
+        assertEq(counterProxy.afterRemoveLiquidityCount(id), 1);
     }
 }