diff --git a/src/UniStaker.sol b/src/UniStaker.sol index 4cd6690..7a29516 100644 --- a/src/UniStaker.sol +++ b/src/UniStaker.sol @@ -566,8 +566,9 @@ contract UniStaker is INotifiableRewardReceiver, Multicall, EIP712, Nonces { /// @notice Claim reward tokens the message sender has earned as a stake beneficiary. Tokens are /// sent to the message sender. - function claimReward() external { - _claimReward(msg.sender); + /// @return Amount of reward tokens claimed. + function claimReward() external returns (uint256) { + return _claimReward(msg.sender); } /// @notice Claim earned reward tokens for a beneficiary, using a signature to validate the @@ -575,8 +576,10 @@ contract UniStaker is INotifiableRewardReceiver, Multicall, EIP712, Nonces { /// @param _beneficiary Address of the beneficiary who will receive the reward. /// @param _deadline The timestamp at which the signature should expire. /// @param _signature Signature of the beneficiary authorizing this reward claim. + /// @return Amount of reward tokens claimed. function claimRewardOnBehalf(address _beneficiary, uint256 _deadline, bytes memory _signature) external + returns (uint256) { _revertIfPastDeadline(_deadline); _revertIfSignatureIsNotValidNow( @@ -588,7 +591,7 @@ contract UniStaker is INotifiableRewardReceiver, Multicall, EIP712, Nonces { ), _signature ); - _claimReward(_beneficiary); + return _claimReward(_beneficiary); } /// @notice Called by an authorized rewards notifier to alert the staking contract that a new @@ -794,18 +797,20 @@ contract UniStaker is INotifiableRewardReceiver, Multicall, EIP712, Nonces { } /// @notice Internal convenience method which claims earned rewards. + /// @return Amount of reward tokens claimed. /// @dev This method must only be called after proper authorization has been completed. /// @dev See public claimReward methods for additional documentation. - function _claimReward(address _beneficiary) internal { + function _claimReward(address _beneficiary) internal returns (uint256) { _checkpointGlobalReward(); _checkpointReward(_beneficiary); uint256 _reward = scaledUnclaimedRewardCheckpoint[_beneficiary] / SCALE_FACTOR; - if (_reward == 0) return; + if (_reward == 0) return 0; scaledUnclaimedRewardCheckpoint[_beneficiary] = 0; emit RewardClaimed(_beneficiary, _reward); SafeERC20.safeTransfer(REWARD_TOKEN, _beneficiary, _reward); + return _reward; } /// @notice Checkpoints the global reward per token accumulator. diff --git a/test/UniStaker.t.sol b/test/UniStaker.t.sol index 9d13901..3f8a5d9 100644 --- a/test/UniStaker.t.sol +++ b/test/UniStaker.t.sol @@ -4905,6 +4905,33 @@ contract ClaimReward is UniStakerRewardsTest { assertEq(rewardToken.balanceOf(_depositor), _earned); } + function testFuzz_ReturnsRewardAmount( + address _depositor, + address _delegatee, + uint256 _stakeAmount, + uint256 _rewardAmount, + uint256 _durationPercent + ) public { + vm.assume(_depositor != address(uniStaker)); + + (_stakeAmount, _rewardAmount) = _boundToRealisticStakeAndReward(_stakeAmount, _rewardAmount); + _durationPercent = bound(_durationPercent, 0, 100); + + // A user deposits staking tokens + _boundMintAndStake(_depositor, _stakeAmount, _delegatee); + // The contract is notified of a reward + _mintTransferAndNotifyReward(_rewardAmount); + // A portion of the duration passes + _jumpAheadByPercentOfRewardDuration(_durationPercent); + + uint256 _earned = uniStaker.unclaimedReward(_depositor); + + vm.prank(_depositor); + uint256 _claimedAmount = uniStaker.claimReward(); + + assertEq(_earned, _claimedAmount); + } + function testFuzz_ResetsTheRewardsEarnedByTheUser( address _depositor, address _delegatee, @@ -5008,6 +5035,56 @@ contract ClaimRewardOnBehalf is UniStakerRewardsTest { assertEq(rewardToken.balanceOf(_beneficiary), _earned); } + function testFuzz_ReturnsClaimedRewardAmount( + uint256 _beneficiaryPrivateKey, + address _sender, + uint256 _depositAmount, + uint256 _durationPercent, + uint256 _rewardAmount, + address _delegatee, + address _depositor, + uint256 _currentNonce, + uint256 _deadline + ) public { + vm.assume(_delegatee != address(0) && _depositor != address(0) && _sender != address(0)); + _beneficiaryPrivateKey = bound(_beneficiaryPrivateKey, 1, 100e18); + address _beneficiary = vm.addr(_beneficiaryPrivateKey); + + UniStaker.DepositIdentifier _depositId; + (_depositAmount, _rewardAmount) = _boundToRealisticStakeAndReward(_depositAmount, _rewardAmount); + _durationPercent = bound(_durationPercent, 0, 100); + + // A user deposits staking tokens + (_depositAmount, _depositId) = + _boundMintAndStake(_depositor, _depositAmount, _delegatee, _beneficiary); + // The contract is notified of a reward + _mintTransferAndNotifyReward(_rewardAmount); + // A portion of the duration passes + _jumpAheadByPercentOfRewardDuration(_durationPercent); + _deadline = bound(_deadline, block.timestamp, type(uint256).max); + + uint256 _earned = uniStaker.unclaimedReward(_beneficiary); + + stdstore.target(address(uniStaker)).sig("nonces(address)").with_key(_beneficiary).checked_write( + _currentNonce + ); + + bytes32 _message = keccak256( + abi.encode( + uniStaker.CLAIM_REWARD_TYPEHASH(), _beneficiary, uniStaker.nonces(_beneficiary), _deadline + ) + ); + + bytes32 _messageHash = + keccak256(abi.encodePacked("\x19\x01", EIP712_DOMAIN_SEPARATOR, _message)); + bytes memory _signature = _sign(_beneficiaryPrivateKey, _messageHash); + + vm.prank(_sender); + uint256 _claimedAmount = uniStaker.claimRewardOnBehalf(_beneficiary, _deadline, _signature); + + assertEq(_earned, _claimedAmount); + } + function testFuzz_RevertIf_WrongNonceIsUsed( uint256 _beneficiaryPrivateKey, address _sender,