From 505b72c28da87b947affcf381dafebb59dddb7cf Mon Sep 17 00:00:00 2001 From: smol-ninja Date: Mon, 22 Apr 2024 23:12:00 +0100 Subject: [PATCH] docs: add DISCLAIMER --- src/StakeSablierNFT.sol | 81 ++++++++++---------- test/stake-sablier-nft/StakeSablierNFT.t.sol | 3 +- test/stake-sablier-nft/claim/claim.t.sol | 41 +++------- test/stake-sablier-nft/stake/stake.t.sol | 2 +- 4 files changed, 54 insertions(+), 73 deletions(-) diff --git a/src/StakeSablierNFT.sol b/src/StakeSablierNFT.sol index a2f2186..fc84b34 100644 --- a/src/StakeSablierNFT.sol +++ b/src/StakeSablierNFT.sol @@ -6,6 +6,10 @@ import { Adminable } from "@sablier/v2-core/src/abstracts/Adminable.sol"; import { ISablierV2Lockup } from "@sablier/v2-core/src/interfaces/ISablierV2Lockup.sol"; /// @title StakeSablierNFT +/// @notice DISCLAIMER: This template has not been audited and is provided "as is" with no warranties of any kind, +/// either express or implied. It is intended solely for demonstration purposes on how to build a staking contract +/// using Sablier NFT. This template should not be used in a production environment. It makes specific assumptions that +/// may not be applicable to your particular needs. /// @dev This template allows users to stake Sablier NFTs and earn staking rewards. /// /// Requirements: @@ -22,8 +26,7 @@ contract StakeSablierNFT is Adminable { //////////////////////////////////////////////////////////////////////////*/ error ClaimAmountExceedsBalance(uint256 claimAmount, uint256 balance); - error NotAuthorized(uint256 tokenId); - error NotStaked(uint256 tokenId); + error InvalidToken(IERC20 streamingToken, IERC20 rewardToken); error NotStreamOwner(address account, uint256 tokenId); error ZeroAmount(); @@ -64,26 +67,6 @@ contract StakeSablierNFT is Adminable { /// @dev The owner of the Sablier stream mapped by tokenId. mapping(uint256 tokenId => address owner) public streamOwner; - /*////////////////////////////////////////////////////////////////////////// - MODIFIERS - //////////////////////////////////////////////////////////////////////////*/ - - modifier onlyStoredStreamOwner(uint256 tokenId) { - // Check: if the `msg.sender` is the stored owner of the Sablier Stream - if (streamOwner[tokenId] != msg.sender) { - revert NotStreamOwner(msg.sender, tokenId); - } - _; - } - - modifier onlyStreamRecipient(uint256 tokenId) { - // Check: if the `msg.sender` is the recipient of the Sablier Stream - if (SABLIER_CONTRACT.getRecipient(tokenId) != msg.sender) { - revert NotStreamOwner(msg.sender, tokenId); - } - _; - } - /*////////////////////////////////////////////////////////////////////////// CONSTRUCTOR //////////////////////////////////////////////////////////////////////////*/ @@ -102,22 +85,27 @@ contract StakeSablierNFT is Adminable { USER-FACING NON-CONSTANT FUNCTIONS //////////////////////////////////////////////////////////////////////////*/ - /// @notice Claim the staking rewards for the `tokenId` when Sablier NFT is staked. - /// @dev This function can only be called by the original owner of the Sablier NFT - /// @param tokenId The tokenId of the Sablier NFT to withdraw rewards for - function claimWhenStaked(uint256 tokenId) public onlyStoredStreamOwner(tokenId) { - // Effect: update the claimable rewards - uint256 claimAmount = updateClaimAmount(tokenId); - - _claim(claimAmount, tokenId); - } - - /// @notice Claim the staking rewards for the `tokenId` after Sablier NFT has been unstaked. - /// @dev This function can only be called by the current recipient of the Sablier NFT + /// @notice Claim the staking rewards for the `tokenId`. + /// @dev + /// If NFT is staked: + /// - This function should be called by the original owner of the Sablier NFT + /// - The staking rewards are updated before claiming + /// If NFT is not staked: + /// - This function should be called by the current recipient of the Sablier NFT + /// - The staking rewards are loaded from storage /// @param tokenId The tokenId of the Sablier NFT to withdraw rewards for - function claimWhenUnstaked(uint256 tokenId) public onlyStreamRecipient(tokenId) { - // Load the staking rewards from storage - uint256 claimAmount = stakingRewards[tokenId]; + function claim(uint256 tokenId) public { + uint256 claimAmount; + + if (streamOwner[tokenId] == msg.sender) { + // Effect: update the claimable rewards since the stream is staked + claimAmount = updateClaimAmount(tokenId); + } else if (SABLIER_CONTRACT.getRecipient(tokenId) == msg.sender) { + // Load the staking rewards from storage since the stream is not staked + claimAmount = stakingRewards[tokenId]; + } else { + revert NotStreamOwner(msg.sender, tokenId); + } _claim(claimAmount, tokenId); } @@ -125,10 +113,16 @@ contract StakeSablierNFT is Adminable { /// @notice Stake a Sablier NFT with specified base asset /// @dev The `msg.sender` must approve the staking contract to spend the Sablier NFT before calling this function /// @param tokenId The tokenId of the Sablier NFT to be staked - function stake(uint256 tokenId) public onlyStreamRecipient(tokenId) { + function stake(uint256 tokenId) public { + // Check: if the `msg.sender` is the recipient of the Sablier Stream + if (SABLIER_CONTRACT.getRecipient(tokenId) != msg.sender) { + revert NotStreamOwner(msg.sender, tokenId); + } + // Check: if the Sablier NFT was minted with the staking asset - if (SABLIER_CONTRACT.getAsset(tokenId) != REWARD_TOKEN) { - revert NotAuthorized(tokenId); + IERC20 streamingAsset = IERC20(SABLIER_CONTRACT.getAsset(tokenId)); + if (streamingAsset != REWARD_TOKEN) { + revert InvalidToken(streamingAsset, REWARD_TOKEN); } // Effect: store the owner of the Sablier NFT @@ -146,7 +140,12 @@ contract StakeSablierNFT is Adminable { /// @notice Unstaking a Sablier NFT will transfer the NFT back to the `msg.sender`. /// @dev This function can only be called by the original owner of the Sablier NFT /// @param tokenId The tokenId of the Sablier NFT to be unstaked - function unstake(uint256 tokenId) public onlyStoredStreamOwner(tokenId) { + function unstake(uint256 tokenId) public { + // Check: if the `msg.sender` is the stored owner of the Sablier Stream + if (streamOwner[tokenId] != msg.sender) { + revert NotStreamOwner(msg.sender, tokenId); + } + // Effect: update the claimable rewards updateClaimAmount(tokenId); diff --git a/test/stake-sablier-nft/StakeSablierNFT.t.sol b/test/stake-sablier-nft/StakeSablierNFT.t.sol index c1e7fac..b91d045 100644 --- a/test/stake-sablier-nft/StakeSablierNFT.t.sol +++ b/test/stake-sablier-nft/StakeSablierNFT.t.sol @@ -10,8 +10,7 @@ import { StakeSablierNFT } from "src/StakeSablierNFT.sol"; abstract contract StakeSablierNFT_Fork_Test is Test { // Errors error ClaimAmountExceedsBalance(uint256 claimAmount, uint256 balance); - error NotAuthorized(uint256); - error NotStaked(uint256); + error InvalidToken(IERC20 streamingToken, IERC20 rewardToken); error NotStreamOwner(address, uint256); error ZeroAmount(); diff --git a/test/stake-sablier-nft/claim/claim.t.sol b/test/stake-sablier-nft/claim/claim.t.sol index 7c37ae4..89cb123 100644 --- a/test/stake-sablier-nft/claim/claim.t.sol +++ b/test/stake-sablier-nft/claim/claim.t.sol @@ -4,31 +4,27 @@ pragma solidity >=0.8.19; import { StakeSablierNFT_Fork_Test } from "../StakeSablierNFT.t.sol"; contract Claim_Test is StakeSablierNFT_Fork_Test { - /*////////////////////////////////////////////////////////////////////////// - claimWhenStaked TEST - //////////////////////////////////////////////////////////////////////////*/ - modifier givenStaked() { stakingContract.stake(existingStreamId); _; } - function test_RevertWhen_CallerNotStaker() external givenStaked { + function test_RevertWhen_CallerUnauthorized() external givenStaked { address unauthorizedCaller = makeAddr("Unauthorized"); // Change the caller to an unauthorized address vm.startPrank({ msgSender: unauthorizedCaller }); vm.expectRevert(abi.encodeWithSelector(NotStreamOwner.selector, unauthorizedCaller, existingStreamId)); - stakingContract.claimWhenStaked(existingStreamId); + stakingContract.claim(existingStreamId); } - modifier whenCallerIsStaker() { + modifier whenAuthorizedCaller() { _; } - function test_RevertWhen_ClaimAmountZero() external whenCallerIsStaker givenStaked { + function test_RevertWhen_ClaimAmountZero() external whenAuthorizedCaller givenStaked { vm.expectRevert(abi.encodeWithSelector(ZeroAmount.selector)); - stakingContract.claimWhenStaked(existingStreamId); + stakingContract.claim(existingStreamId); } modifier whenClaimAmountNotZero() { @@ -37,7 +33,7 @@ contract Claim_Test is StakeSablierNFT_Fork_Test { function test_RevertWhen_ContractBalanceIsLessThanClaimAmount() external - whenCallerIsStaker + whenAuthorizedCaller whenClaimAmountNotZero givenStaked { @@ -60,7 +56,7 @@ contract Claim_Test is StakeSablierNFT_Fork_Test { vm.expectRevert(abi.encodeWithSelector(ClaimAmountExceedsBalance.selector, expectedReward, balance)); // Claim rewards - stakingContract.claimWhenStaked(existingStreamId); + stakingContract.claim(existingStreamId); } modifier whenContractBalanceIsNotLessThanClaimAmount() { @@ -69,7 +65,7 @@ contract Claim_Test is StakeSablierNFT_Fork_Test { function test_Claim_GivenStaked() external - whenCallerIsStaker + whenAuthorizedCaller whenClaimAmountNotZero whenContractBalanceIsNotLessThanClaimAmount givenStaked @@ -96,16 +92,12 @@ contract Claim_Test is StakeSablierNFT_Fork_Test { emit Transfer(address(stakingContract), staker, expectedReward); // Claim rewards - stakingContract.claimWhenStaked(existingStreamId); + stakingContract.claim(existingStreamId); // Assert: staker received the staking rewards assertEq(stakingContract.stakingRewards(existingStreamId), 0); } - /*////////////////////////////////////////////////////////////////////////// - claimWhenUnstaked TEST - //////////////////////////////////////////////////////////////////////////*/ - modifier givenUnstaked() { stakingContract.stake(existingStreamId); @@ -116,25 +108,16 @@ contract Claim_Test is StakeSablierNFT_Fork_Test { _; } - function test_RevertWhen_CallerNotRecipient() external givenUnstaked { - address unauthorizedCaller = makeAddr("Unauthorized"); - // Change the caller to an unauthorized address - vm.startPrank({ msgSender: unauthorizedCaller }); - - vm.expectRevert(abi.encodeWithSelector(NotStreamOwner.selector, unauthorizedCaller, existingStreamId)); - stakingContract.claimWhenUnstaked(existingStreamId); - } - modifier whenCallerIsRecipient() { _; } function test_RevertWhen_ClaimAmountZero_GivenUnstaked() external whenCallerIsRecipient givenUnstaked { // claim rewards so that claim amount becomes zero - stakingContract.claimWhenUnstaked(existingStreamId); + stakingContract.claim(existingStreamId); vm.expectRevert(abi.encodeWithSelector(ZeroAmount.selector)); - stakingContract.claimWhenUnstaked(existingStreamId); + stakingContract.claim(existingStreamId); } function test_Claim_GivenUnstaked() @@ -162,7 +145,7 @@ contract Claim_Test is StakeSablierNFT_Fork_Test { emit Transfer(address(stakingContract), staker, expectedReward); // Claim rewards - stakingContract.claimWhenUnstaked(existingStreamId); + stakingContract.claim(existingStreamId); // Assert: staker received the staking rewards assertEq(stakingContract.stakingRewards(existingStreamId), 0); diff --git a/test/stake-sablier-nft/stake/stake.t.sol b/test/stake-sablier-nft/stake/stake.t.sol index 63c2452..d96e791 100644 --- a/test/stake-sablier-nft/stake/stake.t.sol +++ b/test/stake-sablier-nft/stake/stake.t.sol @@ -30,7 +30,7 @@ contract Stake_Test is StakeSablierNFT_Fork_Test { // Change the caller to the staker again vm.startPrank({ msgSender: staker }); - vm.expectRevert(abi.encodeWithSelector(NotAuthorized.selector, existingStreamId)); + vm.expectRevert(abi.encodeWithSelector(InvalidToken.selector, sablier.getAsset(existingStreamId), token)); stakingContract.stake(existingStreamId); }