// SPDX-License-Identifier: MIT
pragma solidity 0.8.26;

import {Test} from "forge-std/Test.sol";
import {PoolManager} from "v4-core/src/PoolManager.sol";
import {IPoolManager} from "v4-core/src/interfaces/IPoolManager.sol";
import {IHooks} from "v4-core/src/interfaces/IHooks.sol";
import {IUnlockCallback} from "v4-core/src/interfaces/callback/IUnlockCallback.sol";
import {Currency} from "v4-core/src/types/Currency.sol";
import {PoolKey} from "v4-core/src/types/PoolKey.sol";
import {SwapParams, ModifyLiquidityParams} from "v4-core/src/types/PoolOperation.sol";
import {BalanceDelta} from "v4-core/src/types/BalanceDelta.sol";
import {TransientStateLibrary} from "v4-core/src/libraries/TransientStateLibrary.sol";
import {TickMath} from "v4-core/src/libraries/TickMath.sol";
import {EHS} from "src/EHS.sol";
import {EHSHook} from "src/EHSHook.sol";

interface IERC20Like {
    function transfer(address, uint256) external returns (bool);
    function transferFrom(address, address, uint256) external returns (bool);
    function balanceOf(address) external view returns (uint256);
    function approve(address, uint256) external returns (bool);
}

/// @dev Plain ERC-20 standing in for IMD at its mainnet address.
contract PlainIMD {
    string public constant symbol = "IMD";
    uint8 public constant decimals = 18;
    mapping(address => uint256) public balanceOf;
    mapping(address => mapping(address => uint256)) public allowance;
    function mint(address to, uint256 a) external { balanceOf[to] += a; }
    function approve(address s, uint256 a) external returns (bool) { allowance[msg.sender][s] = a; return true; }
    function transfer(address to, uint256 a) external returns (bool) {
        balanceOf[msg.sender] -= a; balanceOf[to] += a; return true;
    }
    function transferFrom(address f, address to, uint256 a) external returns (bool) {
        allowance[f][msg.sender] -= a; balanceOf[f] -= a; balanceOf[to] += a; return true;
    }
}

/// @dev A shared router shaped like the Universal Router's v4 path: the caller supplies no
/// hookData, the router settles its debt from the payer, takes the credit to the user, and has a
/// public sweep for tokens left on it (Universal Router's SWEEP command).
contract SharedRouter is IUnlockCallback {
    using TransientStateLibrary for IPoolManager;
    IPoolManager immutable manager;
    constructor(IPoolManager m) { manager = m; }

    error DeltaNotPositive();

    function swap(PoolKey memory key, SwapParams memory p, address recipient, bytes memory hookData) external {
        manager.unlock(abi.encode(key, p, msg.sender, recipient, hookData));
    }

    function unlockCallback(bytes calldata data) external returns (bytes memory) {
        require(msg.sender == address(manager));
        (PoolKey memory key, SwapParams memory p, address payer, address recipient, bytes memory hookData) =
            abi.decode(data, (PoolKey, SwapParams, address, address, bytes));
        manager.swap(key, p, hookData);
        Currency input = p.zeroForOne ? key.currency0 : key.currency1;
        Currency output = p.zeroForOne ? key.currency1 : key.currency0;
        int256 d = manager.currencyDelta(address(this), input);
        if (d < 0) {
            manager.sync(input);
            IERC20Like(Currency.unwrap(input)).transferFrom(payer, address(manager), uint256(-d));
            manager.settle();
        }
        int256 o = manager.currencyDelta(address(this), output);
        if (o < 0) revert DeltaNotPositive();
        if (o > 0) manager.take(output, recipient, uint256(o));
        return "";
    }

    function sweep(address token, address to) external {
        IERC20Like(token).transfer(to, IERC20Like(token).balanceOf(address(this)));
    }
}

contract RewardToRouterProof is Test, IUnlockCallback {
    address constant IMD = 0xD34a99Bc0f67aE1bbd63C660e6d0b0dd03E263B7;
    uint160 constant Q96 = 79228162514264337593543950336;
    PoolManager manager;
    EHS token;
    EHSHook hook;
    PoolKey key;
    SharedRouter router;
    bool pair0;

    function setUp() public {
        manager = new PoolManager(address(this));
        vm.etch(IMD, address(new PlainIMD()).code);
        token = new EHS();
        bytes memory code = abi.encodePacked(type(EHSHook).creationCode, abi.encode(manager, address(token)));
        bytes32 hash = keccak256(code);
        address deployed;
        for (uint256 i; i < 300_000; ++i) {
            bytes32 salt = bytes32(i);
            address predicted =
                address(uint160(uint256(keccak256(abi.encodePacked(bytes1(0xff), address(this), salt, hash)))));
            if (uint160(predicted) & 0x3fff != 0x20cc) continue;
            assembly ("memory-safe") { deployed := create2(0, add(code, 32), mload(code), salt) }
            break;
        }
        hook = EHSHook(deployed);
        pair0 = IMD < address(token);
        key = PoolKey(
            Currency.wrap(pair0 ? IMD : address(token)),
            Currency.wrap(pair0 ? address(token) : IMD),
            12500,
            60,
            IHooks(deployed)
        );
        manager.initialize(key, Q96);
        router = new SharedRouter(manager);
        // Token-only seed, as the launch factory does.
        int24 lower = pair0 ? int24(-600) : int24(60);
        int24 upper = pair0 ? int24(-60) : int24(600);
        manager.unlock(abi.encode(lower, upper, int256(1_000_000 ether)));
    }

    function unlockCallback(bytes calldata data) external returns (bytes memory) {
        (int24 lower, int24 upper, int256 liq) = abi.decode(data, (int24, int24, int256));
        (BalanceDelta delta,) = manager.modifyLiquidity(key, ModifyLiquidityParams(lower, upper, liq, bytes32(0)), "");
        settle(key.currency0, delta.amount0());
        settle(key.currency1, delta.amount1());
        return "";
    }

    function settle(Currency c, int128 a) internal {
        if (a < 0) {
            manager.sync(c);
            IERC20Like(Currency.unwrap(c)).transfer(address(manager), uint256(uint128(-a)));
            manager.settle();
        } else if (a > 0) {
            manager.take(c, address(this), uint256(uint128(a)));
        }
    }

    /// A buy routed through a shared router with no hookData must not credit the router: the
    /// router is not the buyer, and anyone can pay it with payRecorded(router) and sweep the IMD.
    function test_rewardWithoutHookDataMustNotGoToTheRouter() public {
        address alice = makeAddr("alice");
        PlainIMD(IMD).mint(alice, 100 ether);
        vm.prank(alice);
        IERC20Like(IMD).approve(address(router), type(uint256).max);
        uint160 limit = pair0 ? TickMath.MIN_SQRT_PRICE + 1 : TickMath.MAX_SQRT_PRICE - 1;
        vm.prank(alice);
        router.swap(key, SwapParams(pair0, -int256(100 ether), limit), alice, "");
        assertGt(token.balanceOf(alice), 0, "alice bought EHS");
        assertEq(hook.pending(alice), 0, "alice is unknown to the hook");

        // Current code: pending[router] == 0.980392156862745098e18 and a bot takes it.
        uint256 recorded = hook.pending(address(router));
        address bot = makeAddr("bot");
        vm.prank(bot);
        hook.payRecorded(address(router));
        vm.prank(bot);
        router.sweep(IMD, bot);
        assertEq(IERC20Like(IMD).balanceOf(bot), recorded, "payRecorded(router) pays the router; sweep takes it");

        assertEq(recorded, 0, "reward recorded to the shared router instead of the buyer");
        assertEq(hook.owedTotal(), 0, "an obligation nobody entitled can claim was recorded");
    }
}

// SPDX-License-Identifier: MIT
pragma solidity 0.8.26;

import {Test} from "forge-std/Test.sol";
import {PoolManager} from "v4-core/src/PoolManager.sol";
import {IPoolManager} from "v4-core/src/interfaces/IPoolManager.sol";
import {IHooks} from "v4-core/src/interfaces/IHooks.sol";
import {IUnlockCallback} from "v4-core/src/interfaces/callback/IUnlockCallback.sol";
import {Currency} from "v4-core/src/types/Currency.sol";
import {PoolKey} from "v4-core/src/types/PoolKey.sol";
import {SwapParams, ModifyLiquidityParams} from "v4-core/src/types/PoolOperation.sol";
import {BalanceDelta} from "v4-core/src/types/BalanceDelta.sol";
import {TickMath} from "v4-core/src/libraries/TickMath.sol";
import {TransientStateLibrary} from "v4-core/src/libraries/TransientStateLibrary.sol";
import {EHS} from "src/EHS.sol";
import {EHSHook} from "src/EHSHook.sol";

/// @dev Stand-in for the mainnet IMD ERC-20 (etched at its address).
contract PlainERC20 {
    mapping(address => uint256) public balanceOf;
    mapping(address => mapping(address => uint256)) public allowance;

    function mint(address to, uint256 amount) external {
        balanceOf[to] += amount;
    }

    function approve(address spender, uint256 amount) external returns (bool) {
        allowance[msg.sender][spender] = amount;
        return true;
    }

    function transfer(address to, uint256 amount) external returns (bool) {
        balanceOf[msg.sender] -= amount;
        balanceOf[to] += amount;
        return true;
    }

    function transferFrom(address from, address to, uint256 amount) external returns (bool) {
        allowance[from][msg.sender] -= amount;
        balanceOf[from] -= amount;
        balanceOf[to] += amount;
        return true;
    }
}

/// @dev A shared, permissionless router in the shape of the Universal Router: it swaps for
/// whoever calls it, sends the output to the caller's chosen recipient, passes the hookData the
/// caller gives it (the Uniswap interface gives none), and lets anyone sweep tokens it holds.
contract SharedRouter is IUnlockCallback {
    IPoolManager immutable manager;

    constructor(IPoolManager m) {
        manager = m;
    }

    function swap(PoolKey calldata key, SwapParams calldata params, address recipient, bytes calldata hookData)
        external
    {
        manager.unlock(abi.encode(msg.sender, key, params, recipient, hookData));
    }

    function unlockCallback(bytes calldata data) external returns (bytes memory) {
        require(msg.sender == address(manager));
        (address payer, PoolKey memory key, SwapParams memory params, address recipient, bytes memory hookData) =
            abi.decode(data, (address, PoolKey, SwapParams, address, bytes));
        BalanceDelta delta = manager.swap(key, params, hookData);
        _settle(key.currency0, delta.amount0(), payer, recipient);
        _settle(key.currency1, delta.amount1(), payer, recipient);
        return "";
    }

    function _settle(Currency c, int128 amount, address payer, address recipient) internal {
        if (amount < 0) {
            manager.sync(c);
            PlainERC20(Currency.unwrap(c)).transferFrom(payer, address(manager), uint256(uint128(-amount)));
            manager.settle();
        } else if (amount > 0) {
            manager.take(c, recipient, uint256(uint128(amount)));
        }
    }

    /// @dev Universal Router's SWEEP command: anyone may take whatever the router holds.
    function sweep(address token, address to) external {
        PlainERC20(token).transfer(to, PlainERC20(token).balanceOf(address(this)));
    }
}

contract RouterRewardLeakTest is Test, IUnlockCallback {
    address constant IMD = 0xD34a99Bc0f67aE1bbd63C660e6d0b0dd03E263B7;
    uint160 constant Q96 = 79228162514264337593543950336;

    PoolManager manager;
    EHS token;
    EHSHook hook;
    SharedRouter router;
    PoolKey key;
    bool pair0;

    function setUp() public {
        manager = new PoolManager(address(this));
        token = new EHS();
        vm.etch(IMD, address(new PlainERC20()).code);

        bytes memory code = abi.encodePacked(type(EHSHook).creationCode, abi.encode(manager, address(token)));
        bytes32 hash = keccak256(code);
        address deployed;
        for (uint256 i; i < 300_000; ++i) {
            bytes32 salt = bytes32(i);
            address predicted =
                address(uint160(uint256(keccak256(abi.encodePacked(bytes1(0xff), address(this), salt, hash)))));
            if (uint160(predicted) & 0x3fff != 0x20cc) continue;
            assembly ("memory-safe") {
                deployed := create2(0, add(code, 32), mload(code), salt)
            }
            break;
        }
        require(deployed != address(0), "mining failed");
        hook = EHSHook(deployed);
        pair0 = IMD < address(token);
        key = PoolKey(
            Currency.wrap(pair0 ? IMD : address(token)),
            Currency.wrap(pair0 ? address(token) : IMD),
            12500,
            60,
            IHooks(deployed)
        );
        manager.initialize(key, Q96);
        PlainERC20(IMD).mint(address(this), 10_000_000 ether);
        manager.unlock("");
        router = new SharedRouter(manager);
    }

    // Liquidity provision for the test's own pool.
    function unlockCallback(bytes calldata) external returns (bytes memory) {
        manager.modifyLiquidity(key, ModifyLiquidityParams(-60000, 60000, 1_000_000 ether, bytes32(0)), "");
        _pay(key.currency0);
        _pay(key.currency1);
        return "";
    }

    function _pay(Currency c) internal {
        int256 d = TransientStateLibrary.currencyDelta(manager, address(this), c);
        if (d < 0) {
            manager.sync(c);
            PlainERC20(Currency.unwrap(c)).transfer(address(manager), uint256(-d));
            manager.settle();
        }
    }

    /// A buyer buys through a shared router with the empty hookData the Uniswap interface sends.
    /// The reward for that buy must not be credited to the router, where anyone can take it.
    function test_rewardIsNotCreditedToASharedRouterOnEmptyHookData() public {
        address buyer = makeAddr("buyer");
        address attacker = makeAddr("attacker");
        PlainERC20(IMD).mint(buyer, 1_000 ether);
        vm.prank(buyer);
        PlainERC20(IMD).approve(address(router), type(uint256).max);

        // Exact-output buy of EHS: IMD is the unspecified side, so the fee and the reward
        // follow afterSwap's real fill. The buy fills well above the 10 IMD threshold.
        bool zeroForOne = pair0;
        uint160 limit = zeroForOne ? TickMath.MIN_SQRT_PRICE + 1 : TickMath.MAX_SQRT_PRICE - 1;
        vm.prank(buyer);
        router.swap(key, SwapParams(zeroForOne, int256(500 ether), limit), buyer, "");

        assertGt(token.balanceOf(buyer), 0, "buyer received EHS");
        assertGt(hook.accruedBalance(), 10 ether, "hook took its 2% fee");

        // The buyer is the swap recipient; the router is a shared contract. Nothing should be owed
        // to the router.
        assertEq(hook.pending(address(router)), 0, "reward recorded to the shared router");

        // Anyone can pay the router's record and sweep it away from the buyer.
        vm.startPrank(attacker);
        hook.payRecorded(address(router));
        router.sweep(IMD, attacker);
        vm.stopPrank();
        assertEq(PlainERC20(IMD).balanceOf(attacker), 0, "attacker swept the buyer's reward");
    }
}

// SPDX-License-Identifier: MIT
pragma solidity 0.8.26;

import {Test} from "forge-std/Test.sol";
import {PoolManager} from "v4-core/src/PoolManager.sol";
import {IPoolManager} from "v4-core/src/interfaces/IPoolManager.sol";
import {IHooks} from "v4-core/src/interfaces/IHooks.sol";
import {IUnlockCallback} from "v4-core/src/interfaces/callback/IUnlockCallback.sol";
import {Currency} from "v4-core/src/types/Currency.sol";
import {PoolKey} from "v4-core/src/types/PoolKey.sol";
import {SwapParams, ModifyLiquidityParams} from "v4-core/src/types/PoolOperation.sol";
import {TransientStateLibrary} from "v4-core/src/libraries/TransientStateLibrary.sol";
import {TickMath} from "v4-core/src/libraries/TickMath.sol";
import {EHS} from "src/EHS.sol";
import {EHSHook} from "src/EHSHook.sol";

/// @dev Minimal ERC-20 standing in for IMD at its mainnet address.
contract StubIMD {
    string public constant symbol = "IMD";
    uint8 public constant decimals = 18;
    mapping(address => uint256) public balanceOf;

    function mint(address to, uint256 amount) external {
        balanceOf[to] += amount;
    }

    function transfer(address to, uint256 amount) external returns (bool) {
        balanceOf[msg.sender] -= amount;
        balanceOf[to] += amount;
        return true;
    }
}

/// @dev A shared, stateless router in the shape of the Universal Router: anyone may route a swap
/// through it with any hookData (the Uniswap app sends empty hookData), and anyone may SWEEP whatever
/// token balance sits in it (Universal Router command 0x04, no access control).
contract SharedRouter is IUnlockCallback {
    using TransientStateLibrary for IPoolManager;

    IPoolManager internal immutable manager;
    address internal immutable imd;

    constructor(IPoolManager manager_, address imd_) {
        manager = manager_;
        imd = imd_;
    }

    function swap(PoolKey memory key, SwapParams memory params, bytes memory hookData, address recipient) external {
        manager.unlock(abi.encode(msg.sender, recipient, key, params, hookData));
    }

    function sweep(address token, address to) external {
        StubIMD(token).transfer(to, StubIMD(token).balanceOf(address(this)));
    }

    function unlockCallback(bytes calldata data) external returns (bytes memory) {
        require(msg.sender == address(manager), "manager only");
        (address payer, address recipient, PoolKey memory key, SwapParams memory params, bytes memory hookData) =
            abi.decode(data, (address, address, PoolKey, SwapParams, bytes));
        manager.swap(key, params, hookData);
        settle(key.currency0, manager.currencyDelta(address(this), key.currency0), payer, recipient);
        settle(key.currency1, manager.currencyDelta(address(this), key.currency1), payer, recipient);
        return "";
    }

    function settle(Currency currency, int256 amount, address payer, address recipient) internal {
        if (amount < 0) {
            manager.sync(currency);
            // The payer pre-funded the router (Permit2 pull in the real router).
            (bool ok,) = Currency.unwrap(currency).call(
                abi.encodeWithSignature("transfer(address,uint256)", address(manager), uint256(-amount))
            );
            require(ok, "transfer");
            manager.settle();
            payer;
        } else if (amount > 0) {
            manager.take(currency, recipient, uint256(amount));
        }
    }
}

contract RouterRewardSweepTest is Test, IUnlockCallback {
    address internal constant IMD = 0xD34a99Bc0f67aE1bbd63C660e6d0b0dd03E263B7;
    uint160 internal constant Q96 = 79228162514264337593543950336;
    IPoolManager internal manager;
    EHS internal token;
    EHSHook internal hook;
    PoolKey internal key;
    SharedRouter internal router;
    bool internal pair0;

    function setUp() public {
        manager = IPoolManager(address(new PoolManager(address(this))));
        vm.etch(IMD, address(new StubIMD()).code);
        StubIMD(IMD).mint(address(this), 1e30);
        token = new EHS();
        bytes memory code = abi.encodePacked(type(EHSHook).creationCode, abi.encode(manager, address(token)));
        bytes32 hash = keccak256(code);
        address deployed;
        for (uint256 i; i < 300_000; ++i) {
            address predicted =
                address(uint160(uint256(keccak256(abi.encodePacked(bytes1(0xff), address(this), bytes32(i), hash)))));
            if (uint160(predicted) & 0x3fff != 0x20cc) continue;
            bytes32 salt = bytes32(i);
            assembly ("memory-safe") {
                deployed := create2(0, add(code, 32), mload(code), salt)
            }
            break;
        }
        require(deployed != address(0), "mining failed");
        hook = EHSHook(deployed);
        pair0 = IMD < address(token);
        key = PoolKey(
            Currency.wrap(pair0 ? IMD : address(token)),
            Currency.wrap(pair0 ? address(token) : IMD),
            12500,
            60,
            IHooks(deployed)
        );
        manager.initialize(key, Q96);
        manager.unlock(abi.encode(int24(-60000), int24(60000), int256(1_000_000 ether)));
        router = new SharedRouter(manager, IMD);
    }

    function unlockCallback(bytes calldata data) external returns (bytes memory) {
        (int24 lower, int24 upper, int256 liquidity) = abi.decode(data, (int24, int24, int256));
        manager.modifyLiquidity(key, ModifyLiquidityParams(lower, upper, liquidity, bytes32(0)), "");
        settle(key.currency0);
        settle(key.currency1);
        return "";
    }

    function settle(Currency currency) internal {
        int256 amount = TransientStateLibrary.currencyDelta(manager, address(this), currency);
        if (amount < 0) {
            manager.sync(currency);
            currency.transfer(address(manager), uint256(-amount));
            manager.settle();
        } else if (amount > 0) {
            manager.take(currency, address(this), uint256(amount));
        }
    }

    /// A buyer trades through a shared router with empty hookData (the Uniswap app's default). The
    /// hook records the buyer's 1% reward to the router. A stranger then pays that record to the
    /// router and sweeps it out. Expected: the buyer's reward is not payable to a stranger.
    function test_rewardRecordedToSharedRouterIsSweptByAnyone() public {
        address buyer = makeAddr("buyer");
        address stranger = makeAddr("stranger");
        StubIMD(IMD).mint(address(router), 1_000 ether); // the buyer's funds, pulled by the router
        vm.prank(buyer);
        router.swap(key, SwapParams(pair0, -int256(1_000 ether), pair0 ? TickMath.MIN_SQRT_PRICE + 1 : TickMath.MAX_SQRT_PRICE - 1), "", buyer);
        assertGt(token.balanceOf(buyer), 0, "buyer received EHS");
        assertEq(hook.pending(buyer), 0, "buyer has no record");
        emit log_named_uint("reward recorded to the shared router", hook.pending(address(router)));

        vm.startPrank(stranger);
        hook.payRecorded(address(router));
        router.sweep(IMD, stranger);
        vm.stopPrank();

        assertEq(StubIMD(IMD).balanceOf(stranger), 0, "a stranger must not receive the buyer's reward");
    }
}

// SPDX-License-Identifier: MIT
pragma solidity 0.8.26;

import {Test} from "forge-std/Test.sol";
import {PoolManager} from "v4-core/src/PoolManager.sol";
import {IPoolManager} from "v4-core/src/interfaces/IPoolManager.sol";
import {IHooks} from "v4-core/src/interfaces/IHooks.sol";
import {IUnlockCallback} from "v4-core/src/interfaces/callback/IUnlockCallback.sol";
import {Currency} from "v4-core/src/types/Currency.sol";
import {PoolKey} from "v4-core/src/types/PoolKey.sol";
import {SwapParams, ModifyLiquidityParams} from "v4-core/src/types/PoolOperation.sol";
import {BalanceDelta} from "v4-core/src/types/BalanceDelta.sol";
import {TickMath} from "v4-core/src/libraries/TickMath.sol";
import {TransientStateLibrary} from "v4-core/src/libraries/TransientStateLibrary.sol";
import {EHS} from "src/EHS.sol";
import {EHSHook} from "src/EHSHook.sol";

/// @dev Stand-in for the mainnet IMD ERC-20 (etched at its address).
contract PlainERC20 {
    mapping(address => uint256) public balanceOf;
    mapping(address => mapping(address => uint256)) public allowance;

    function mint(address to, uint256 amount) external {
        balanceOf[to] += amount;
    }

    function approve(address spender, uint256 amount) external returns (bool) {
        allowance[msg.sender][spender] = amount;
        return true;
    }

    function transfer(address to, uint256 amount) external returns (bool) {
        balanceOf[msg.sender] -= amount;
        balanceOf[to] += amount;
        return true;
    }

    function transferFrom(address from, address to, uint256 amount) external returns (bool) {
        allowance[from][msg.sender] -= amount;
        balanceOf[from] -= amount;
        balanceOf[to] += amount;
        return true;
    }
}

/// @dev A shared, permissionless router in the shape of the Universal Router: it swaps for
/// whoever calls it, sends the output to the caller's chosen recipient, passes the hookData the
/// caller gives it (the Uniswap interface gives none), and lets anyone sweep tokens it holds.
contract SharedRouter is IUnlockCallback {
    IPoolManager immutable manager;

    constructor(IPoolManager m) {
        manager = m;
    }

    function swap(PoolKey calldata key, SwapParams calldata params, address recipient, bytes calldata hookData)
        external
    {
        manager.unlock(abi.encode(msg.sender, key, params, recipient, hookData));
    }

    function unlockCallback(bytes calldata data) external returns (bytes memory) {
        require(msg.sender == address(manager));
        (address payer, PoolKey memory key, SwapParams memory params, address recipient, bytes memory hookData) =
            abi.decode(data, (address, PoolKey, SwapParams, address, bytes));
        BalanceDelta delta = manager.swap(key, params, hookData);
        _settle(key.currency0, delta.amount0(), payer, recipient);
        _settle(key.currency1, delta.amount1(), payer, recipient);
        return "";
    }

    function _settle(Currency c, int128 amount, address payer, address recipient) internal {
        if (amount < 0) {
            manager.sync(c);
            PlainERC20(Currency.unwrap(c)).transferFrom(payer, address(manager), uint256(uint128(-amount)));
            manager.settle();
        } else if (amount > 0) {
            manager.take(c, recipient, uint256(uint128(amount)));
        }
    }

    /// @dev Universal Router's SWEEP command: anyone may take whatever the router holds.
    function sweep(address token, address to) external {
        PlainERC20(token).transfer(to, PlainERC20(token).balanceOf(address(this)));
    }
}

contract RouterRewardLeakTest is Test, IUnlockCallback {
    address constant IMD = 0xD34a99Bc0f67aE1bbd63C660e6d0b0dd03E263B7;
    uint160 constant Q96 = 79228162514264337593543950336;

    PoolManager manager;
    EHS token;
    EHSHook hook;
    SharedRouter router;
    PoolKey key;
    bool pair0;

    function setUp() public {
        manager = new PoolManager(address(this));
        token = new EHS();
        vm.etch(IMD, address(new PlainERC20()).code);

        bytes memory code = abi.encodePacked(type(EHSHook).creationCode, abi.encode(manager, address(token)));
        bytes32 hash = keccak256(code);
        address deployed;
        for (uint256 i; i < 300_000; ++i) {
            bytes32 salt = bytes32(i);
            address predicted =
                address(uint160(uint256(keccak256(abi.encodePacked(bytes1(0xff), address(this), salt, hash)))));
            if (uint160(predicted) & 0x3fff != 0x20cc) continue;
            assembly ("memory-safe") {
                deployed := create2(0, add(code, 32), mload(code), salt)
            }
            break;
        }
        require(deployed != address(0), "mining failed");
        hook = EHSHook(deployed);
        pair0 = IMD < address(token);
        key = PoolKey(
            Currency.wrap(pair0 ? IMD : address(token)),
            Currency.wrap(pair0 ? address(token) : IMD),
            12500,
            60,
            IHooks(deployed)
        );
        manager.initialize(key, Q96);
        PlainERC20(IMD).mint(address(this), 10_000_000 ether);
        manager.unlock("");
        router = new SharedRouter(manager);
    }

    // Liquidity provision for the test's own pool.
    function unlockCallback(bytes calldata) external returns (bytes memory) {
        manager.modifyLiquidity(key, ModifyLiquidityParams(-60000, 60000, 1_000_000 ether, bytes32(0)), "");
        _pay(key.currency0);
        _pay(key.currency1);
        return "";
    }

    function _pay(Currency c) internal {
        int256 d = TransientStateLibrary.currencyDelta(manager, address(this), c);
        if (d < 0) {
            manager.sync(c);
            PlainERC20(Currency.unwrap(c)).transfer(address(manager), uint256(-d));
            manager.settle();
        }
    }

    /// A buyer buys through a shared router with the empty hookData the Uniswap interface sends.
    /// The reward for that buy must not be credited to the router, where anyone can take it.
    function test_rewardIsNotCreditedToASharedRouterOnEmptyHookData() public {
        address buyer = makeAddr("buyer");
        address attacker = makeAddr("attacker");
        PlainERC20(IMD).mint(buyer, 1_000 ether);
        vm.prank(buyer);
        PlainERC20(IMD).approve(address(router), type(uint256).max);

        // Exact-output buy of EHS: IMD is the unspecified side, so the fee and the reward
        // follow afterSwap's real fill. The buy fills well above the 10 IMD threshold.
        bool zeroForOne = pair0;
        uint160 limit = zeroForOne ? TickMath.MIN_SQRT_PRICE + 1 : TickMath.MAX_SQRT_PRICE - 1;
        vm.prank(buyer);
        router.swap(key, SwapParams(zeroForOne, int256(500 ether), limit), buyer, "");

        assertGt(token.balanceOf(buyer), 0, "buyer received EHS");
        assertGt(hook.accruedBalance(), 10 ether, "hook took its 2% fee");

        // The buyer is the swap recipient; the router is a shared contract. Nothing should be owed
        // to the router.
        assertEq(hook.pending(address(router)), 0, "reward recorded to the shared router");

        // Anyone can pay the router's record and sweep it away from the buyer.
        vm.startPrank(attacker);
        hook.payRecorded(address(router));
        router.sweep(IMD, attacker);
        vm.stopPrank();
        assertEq(PlainERC20(IMD).balanceOf(attacker), 0, "attacker swept the buyer's reward");
    }
}