start frontrun test hook

contracts/middleware/BaseMiddleware.sol

@@ -8,7 +8,6 @@ import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
 import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
 import {BeforeSwapDelta} from "@uniswap/v4-core/src/types/BeforeSwapDelta.sol";
 import {Proxy} from "@openzeppelin/contracts/proxy/Proxy.sol";
-import {console} from "../../lib/forge-std/src/console.sol";
 
 contract BaseMiddleware is Proxy {
     /// @notice The address of the pool manager
@@ -21,11 +20,10 @@ contract BaseMiddleware is Proxy {
     }
 
     function _implementation() internal view override returns (address) {
-        console.logAddress(implementation);
         return implementation;
     }
 
-    // yo how do i remove this warning
+    // yo i wanna delete this function but how do i remove this warning
     receive() external payable {
         _delegate(_implementation());
     }

contracts/middleware/MiddlewareProtect.sol

@@ -81,10 +81,11 @@ contract MiddlewareProtect is BaseMiddleware {
         address,
         PoolKey calldata,
         IPoolManager.ModifyLiquidityParams calldata,
+        BalanceDelta,
         bytes calldata
     ) external returns (bytes4, BalanceDelta) {
         console.log("afterRemoveLiquidity middleware");
         implementation.delegatecall{gas: gasLimit}(msg.data);
         return (BaseHook.afterRemoveLiquidity.selector, BalanceDeltaLibrary.ZERO_DELTA);
     }
-}
+}

contracts/middleware/MiddlewareProtectFactory.sol

@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import {IMiddlewareFactory} from "../interfaces/IMiddlewareFactory.sol";
+import {MiddlewareProtect} from "./MiddlewareProtect.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+
+contract MiddlewareProtectFactory is IMiddlewareFactory {
+    mapping(address => address) private _implementations;
+
+    IPoolManager public immutable poolManager;
+
+    constructor(IPoolManager _poolManager) {
+        poolManager = _poolManager;
+    }
+
+    function getImplementation(address middleware) external view override returns (address implementation) {
+        return _implementations[middleware];
+    }
+
+    function createMiddleware(address implementation, bytes32 salt) external override returns (address middleware) {
+        middleware = address(new MiddlewareProtect{salt: salt}(poolManager, implementation));
+        _implementations[middleware] = implementation;
+        emit MiddlewareCreated(implementation, middleware);
+    }
+}

test/MiddlewareProtectFactory.t.sol

@@ -0,0 +1,101 @@
+// 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 {FeeTakingLite} from "./middleware/FeeTakingLite.sol";
+import {MiddlewareProtect} from "../contracts/middleware/MiddlewareProtect.sol";
+import {PoolManager} from "@uniswap/v4-core/src/PoolManager.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.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 {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 {HooksRevert} from "./middleware/HooksRevert.sol";
+import {HooksOutOfGas} from "./middleware/HooksOutOfGas.sol";
+import {MiddlewareProtectFactory} from "./../contracts/middleware/MiddlewareProtectFactory.sol";
+import {HookMiner} from "./utils/HookMiner.sol";
+
+contract MiddlewareProtectFactoryTest is Test, Deployers {
+    using PoolIdLibrary for PoolKey;
+    using StateLibrary for IPoolManager;
+
+    uint160 constant SQRT_RATIO_10_1 = 250541448375047931186413801569;
+
+    address constant TREASURY = address(0x1234567890123456789012345678901234567890);
+    uint128 private constant TOTAL_BIPS = 10000;
+
+    HookEnabledSwapRouter router;
+    TestERC20 token0;
+    TestERC20 token1;
+    PoolId id;
+
+    MiddlewareProtectFactory factory;
+
+    function setUp() public {
+        deployFreshManagerAndRouters();
+        (currency0, currency1) = deployMintAndApprove2Currencies();
+
+        router = new HookEnabledSwapRouter(manager);
+        token0 = TestERC20(Currency.unwrap(currency0));
+        token1 = TestERC20(Currency.unwrap(currency1));
+
+        token0.approve(address(router), type(uint256).max);
+        token1.approve(address(router), type(uint256).max);
+
+        factory = new MiddlewareProtectFactory(manager);
+    }
+
+    function testVariousProtectFactory() public {
+        FeeTakingLite feeTakingLite = new FeeTakingLite(manager);
+        uint160 flags =
+            uint160(Hooks.AFTER_SWAP_FLAG | Hooks.AFTER_SWAP_RETURNS_DELTA_FLAG | Hooks.BEFORE_REMOVE_LIQUIDITY_FLAG);
+        (address hookAddress, bytes32 salt) = HookMiner.find(
+            address(factory),
+            flags,
+            type(MiddlewareProtect).creationCode,
+            abi.encode(address(manager), address(feeTakingLite))
+        );
+        testOn(address(feeTakingLite), salt);
+
+        HooksRevert hooksRevert = new HooksRevert(manager);
+        flags = uint160(Hooks.BEFORE_REMOVE_LIQUIDITY_FLAG | Hooks.AFTER_REMOVE_LIQUIDITY_FLAG);
+        (hookAddress, salt) = HookMiner.find(
+            address(factory),
+            flags,
+            type(MiddlewareProtect).creationCode,
+            abi.encode(address(manager), address(hooksRevert))
+        );
+        testOn(address(hooksRevert), salt);
+
+        HooksOutOfGas hooksOutOfGas = new HooksOutOfGas(manager);
+        flags = uint160(Hooks.BEFORE_REMOVE_LIQUIDITY_FLAG | Hooks.AFTER_REMOVE_LIQUIDITY_FLAG);
+        (hookAddress, salt) = HookMiner.find(
+            address(factory),
+            flags,
+            type(MiddlewareProtect).creationCode,
+            abi.encode(address(manager), address(hooksOutOfGas))
+        );
+        testOn(address(hooksOutOfGas), salt);
+    }
+
+    // creates a middleware on an implementation
+    function testOn(address implementation, bytes32 salt) internal {
+        address hookAddress = factory.createMiddleware(implementation, salt);
+        MiddlewareProtect middlewareProtect = MiddlewareProtect(payable(hookAddress));
+
+        (key, id) = initPoolAndAddLiquidity(
+            currency0, currency1, IHooks(address(middlewareProtect)), 3000, SQRT_PRICE_1_1, ZERO_BYTES
+        );
+
+        removeLiquidity(currency0, currency1, IHooks(address(middlewareProtect)), 3000, SQRT_PRICE_1_1, ZERO_BYTES);
+
+        assertEq(factory.getImplementation(hookAddress), implementation);
+    }
+}

test/MiddlewareRemove.t.sol

@@ -5,7 +5,6 @@ import {Test} from "forge-std/Test.sol";
 import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol";
 import {FeeTakingLite} from "./middleware/FeeTakingLite.sol";
 import {MiddlewareRemove} from "../contracts/middleware/MiddlewareRemove.sol";
-import {MiddlewareRemoveImplementation} from "./shared/implementation/MiddlewareRemoveImplementation.sol";
 import {PoolManager} from "@uniswap/v4-core/src/PoolManager.sol";
 import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
 import {Deployers} from "@uniswap/v4-core/test/utils/Deployers.sol";
@@ -70,8 +69,7 @@ contract MiddlewareRemoveTest is Test, Deployers {
         MiddlewareRemove middlewareRemove = MiddlewareRemove(payable(address(nonce << 20 | flags)));
         nonce++;
         vm.record();
-        MiddlewareRemoveImplementation impl =
-            new MiddlewareRemoveImplementation(manager, implementation, middlewareRemove);
+        MiddlewareRemove impl = new MiddlewareRemove(manager, implementation);
         (, bytes32[] memory writes) = vm.accesses(address(impl));
         vm.etch(address(middlewareRemove), address(impl).code);
         unchecked {

test/MiddlewareRemoveFactory.t.sol

@@ -5,7 +5,6 @@ import {Test} from "forge-std/Test.sol";
 import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol";
 import {FeeTakingLite} from "./middleware/FeeTakingLite.sol";
 import {MiddlewareRemove} from "../contracts/middleware/MiddlewareRemove.sol";
-import {MiddlewareRemoveImplementation} from "./shared/implementation/MiddlewareRemoveImplementation.sol";
 import {PoolManager} from "@uniswap/v4-core/src/PoolManager.sol";
 import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
 import {Deployers} from "@uniswap/v4-core/test/utils/Deployers.sol";

test/middleware/HookFrontrun.sol

@@ -0,0 +1,61 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import {BaseHook} from "./../../contracts/BaseHook.sol";
+import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
+import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
+import {Currency, CurrencyLibrary} from "@uniswap/v4-core/src/types/Currency.sol";
+import {SafeCast} from "@uniswap/v4-core/src/libraries/SafeCast.sol";
+
+abstract contract FeeTaker is BaseHook {
+    using SafeCast for uint256;
+
+    bytes internal constant ZERO_BYTES = bytes("");
+
+    constructor(IPoolManager _poolManager) BaseHook(_poolManager) {}
+
+    function getHookPermissions() public pure virtual override returns (Hooks.Permissions memory) {
+        return Hooks.Permissions({
+            beforeInitialize: false,
+            afterInitialize: false,
+            beforeAddLiquidity: false,
+            afterAddLiquidity: false,
+            beforeRemoveLiquidity: false,
+            afterRemoveLiquidity: false,
+            beforeSwap: true,
+            afterSwap: true,
+            beforeDonate: false,
+            afterDonate: false,
+            beforeSwapReturnDelta: false,
+            afterSwapReturnDelta: false,
+            afterAddLiquidityReturnDelta: false,
+            afterRemoveLiquidityReturnDelta: false
+        });
+    }
+
+    function afterSwap(
+        address sender,
+        PoolKey calldata key,
+        IPoolManager.SwapParams calldata params,
+        BalanceDelta delta,
+        bytes calldata hookData
+    ) external override onlyByManager returns (bytes4, int128) {
+        //(Currency currencyUnspecified, amountUnspecified) = key.getUnspecified(params);
+
+        // fee will be in the unspecified token of the swap
+        bool currency0Specified = (params.amountSpecified < 0 == params.zeroForOne);
+        (Currency currencyUnspecified, int128 amountUnspecified) =
+            (currency0Specified) ? (key.currency1, delta.amount1()) : (key.currency0, delta.amount0());
+        // if exactOutput swap, get the absolute output amount
+        if (amountUnspecified < 0) amountUnspecified = -amountUnspecified;
+
+        uint256 feeAmount = _feeAmount(amountUnspecified);
+        // mint ERC6909 instead of take to avoid edge case where PM doesn't have enough balance
+        manager.mint(address(this), CurrencyLibrary.toId(currencyUnspecified), feeAmount);
+
+        (bytes4 selector, int128 amount) = _afterSwap(sender, key, params, delta, hookData);
+        return (selector, feeAmount.toInt128() + amount);
+    }
+}