Maia DAO - restakeToken is not permissionless PoC (instructions in the comments)

Boilerplate.sol

// SPDX-License-Identifier: MIT

pragma solidity ^0.8.15;

import {Test, console2 as console} from "forge-std/Test.sol";

import {FixedPointMathLib} from "solmate/utils/FixedPointMathLib.sol";
import {SafeCastLib} from "solmate/utils/SafeCastLib.sol";
import {SafeTransferLib} from "solmate/utils/SafeTransferLib.sol";
import {ERC20} from "solmate/tokens/ERC20.sol";
import {MockERC20} from "solmate/test/utils/mocks/MockERC20.sol";
import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";

// Imports for minting UNI-V3 positions
import "@uniswap/v3-core/contracts/interfaces/IUniswapV3Pool.sol";
import "@uniswap/v3-core/contracts/libraries/TickMath.sol";
import "@openzeppelin/contracts/token/ERC721/IERC721Receiver.sol";
import "@uniswap/v3-periphery/contracts/interfaces/ISwapRouter.sol";
import "@uniswap/v3-periphery/contracts/interfaces/INonfungiblePositionManager.sol";
import "@uniswap/v3-periphery/contracts/libraries/TransferHelper.sol";
import "@uniswap/v3-periphery/contracts/base/LiquidityManagement.sol";

import {UniswapV3Factory, UniswapV3Pool} from "@uniswap/v3-core/contracts/UniswapV3Factory.sol";

import {IWETH9} from "@uniswap/v3-periphery/contracts/interfaces/external/IWETH9.sol";
import {NonfungiblePositionManager} from "@uniswap/v3-periphery/contracts/NonfungiblePositionManager.sol";

import {
    UniswapV3GaugeFactory,
    FlywheelGaugeRewards,
    BaseV2GaugeManager
} from "@gauges/factories/UniswapV3GaugeFactory.sol";
import {BribesFactory, FlywheelBoosterGaugeWeight} from "@gauges/factories/BribesFactory.sol";
import {UniswapV3Gauge, BaseV2Gauge} from "@gauges/UniswapV3Gauge.sol";

import {BaseV2Minter} from "@hermes/minters/BaseV2Minter.sol";
import {bHermes} from "@hermes/bHermes.sol";
import {HERMES} from "@hermes/tokens/HERMES.sol";

import {PoolVariables} from "@talos/libraries/PoolVariables.sol";

import {IUniswapV3Pool, UniswapV3Staker, IUniswapV3Staker, IncentiveTime} from "@v3-staker/UniswapV3Staker.sol";

contract Boilerplate is Test {
    ////////////////////////////////////////////////////////////////////
    //                   Original UniV3StakerTest setup               //
    ////////////////////////////////////////////////////////////////////

    using FixedPointMathLib for uint256;
    using FixedPointMathLib for uint160;
    using FixedPointMathLib for uint128;
    using SafeCastLib for uint256;
    using SafeCastLib for int256;
    using SafeTransferLib for ERC20;

    struct Deposit {
        address owner;
        uint128 liquidity;
        address token0;
        address token1;
    }

    mapping(uint256 => Deposit) public deposits;

    bHermes bHermesToken;

    BaseV2Minter baseV2Minter;
    BaseV2GaugeManager baseV2GaugeManager;

    FlywheelGaugeRewards flywheelGaugeRewards;
    BribesFactory bribesFactory;

    FlywheelBoosterGaugeWeight flywheelGaugeWeightBooster;

    UniswapV3GaugeFactory uniswapV3GaugeFactory;
    UniswapV3Gauge gauge;

    HERMES rewardToken;

    IUniswapV3Staker uniswapV3Staker;
    UniswapV3Staker uniswapV3StakerContract;

    IUniswapV3Staker.IncentiveKey key;
    bytes32 incentiveId;

    // Pool fee on arbitrum DAI/USDC pool is 0.01%
    uint24 constant poolFee = 100;

    ////////////////////////////////////////////////////////////////////
    //                      Testing Boilerplate                       //
    ////////////////////////////////////////////////////////////////////

    address public ATTACKER;
    address public DEPLOYER;
    address public ADMIN;
    address public ALICE;
    address public BOB;
    address public CHARLIE;
    address public EVE;

    bool activePrank;

    // Arbitrum forked by anvil, this way we can reduce calls to Alchemy
    // Make sure to first spin up anvil with --fork-url of arbitrum mainnet
    string internal localhost = vm.envString("LOCALHOST_RPC_URL");

    // Initialize is used instead of setUp() to have additional setUp() available
    // in each test contract that inherits from this one
    function initializeBoilerplate() public {
        makeAddr();
        vm.label(address(this), "Test contract");
        vm.createSelectFork(localhost);
        console.log("[BOILERPLATE] Current fork: %s", localhost);
    }

    ////////////////////////////////////////////////////////////////////
    //                      Addresses Arbitrum                        //
    ////////////////////////////////////////////////////////////////////

    // Old Arbitrum USDC
    address constant USDC = 0xFF970A61A04b1cA14834A43f5dE4533eBDDB5CC8;
    address constant DAI = 0xDA10009cBd5D07dd0CeCc66161FC93D7c9000da1;

    UniswapV3Factory uniswapV3Factory = UniswapV3Factory(0x1F98431c8aD98523631AE4a59f267346ea31F984);

    INonfungiblePositionManager nonfungiblePositionManager =
        INonfungiblePositionManager(payable(0xC36442b4a4522E871399CD717aBDD847Ab11FE88));

    // Token0: DAI, Token1: USDC
    UniswapV3Pool DAI_USDC_pool = UniswapV3Pool(0xF0428617433652c9dc6D1093A42AdFbF30D29f74);

    ////////////////////////////////////////////////////////////////////
    //                          Utilities                             //
    ////////////////////////////////////////////////////////////////////

    function makeAddr() public {
        ATTACKER = address(0x1337);
        vm.label(ATTACKER, "ATTACKER");
        DEPLOYER = address(0xDEADBEEF);
        vm.label(DEPLOYER, "DEPLOYER");
        ADMIN = address(0xCAFE);
        vm.label(ADMIN, "ADMIN");
        ALICE = address(0x1111);
        vm.label(ALICE, "ALICE");
        BOB = address(0x2222);
        vm.label(BOB, "BOB");
        CHARLIE = address(0x3333);
        vm.label(CHARLIE, "CHARLIE");
        EVE = address(0x4444);
        vm.label(EVE, "EVE");

        vm.label(address(nonfungiblePositionManager), "nonfungiblePositionManager");
        vm.label(address(DAI_USDC_pool), "DAI_USDC_pool");
        vm.label(address(uniswapV3Factory), "uniswapV3Factory");
    }
}

BrokenRestake.t.sol

// SPDX-License-Identifier: AGPL-3.0-only
pragma solidity ^0.8.0;
pragma abicoder v2;

import "./Boilerplate.sol";

contract BrokenRestake is Boilerplate, IERC721Receiver {
    uint256 tokenIdAlice;
    uint256 tokenIdBob;

    function setUp() public {
        initializeBoilerplate();

        // In HERMES deployment the owner will be set to BaseV2Minter
        // In this test ADMIN is set as a temporary owner just to mint
        // initial tokens to get the reward system rolling.
        rewardToken = new HERMES({_owner: ADMIN});

        vm.startPrank(ADMIN);
        rewardToken.mint(ALICE, 1_000e18);
        rewardToken.mint(BOB, 1_000e18);
        vm.stopPrank();

        bHermesToken =
        new bHermes({ _hermes: rewardToken, _owner: ADMIN, _gaugeCycleLength: 1 weeks, _incrementFreezeWindow: 12 hours });

        flywheelGaugeWeightBooster = new FlywheelBoosterGaugeWeight({ _bHermesGauges: bHermesToken.gaugeWeight() });

        bribesFactory = new BribesFactory(
            BaseV2GaugeManager(address(this)),
            flywheelGaugeWeightBooster,
            1 weeks,
            ADMIN
        );

        baseV2Minter = new BaseV2Minter({ _vault:address(bHermesToken), _dao: address(0), _owner: ADMIN });

        // Transfer ownership to BaseV2Minter to mimic real deployment
        vm.prank(ADMIN);
        rewardToken.transferOwnership(address(baseV2Minter));

        flywheelGaugeRewards =
        new FlywheelGaugeRewards({_rewardToken: address(rewardToken), _owner:ADMIN, _gaugeToken:bHermesToken.gaugeWeight(), _minter:baseV2Minter});

        baseV2Minter.initialize(flywheelGaugeRewards);

        uniswapV3GaugeFactory = new UniswapV3GaugeFactory(
            BaseV2GaugeManager(address(0)),
            bHermesToken.gaugeBoost(),
            uniswapV3Factory,
            nonfungiblePositionManager,
            flywheelGaugeRewards,
            bribesFactory,
            // Owner
            ADMIN
        );

        vm.mockCall(address(0), abi.encodeWithSignature("addGauge(address)"), abi.encode(""));

        uniswapV3StakerContract = uniswapV3GaugeFactory.uniswapV3Staker();

        uniswapV3Staker = IUniswapV3Staker(address(uniswapV3StakerContract));

        deal(address(DAI), ALICE, 1_000e18);
        deal(address(USDC), ALICE, 1_000e6);
        deal(address(DAI), BOB, 1_000e18);
        deal(address(USDC), BOB, 1_000e6);

        vm.startPrank(ADMIN);
        gauge = createGaugeAndAddToGaugeBoost({_pool: DAI_USDC_pool, minWidth: 1});
        vm.stopPrank();

        vm.startPrank(ALICE);
        rewardToken.approve(address(bHermesToken), 500e18);
        bHermesToken.deposit(500e18, ALICE);
        bHermesToken.claimMultiple(500e18);
        bHermesToken.gaugeWeight().delegate(ALICE);
        bHermesToken.gaugeWeight().incrementGauge(address(gauge), uint112(500e18));
        vm.stopPrank();

        vm.startPrank(BOB);
        rewardToken.approve(address(bHermesToken), 500e18);
        bHermesToken.deposit(500e18, BOB);
        bHermesToken.claimMultiple(500e18);
        bHermesToken.gaugeWeight().delegate(BOB);
        bHermesToken.gaugeWeight().incrementGauge(address(gauge), uint112(500e18));
        vm.stopPrank();

        // newEpoch() only triggers reward accrual if previous epoch < new epoch
        // Since the whole setup is happening in the previous epoch, we need to
        // warp to the next epoch
        vm.warp(block.timestamp + 1 weeks);
        // This call will pull HERMES token rewards from the minter
        // and create the initial incentive
        gauge.newEpoch();

        // This is how incentive start time is calculated in createIncentiveFromGauge
        // This is also the start of the new week
        uint256 incentiveStartTime = IncentiveTime.computeEnd(block.timestamp);
        vm.warp(incentiveStartTime);
        // That's why newEpoch() can be called
        gauge.newEpoch();

        vm.startPrank(ALICE);
        (tokenIdAlice,,,) = mintNewPosition({amount0ToMint: 1_000e18, amount1ToMint: 1_000e6, _recipient: ALICE});
        vm.stopPrank();

        vm.startPrank(BOB);
        (tokenIdBob,,,) = mintNewPosition({amount0ToMint: 1_000e18, amount1ToMint: 1_000e6, _recipient: BOB});
        vm.stopPrank();
    }

    function testRestake_RestakeIsNotPermissionless() public {
        vm.startPrank(ALICE);
        // 1.a Alice stakes her NFT (at incentive StartTime)
        nonfungiblePositionManager.safeTransferFrom(ALICE, address(uniswapV3Staker), tokenIdAlice);
        vm.stopPrank();

        vm.startPrank(BOB);
        // 1.b Bob stakes his NFT (at incentive StartTime)
        nonfungiblePositionManager.safeTransferFrom(BOB, address(uniswapV3Staker), tokenIdBob);
        vm.stopPrank();

        vm.warp(block.timestamp + 1 weeks); // 2.a Warp to incentive end time
        gauge.newEpoch();                   // 2.b Queue minter rewards for the next cycle

        vm.startPrank(BOB);
        uniswapV3Staker.restakeToken(tokenIdBob); // 3.a Bob can restake his own token
        vm.stopPrank();

        vm.startPrank(CHARLIE);
        vm.expectRevert(bytes4(keccak256("NotCalledByOwner()")));
        uniswapV3Staker.restakeToken(tokenIdAlice); // 3.b Charlie cannot restake Alice's token
        vm.stopPrank();

        uint256 rewardsBob = uniswapV3Staker.rewards(BOB);
        uint256 rewardsAlice = uniswapV3Staker.rewards(ALICE);

        assertNotEq(rewardsBob, 0, "Bob should have rewards");
        assertEq(rewardsAlice, 0, "Alice should not have rewards");

        console.log("");
        console.log("=================");
        console.log("Bob's rewards   : %s", rewardsBob);
        console.log("Alice's rewards : %s", rewardsAlice);
        console.log("=================");
    }

    ////////////////////////////////////////////////////////////////////
    //                            Utilities                           //
    ////////////////////////////////////////////////////////////////////

    /// @notice Calls the mint function defined in periphery, mints the same amount of each token. For this example we are providing 1000 WETH and 1000 USDC in liquidity
    /// @return tokenId The id of the newly minted ERC721
    /// @return liquidity The amount of liquidity for the position
    /// @return amount0 The amount of token0
    /// @return amount1 The amount of token1
    function mintNewPosition(uint256 amount0ToMint, uint256 amount1ToMint, address _recipient)
        public
        returns (uint256 tokenId, uint128 liquidity, uint256 amount0, uint256 amount1)
    {
        TransferHelper.safeApprove(DAI, address(nonfungiblePositionManager), amount0ToMint);
        TransferHelper.safeApprove(USDC, address(nonfungiblePositionManager), amount1ToMint);

        // Current tick -276325
        INonfungiblePositionManager.MintParams memory params = INonfungiblePositionManager.MintParams({
            token0: DAI,
            token1: USDC,
            fee: poolFee,
            tickLower: -276326,
            tickUpper: -276324,
            amount0Desired: amount0ToMint,
            amount1Desired: amount1ToMint,
            amount0Min: 0,
            amount1Min: 0,
            recipient: _recipient,
            deadline: block.timestamp
        });

        (tokenId, liquidity, amount0, amount1) = nonfungiblePositionManager.mint(params);

        // Create a deposit
        _createDeposit(msg.sender, tokenId);

        // Remove allowance and refund in both assets.

        if (amount0 < amount0ToMint) {
            TransferHelper.safeApprove(DAI, address(nonfungiblePositionManager), 0);
            uint256 refund0 = amount0ToMint - amount0;
            TransferHelper.safeTransfer(DAI, msg.sender, refund0);
        }

        if (amount1 < amount1ToMint) {
            TransferHelper.safeApprove(USDC, address(nonfungiblePositionManager), 0);
            uint256 refund1 = amount1ToMint - amount1;
            TransferHelper.safeTransfer(USDC, msg.sender, refund1);
        }
    }

    // Implementing `onERC721Received` so this contract can receive custody of erc721 tokens
    function onERC721Received(address operator, address, uint256 tokenId, bytes calldata)
        external
        override
        returns (bytes4)
    {
        _createDeposit(operator, tokenId);
        return this.onERC721Received.selector;
    }

    function _createDeposit(address owner, uint256 tokenId) internal {
        (,, address token0, address token1,,,, uint128 liquidity,,,,) = nonfungiblePositionManager.positions(tokenId);

        assertEq(token0, DAI, "Token0 should be DAI");
        assertEq(token1, USDC, "Token1 should be USDC");
        // set the owner and data for position
        // operator is msg.sender
        deposits[tokenId] = Deposit({owner: owner, liquidity: liquidity, token0: token0, token1: token1});
    }

    // Create a new Uniswap V3 Gauge from a Uniswap V3 pool
    function createGaugeAndAddToGaugeBoost(IUniswapV3Pool _pool, uint256 minWidth)
        internal
        returns (UniswapV3Gauge _gauge)
    {
        uniswapV3GaugeFactory.createGauge(address(_pool), abi.encode(uint24(minWidth)));
        _gauge = UniswapV3Gauge(address(uniswapV3GaugeFactory.strategyGauges(address(_pool))));
        bHermesToken.gaugeBoost().addGauge(address(_gauge));
        // The following 3 lines has been added on top of the original code
        bHermesToken.gaugeWeight().addGauge(address(_gauge));
        bHermesToken.gaugeWeight().setMaxGauges(4);
        bHermesToken.gaugeWeight().setMaxDelegates(1);
    }

    // Create a Uniswap V3 Staker incentive
    function createIncentive(IUniswapV3Staker.IncentiveKey memory _key, uint256 amount) internal {
        uniswapV3Staker.createIncentive(_key, amount);
    }
}

1. Create BrokenRestake.sol in the test/ folder and copy the content from the Gist above
2. Create Boilerplate.sol in the test/ folder and copy the content from the Gist above
3. Since the restakeToken function is missing from the IUniswapV3Staker.sol, add it:

diff --git a/src/uni-v3-staker/interfaces/IUniswapV3Staker.sol b/src/uni-v3-staker/interfaces/IUniswapV3Staker.sol
index 895c505..31c4153 100644
--- a/src/uni-v3-staker/interfaces/IUniswapV3Staker.sol
+++ b/src/uni-v3-staker/interfaces/IUniswapV3Staker.sol
@@ -239,6 +239,8 @@ interface IUniswapV3Staker is IERC721Receiver {
/// @param tokenId The ID of the token to stake
function stakeToken(uint256 tokenId) external;

+    function restakeToken(uint256 tokenId) external;
+
/*//////////////////////////////////////////////////////////////
                   GAUGE UPDATE LOGIC
//////////////////////////////////////////////////////////////*/

4. To remove the INIT_CODE_HASH issue (solved by the python script in the original repo), hardcode the pool address with the address of the DAI/USDC pool on Arbitrum Mainnet.
   Add the following line to the /src/uni-v3-staker/libraries/NFTPositionInfo.sol

diff --git a/src/uni-v3-staker/libraries/NFTPositionInfo.sol b/src/uni-v3-staker/libraries/NFTPositionInfo.sol
index 7d82ece..9543ebe 100644
--- a/src/uni-v3-staker/libraries/NFTPositionInfo.sol
+++ b/src/uni-v3-staker/libraries/NFTPositionInfo.sol
@@ -34,5 +34,6 @@ library NFTPositionInfo {
                 address(factory), PoolAddress.PoolKey({token0: token0, token1: token1, fee: fee})
             )
         );
+        pool = IUniswapV3Pool(0xF0428617433652c9dc6D1093A42AdFbF30D29f74);
     }
 }

5. Spin up a local anvil chain with anvil --fork-url $ARBITRUM_MAINNET_RPC_URL, where $ARBITRUM_MAINNET_RPC_URL is the URL provided by a service like Alchemy or Infura.
6. Run the test with forge test --mp test/BrokenRestake.sol --fork-url http://127.0.0.1:8545 -vvv