Skip to content

Commit

Permalink
Support multiple reward notifiers to the staker contract controlled b…
Browse files Browse the repository at this point in the history 
…y an admin
  • Loading branch information
@wildmolasses
wildmolasses authored Jan 31, 2024
1 parent fd2152d commit 57721b6
Show file tree
Hide file tree
Showing 2 changed files with 194 additions and 16 deletions.
54 changes: 48 additions & 6 deletions src/UniStaker.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,12 @@ contract UniStaker is INotifiableRewardReceiver, ReentrancyGuard {
/// @notice Thrown if a caller attempts to specify address zero for certain designated addresses.
error UniStaker__InvalidAddress();

/// @notice Emitted when the admin address is set.
event AdminSet(address indexed oldAdmin, address indexed newAdmin);

/// @notice Emitted when a rewards notifier address is enabled or disabled.
event RewardsNotifierSet(address indexed account, bool isEnabled);

/// @notice Metadata associated with a discrete staking deposit.
/// @param balance The deposit's staked balance.
/// @param owner The owner of this deposit.
Expand All @@ -62,8 +68,6 @@ contract UniStaker is INotifiableRewardReceiver, ReentrancyGuard {
/// @notice Delegable governance token which users stake to earn rewards.
IERC20Delegates public immutable STAKE_TOKEN;

address public immutable REWARDS_NOTIFIER;

/// @notice Length of time over which rewards sent to this contract are distributed to stakers.
uint256 public constant REWARD_DURATION = 7 days;

Expand All @@ -74,6 +78,9 @@ contract UniStaker is INotifiableRewardReceiver, ReentrancyGuard {
/// @dev Unique identifier that will be used for the next deposit.
DepositIdentifier private nextDepositId;

/// @notice Permissioned actor that can enable/disable `rewardsNotifier` addresses.
address public admin;

/// @notice Global amount currently staked across all user deposits.
uint256 public totalSupply;

Expand Down Expand Up @@ -116,13 +123,34 @@ contract UniStaker is INotifiableRewardReceiver, ReentrancyGuard {
/// account claims their earned rewards.
mapping(address account => uint256 amount) public rewards;

/// @notice Maps addresses to whether they are authorized to call `notifyRewardsAmount`.
mapping(address rewardsNotifier => bool) public isRewardsNotifier;

/// @param _rewardsToken ERC20 token in which rewards will be denominated.
/// @param _stakeToken Delegable governance token which users will stake to earn rewards.
/// @param _rewardsNotifier DEPRECATED
constructor(IERC20 _rewardsToken, IERC20Delegates _stakeToken, address _rewardsNotifier) {
/// @param _admin Address which will have permission to manage rewardsNotifiers.
constructor(IERC20 _rewardsToken, IERC20Delegates _stakeToken, address _admin) {
REWARDS_TOKEN = _rewardsToken;
STAKE_TOKEN = _stakeToken;
REWARDS_NOTIFIER = _rewardsNotifier;
_setAdmin(_admin);
}

/// @notice Set the admin address.
/// @param _newAdmin Address of the new admin.
/// @dev Caller must be the current admin.
function setAdmin(address _newAdmin) external {
_revertIfNotAdmin();
_setAdmin(_newAdmin);
}

/// @notice Enables or disables a rewards notifier address.
/// @param _rewardsNotifier Address of the rewards notifier.
/// @param _isEnabled `true` to enable the `_rewardsNotifier`, or `false` to disable.
/// @dev Caller must be the current admin.
function setRewardsNotifier(address _rewardsNotifier, bool _isEnabled) external {
_revertIfNotAdmin();
isRewardsNotifier[_rewardsNotifier] = _isEnabled;
emit RewardsNotifierSet(_rewardsNotifier, _isEnabled);
}

/// @notice Timestamp representing the last time at which rewards have been distributed, which is
Expand Down Expand Up @@ -279,7 +307,7 @@ contract UniStaker is INotifiableRewardReceiver, ReentrancyGuard {
/// staking contract before the rewards notifier calls this method.
/// @param _amount Quantity of reward tokens the staking contract is being notified of.
function notifyRewardsAmount(uint256 _amount) external {
if (msg.sender != REWARDS_NOTIFIER) revert UniStaker__Unauthorized("not notifier", msg.sender);
if (!isRewardsNotifier[msg.sender]) revert UniStaker__Unauthorized("not notifier", msg.sender);
// TODO: It looks like the only thing we actually need to do here is update the
// rewardPerTokenStored value. Can we save gas by doing only that?
_updateReward(address(0));
Expand Down Expand Up @@ -377,6 +405,20 @@ contract UniStaker is INotifiableRewardReceiver, ReentrancyGuard {
userRewardPerTokenPaid[_beneficiary] = rewardPerTokenStored;
}

/// @notice Internal helper method which sets the admin address.
/// @param _newAdmin Address of the new admin.
function _setAdmin(address _newAdmin) internal {
_revertIfAddressZero(_newAdmin);
emit AdminSet(admin, _newAdmin);
admin = _newAdmin;
}

/// @notice Internal helper method which reverts UniStaker__Unauthorized if the message sender is
/// not the admin.
function _revertIfNotAdmin() internal view {
if (msg.sender != admin) revert UniStaker__Unauthorized("not admin", msg.sender);
}

/// @notice Internal helper method which reverts UniStaker__Unauthorized if the message sender is
/// not the owner of the deposit.
/// @param deposit Deposit to validate.
Expand Down
156 changes: 146 additions & 10 deletions test/UniStaker.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,13 @@ import {ERC20Fake} from "test/fakes/ERC20Fake.sol";
contract UniStakerTest is Test {
ERC20Fake rewardToken;
ERC20VotesMock govToken;
address admin;
address rewardsNotifier;
UniStaker uniStaker;

event RewardsNotifierSet(address indexed account, bool isEnabled);
event AdminSet(address indexed oldAdmin, address indexed newAdmin);

function setUp() public {
// Set the block timestamp to an arbitrary value to avoid introducing assumptions into tests
// based on a starting timestamp of 0, which is the default.
Expand All @@ -26,15 +30,21 @@ contract UniStakerTest is Test {
rewardsNotifier = address(0xaffab1ebeef);
vm.label(rewardsNotifier, "Rewards Notifier");

uniStaker = new UniStaker(rewardToken, govToken, rewardsNotifier);
admin = makeAddr("admin");

uniStaker = new UniStaker(rewardToken, govToken, admin);

vm.prank(admin);
uniStaker.setRewardsNotifier(rewardsNotifier, true);

vm.label(address(uniStaker), "UniStaker");
}

function _jumpAhead(uint256 _seconds) public {
vm.warp(block.timestamp + _seconds);
}

function _boundMintAmount(uint256 _amount) internal view returns (uint256) {
function _boundMintAmount(uint256 _amount) internal pure returns (uint256) {
return bound(_amount, 0, 100_000_000_000e18);
}

Expand All @@ -45,7 +55,7 @@ contract UniStakerTest is Test {

function _boundToRealisticStake(uint256 _stakeAmount)
public
view
pure
returns (uint256 _boundedStakeAmount)
{
_boundedStakeAmount = bound(_stakeAmount, 0.1e18, 25_000_000e18);
Expand Down Expand Up @@ -115,19 +125,20 @@ contract Constructor is UniStakerTest {
function test_SetsTheRewardTokenStakeTokenAndRewardsNotifier() public {
assertEq(address(uniStaker.REWARDS_TOKEN()), address(rewardToken));
assertEq(address(uniStaker.STAKE_TOKEN()), address(govToken));
assertEq(uniStaker.REWARDS_NOTIFIER(), rewardsNotifier);
assertEq(uniStaker.admin(), admin);
}

function testFuzz_SetsTheRewardsTokenStakeTokenAndRewardsNotifierToArbitraryAddresses(
function testFuzz_SetsTheRewardsTokenStakeTokenAndOwnerToArbitraryAddresses(
address _rewardsToken,
address _stakeToken,
address _rewardsNotifier
address _admin
) public {
vm.assume(_admin != address(0));
UniStaker _uniStaker =
new UniStaker(IERC20(_rewardsToken), IERC20Delegates(_stakeToken), _rewardsNotifier);
new UniStaker(IERC20(_rewardsToken), IERC20Delegates(_stakeToken), _admin);
assertEq(address(_uniStaker.REWARDS_TOKEN()), address(_rewardsToken));
assertEq(address(_uniStaker.STAKE_TOKEN()), address(_stakeToken));
assertEq(_uniStaker.REWARDS_NOTIFIER(), _rewardsNotifier);
assertEq(_uniStaker.admin(), _admin);
}
}

Expand Down Expand Up @@ -1154,6 +1165,88 @@ contract Withdraw is UniStakerTest {
}
}

contract SetRewardsNotifier is UniStakerTest {
function testFuzz_AllowsAdminToSetRewardsNotifier(address _rewardsNotifier, bool _isEnabled)
public
{
vm.prank(admin);
uniStaker.setRewardsNotifier(_rewardsNotifier, _isEnabled);

assertEq(uniStaker.isRewardsNotifier(_rewardsNotifier), _isEnabled);
}

function test_AllowsTheAdminToDisableAnActiveRewardsNotifier() public {
vm.prank(admin);
uniStaker.setRewardsNotifier(rewardsNotifier, false);

assertFalse(uniStaker.isRewardsNotifier(rewardsNotifier));
}

function testFuzz_EmitsEventWhenRewardsNotifierIsSet(address _rewardsNotifier, bool _isEnabled)
public
{
vm.expectEmit();
emit RewardsNotifierSet(_rewardsNotifier, _isEnabled);
vm.prank(admin);
uniStaker.setRewardsNotifier(_rewardsNotifier, _isEnabled);
}

function testFuzz_RevertIf_TheCallerIsNotTheAdmin(
address _notAdmin,
address _newRewardsNotifier,
bool _isEnabled
) public {
vm.assume(_notAdmin != uniStaker.admin());

vm.prank(_notAdmin);
vm.expectRevert(
abi.encodeWithSelector(
UniStaker.UniStaker__Unauthorized.selector, bytes32("not admin"), _notAdmin
)
);
uniStaker.setRewardsNotifier(_newRewardsNotifier, _isEnabled);
}
}

contract SetAdmin is UniStakerTest {
function testFuzz_AllowsAdminToSetAdmin(address _newAdmin) public {
vm.assume(_newAdmin != address(0));

vm.prank(admin);
uniStaker.setAdmin(_newAdmin);

assertEq(uniStaker.admin(), _newAdmin);
}

function testFuzz_EmitsEventWhenAdminIsSet(address _newAdmin) public {
vm.assume(_newAdmin != address(0));

vm.expectEmit();
emit AdminSet(admin, _newAdmin);

vm.prank(admin);
uniStaker.setAdmin(_newAdmin);
}

function testFuzz_RevertIf_TheCallerIsNotTheAdmin(address _notAdmin, address _newAdmin) public {
vm.assume(_notAdmin != uniStaker.admin());

vm.prank(_notAdmin);
vm.expectRevert(
abi.encodeWithSelector(
UniStaker.UniStaker__Unauthorized.selector, bytes32("not admin"), _notAdmin
)
);
uniStaker.setAdmin(_newAdmin);
}

function test_RevertIf_NewAdminAddressIsZeroAddress() public {
vm.prank(admin);
vm.expectRevert(UniStaker.UniStaker__InvalidAddress.selector);
uniStaker.setAdmin(address(0));
}
}

contract UniStakerRewardsTest is UniStakerTest {
// Because there will be (expected) rounding errors in the amount of rewards earned, this helper
// checks that the truncated number is lesser and within 1% of the expected number.
Expand Down Expand Up @@ -1270,6 +1363,16 @@ contract UniStakerRewardsTest is UniStakerTest {
uniStaker.notifyRewardsAmount(_amount);
vm.stopPrank();
}

function _mintTransferAndNotifyReward(address _rewardsNotifier, uint256 _amount) public {
vm.assume(_rewardsNotifier != address(0));
rewardToken.mint(_rewardsNotifier, _amount);

vm.startPrank(_rewardsNotifier);
rewardToken.transfer(address(uniStaker), _amount);
uniStaker.notifyRewardsAmount(_amount);
vm.stopPrank();
}
}

contract NotifyRewardsAmount is UniStakerRewardsTest {
Expand Down Expand Up @@ -1338,10 +1441,42 @@ contract NotifyRewardsAmount is UniStakerRewardsTest {
assertEq(uniStaker.rewardPerTokenStored(), uniStaker.rewardPerToken());
}

function testFuzz_RevertIf_CallerIsNotTheRewardsNotifier(uint256 _amount, address _notNotifier)
function testFuzz_AllowsMultipleApprovedRewardsNotifiersToNotifyOfRewards(
uint256 _amount1,
uint256 _amount2,
uint256 _amount3,
address _rewardsNotifier1,
address _rewardsNotifier2,
address _rewardsNotifier3
) public {
_amount1 = _boundToRealisticReward(_amount1);
_amount2 = _boundToRealisticReward(_amount2);
_amount3 = _boundToRealisticReward(_amount3);

vm.startPrank(admin);
uniStaker.setRewardsNotifier(_rewardsNotifier1, true);
uniStaker.setRewardsNotifier(_rewardsNotifier2, true);
uniStaker.setRewardsNotifier(_rewardsNotifier3, true);
vm.stopPrank();

// The first notifier notifies
_mintTransferAndNotifyReward(_rewardsNotifier1, _amount1);

// The second notifier notifies
_mintTransferAndNotifyReward(_rewardsNotifier2, _amount2);

// The third notifier notifies
_mintTransferAndNotifyReward(_rewardsNotifier3, _amount3);
uint256 _expectedRewardRate = (_amount1 + _amount2 + _amount3) / uniStaker.REWARD_DURATION();
// because we summed 3 amounts, the rounding error can be as much as 2 units
assertApproxEqAbs(uniStaker.rewardRate(), _expectedRewardRate, 2);
assertLe(uniStaker.rewardRate(), _expectedRewardRate);
}

function testFuzz_RevertIf_CallerIsNotARewardsNotifier(uint256 _amount, address _notNotifier)
public
{
vm.assume(_notNotifier != rewardsNotifier && _notNotifier != address(0));
vm.assume(!uniStaker.isRewardsNotifier(_notNotifier) && _notNotifier != address(0));
_amount = _boundToRealisticReward(_amount);

rewardToken.mint(_notNotifier, _amount);
Expand Down Expand Up @@ -1378,6 +1513,7 @@ contract NotifyRewardsAmount is UniStakerRewardsTest {
// an amount - 1 because rounding errors when calculating the reward rate, which favor the
// staking contract can actually allow for something just below the amount to meet the criteria
_transferPercent = _bound(_transferPercent, 1, 99);

uint256 _transferAmount = _percentOf(_amount, _transferPercent);
rewardToken.mint(rewardsNotifier, _amount);

Expand Down

0 comments on commit 57721b6

Please sign in to comment.