refactor: update staking contract for v2.2

src/StakeSablierNFT.sol

@@ -1,10 +1,12 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
-pragma solidity >=0.8.19;
+pragma solidity >=0.8.22;
 
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import { ERC721Holder } from "@openzeppelin/contracts/token/ERC721/utils/ERC721Holder.sol";
+import { IERC165 } from "@openzeppelin/contracts/utils/introspection/IERC165.sol";
 import { Adminable } from "@sablier/v2-core/src/abstracts/Adminable.sol";
+import { ISablierLockupRecipient } from "@sablier/v2-core/src/interfaces/ISablierLockupRecipient.sol";
 import { ISablierV2Lockup } from "@sablier/v2-core/src/interfaces/ISablierV2Lockup.sol";
 
 /// @title StakeSablierNFT
@@ -22,7 +24,7 @@ import { ISablierV2Lockup } from "@sablier/v2-core/src/interfaces/ISablierV2Lock
 ///   - The staking contract supports only one type of stream at a time, either Lockup Dynamic or Lockup Linear.
 ///   - The Sablier NFT must be transferable because staking requires transferring the NFT to the staking contract.
 ///   - This staking contract assumes that one user can only stake one NFT at a time.
-contract StakeSablierNFT is Adminable, ERC721Holder {
+contract StakeSablierNFT is Adminable, ERC721Holder, ISablierLockupRecipient {
     using SafeERC20 for IERC20;
 
     /*//////////////////////////////////////////////////////////////////////////
@@ -156,6 +158,11 @@ contract StakeSablierNFT is Adminable, ERC721Holder {
         return totalRewardPaidPerERC20Token + totalRewardsPerERC20InCurrentPeriod;
     }
 
+    // {IERC165-supportsInterface} implementation as required by `ISablierLockupRecipient` interface.
+    function supportsInterface(bytes4 interfaceId) public pure override(IERC165) returns (bool) {
+        return interfaceId == 0xf8ee98d3;
+    }
+
     /*//////////////////////////////////////////////////////////////////////////
                          USER-FACING NON-CONSTANT FUNCTIONS
     //////////////////////////////////////////////////////////////////////////*/
@@ -172,17 +179,19 @@ contract StakeSablierNFT is Adminable, ERC721Holder {
         }
     }
 
-    /// @notice Implements the hook to handle the cancelation of the stream.
+    /// @notice Implements the hook to handle cancelation events. This will be called by Sablier contract when a stream
+    /// is canceled by the sender.
     /// @dev This function subtracts the amount refunded to the sender from `totalERC20StakedSupply`.
     ///   - This function also updates the rewards for the staker.
-    function onStreamCanceled(
+    function onSablierLockupCancel(
         uint256 streamId,
-        address,
+        address, /* sender */
         uint128 senderAmount,
-        uint128
+        uint128 /* recipientAmount */
     )
         external
         updateReward(stakedAssets[streamId])
+        returns (bytes4 selector)
     {
         // Check: the caller is the lockup contract.
         if (msg.sender != address(sablierLockup)) {
@@ -191,18 +200,22 @@ contract StakeSablierNFT is Adminable, ERC721Holder {
 
         // Effect: update the total staked amount.
         totalERC20StakedSupply -= senderAmount;
+
+        return ISablierLockupRecipient.onSablierLockupCancel.selector;
     }
 
-    /// @notice Implements the hook to handle the withdrawn amount if sender calls the withdraw.
+    /// @notice Implements the hook to handle withdraw events. This will be called by Sablier contract when withdraw is
+    /// called on a stream.
     /// @dev This function transfers `amount` to the original staker.
-    function onStreamWithdrawn(
+    function onSablierLockupWithdraw(
         uint256 streamId,
-        address,
-        address,
+        address, /* caller */
+        address, /* recipient */
         uint128 amount
     )
         external
         updateReward(stakedAssets[streamId])
+        returns (bytes4 selector)
     {
         // Check: the caller is the lockup contract
         if (msg.sender != address(sablierLockup)) {
@@ -221,6 +234,8 @@ contract StakeSablierNFT is Adminable, ERC721Holder {
 
         // Interaction: transfer the withdrawn amount to the original staker.
         rewardERC20Token.safeTransfer(staker, amount);
+
+        return ISablierLockupRecipient.onSablierLockupWithdraw.selector;
     }
 
     /// @notice Stake a Sablier NFT with specified base asset.

test/stake-sablier-nft/StakeSablierNFT.t.sol

@@ -8,7 +8,6 @@ import { Broker, LockupLinear } from "@sablier/v2-core/src/types/DataTypes.sol";
 import { Test } from "forge-std/src/Test.sol";
 
 import { StakeSablierNFT } from "src/StakeSablierNFT.sol";
-import { console2 } from "forge-std/src/console2.sol";
 
 struct StreamOwner {
     address addr;
@@ -18,12 +17,12 @@ struct StreamOwner {
 struct Users {
     // Creator of the NFT staking contract.
     address admin;
-    // Alice is authorized to stake.
+    // Alice has already staked her NFT.
     StreamOwner alice;
     // Bob is unauthorized to stake.
     StreamOwner bob;
-    // Staker is the user we will test the contract for.
-    StreamOwner staker;
+    // Joe wants to stake his NFT.
+    StreamOwner joe;
 }
 
 abstract contract StakeSablierNFT_Fork_Test is Test {
@@ -44,12 +43,12 @@ abstract contract StakeSablierNFT_Fork_Test is Test {
     event Staked(address indexed user, uint256 tokenId);
     event Unstaked(address indexed user, uint256 tokenId);
 
-    IERC20 public constant DAI = IERC20(0x6B175474E89094C44Da98b954EedeAC495271d0F);
-    IERC20 public constant USDC = IERC20(0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48);
+    IERC20 public constant DAI = IERC20(0x776b6fC2eD15D6Bb5Fc32e0c89DE68683118c62A);
+    IERC20 public constant USDC = IERC20(0x1c7D4B196Cb0C7B01d743Fbc6116a902379C7238);
 
     // Get the latest deployment address from the docs: https://docs.sablier.com/contracts/v2/deployments.
     ISablierV2LockupLinear internal constant SABLIER =
-        ISablierV2LockupLinear(0xAFb979d9afAd1aD27C5eFf4E27226E3AB9e5dCC9);
+        ISablierV2LockupLinear(0x3E435560fd0a03ddF70694b35b673C25c65aBB6C);
 
     // Set a stream ID to stake.
     uint256 internal stakingStreamId = 2;
@@ -62,19 +61,19 @@ abstract contract StakeSablierNFT_Fork_Test is Test {
 
     StakeSablierNFT internal stakingContract;
 
-    uint256 internal tokenAmountsInStream;
+    uint256 internal constant AMOUNT_IN_STREAM = 1000e18;
 
     Users internal users;
 
     function setUp() public {
         // Fork Ethereum Mainnet.
-        vm.createSelectFork({ blockNumber: 19_689_210, urlOrAlias: "mainnet" });
+        vm.createSelectFork({ blockNumber: 6_239_031, urlOrAlias: "sepolia" });
 
         // Create users.
         users.admin = makeAddr("admin");
         users.alice.addr = makeAddr("alice");
         users.bob.addr = makeAddr("bob");
-        users.staker.addr = makeAddr("staker");
+        users.joe.addr = makeAddr("joe");
 
         // Mint some reward tokens to the admin address which will be used to deposit to the staking contract.
         deal({ token: address(rewardToken), to: users.admin, give: 10_000e18 });
@@ -86,22 +85,22 @@ abstract contract StakeSablierNFT_Fork_Test is Test {
         stakingContract =
             new StakeSablierNFT({ initialAdmin: users.admin, rewardERC20Token_: rewardToken, sablierLockup_: SABLIER });
 
+        // Set expected reward rate.
+        rewardRate = 10_000e18 / uint256(1 weeks);
+
         // Fund the staking contract with some reward tokens.
         rewardToken.transfer(address(stakingContract), 10_000e18);
 
         // Start the staking period.
         stakingContract.startStakingPeriod(10_000e18, 1 weeks);
 
-        // Set expected reward rate.
-        rewardRate = 10_000e18 / uint256(1 weeks);
-
         // Stake some streams.
         _createAndStakeStreamBy({ recipient: users.alice, asset: DAI, stake: true });
         _createAndStakeStreamBy({ recipient: users.bob, asset: USDC, stake: false });
-        _createAndStakeStreamBy({ recipient: users.staker, asset: DAI, stake: false });
+        _createAndStakeStreamBy({ recipient: users.joe, asset: DAI, stake: false });
 
         // Make the stream owner the `msg.sender` in all the subsequent calls.
-        resetPrank({ msgSender: users.staker.addr });
+        resetPrank({ msgSender: users.joe.addr });
 
         // Approve the staking contract to spend the NFT.
         SABLIER.setApprovalForAll(address(stakingContract), true);
@@ -114,7 +113,7 @@ abstract contract StakeSablierNFT_Fork_Test is Test {
     }
 
     function _createLockupLinearStreams(address recipient, IERC20 asset) private returns (uint256 streamId) {
-        deal({ token: address(asset), to: users.admin, give: 1000e18 });
+        deal({ token: address(asset), to: users.admin, give: AMOUNT_IN_STREAM });
 
         resetPrank({ msgSender: users.admin });
 
@@ -126,7 +125,7 @@ abstract contract StakeSablierNFT_Fork_Test is Test {
         // Declare the function parameters
         params.sender = users.admin; // The sender will be able to cancel the stream
         params.recipient = recipient; // The recipient of the streamed assets
-        params.totalAmount = uint128(1000e18); // Total amount is the amount inclusive of all fees
+        params.totalAmount = uint128(AMOUNT_IN_STREAM); // Total amount is the amount inclusive of all fees
         params.asset = asset; // The streaming asset
         params.cancelable = true; // Whether the stream will be cancelable or not
         params.transferable = true; // Whether the stream will be transferable or not

test/stake-sablier-nft/claim-rewards/claimRewards.t.sol

@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity >=0.8.19;
+
+import { StakeSablierNFT_Fork_Test } from "../StakeSablierNFT.t.sol";
+
+contract ClaimRewards_Test is StakeSablierNFT_Fork_Test {
+    function test_ClaimRewards_WhenNonStaker() external {
+        // Change the caller to a staker.
+        resetPrank({ msgSender: users.joe.addr });
+
+        // Expect no transfer.
+        vm.expectCall({
+            callee: address(rewardToken),
+            data: abi.encodeCall(rewardToken.transfer, (users.joe.addr, 0)),
+            count: 0
+        });
+
+        // Claim rewards.
+        stakingContract.claimRewards();
+    }
+
+    modifier givenStaked() {
+        // Change the caller to a staker.
+        resetPrank({ msgSender: users.alice.addr });
+
+        vm.warp(block.timestamp + 1 days);
+        _;
+    }
+
+    function test_ClaimRewards() external givenStaked {
+        uint256 expectedReward = 1 days * rewardRate;
+        uint256 initialBalance = rewardToken.balanceOf(users.alice.addr);
+
+        // Claim the rewards.
+        stakingContract.claimRewards();
+
+        // Assert balance increased by the expected reward.
+        uint256 finalBalance = rewardToken.balanceOf(users.alice.addr);
+        assertApproxEqAbs(finalBalance - initialBalance, expectedReward, 0.0001e18);
+
+        // Assert rewards has been set to 0.
+        assertEq(stakingContract.rewards(users.alice.addr), 0);
+    }
+}

test/stake-sablier-nft/claim-rewards/claimRewards.tree

@@ -0,0 +1,5 @@
+claimRewards.t.sol
+├── given the caller is not a staker
+│   └── it should not transfer the rewards
+└── given the caller is a staker
+    └── it should transfer the rewards