From 5f77dcadfdde993fce2100491bd4bc0600cae434 Mon Sep 17 00:00:00 2001 From: Junion <69495294+Jun1on@users.noreply.github.com> Date: Tue, 2 Jul 2024 01:31:58 -0400 Subject: [PATCH] start frontrun test hook --- contracts/middleware/BaseMiddleware.sol | 4 +- contracts/middleware/MiddlewareProtect.sol | 3 +- .../middleware/MiddlewareProtectFactory.sol | 26 +++++ test/MiddlewareProtectFactory.t.sol | 101 ++++++++++++++++++ test/MiddlewareRemove.t.sol | 4 +- test/MiddlewareRemoveFactory.t.sol | 1 - test/middleware/HookFrontrun.sol | 61 +++++++++++ .../MiddlewareRemoveImplementation.sol | 19 ---- 8 files changed, 192 insertions(+), 27 deletions(-) create mode 100644 contracts/middleware/MiddlewareProtectFactory.sol create mode 100644 test/MiddlewareProtectFactory.t.sol create mode 100644 test/middleware/HookFrontrun.sol delete mode 100644 test/shared/implementation/MiddlewareRemoveImplementation.sol diff --git a/contracts/middleware/BaseMiddleware.sol b/contracts/middleware/BaseMiddleware.sol index ce4a7004..ebe5c609 100644 --- a/contracts/middleware/BaseMiddleware.sol +++ b/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()); } diff --git a/contracts/middleware/MiddlewareProtect.sol b/contracts/middleware/MiddlewareProtect.sol index 0ee8a20d..1e800878 100644 --- a/contracts/middleware/MiddlewareProtect.sol +++ b/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); } -} \ No newline at end of file +} diff --git a/contracts/middleware/MiddlewareProtectFactory.sol b/contracts/middleware/MiddlewareProtectFactory.sol new file mode 100644 index 00000000..3287a629 --- /dev/null +++ b/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); + } +} diff --git a/test/MiddlewareProtectFactory.t.sol b/test/MiddlewareProtectFactory.t.sol new file mode 100644 index 00000000..b5a472ac --- /dev/null +++ b/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); + } +} diff --git a/test/MiddlewareRemove.t.sol b/test/MiddlewareRemove.t.sol index 6bafa59b..0ecae3f5 100644 --- a/test/MiddlewareRemove.t.sol +++ b/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 { diff --git a/test/MiddlewareRemoveFactory.t.sol b/test/MiddlewareRemoveFactory.t.sol index d769e077..dea6b7ae 100644 --- a/test/MiddlewareRemoveFactory.t.sol +++ b/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"; diff --git a/test/middleware/HookFrontrun.sol b/test/middleware/HookFrontrun.sol new file mode 100644 index 00000000..615dfb8b --- /dev/null +++ b/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); + } +} diff --git a/test/shared/implementation/MiddlewareRemoveImplementation.sol b/test/shared/implementation/MiddlewareRemoveImplementation.sol deleted file mode 100644 index 33a190ac..00000000 --- a/test/shared/implementation/MiddlewareRemoveImplementation.sol +++ /dev/null @@ -1,19 +0,0 @@ -// SPDX-License-Identifier: UNLICENSED -pragma solidity ^0.8.19; - -import {BaseHook} from "../../../contracts/BaseHook.sol"; -import {MiddlewareRemove} from "../../../contracts/middleware/MiddlewareRemove.sol"; -import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol"; -import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol"; -import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol"; - -contract MiddlewareRemoveImplementation is MiddlewareRemove { - constructor(IPoolManager _poolManager, address _implementation, MiddlewareRemove addressToEtch) - MiddlewareRemove(_poolManager, _implementation) - { - //Hooks.validateHookPermissions(addressToEtch, getHookPermissions()); - } - - // make this a no-op in testing - //function validateHookAddress(BaseHook _this) internal pure override {} -}