Distribute rewards tests

src/RewardsDistributor.sol

@@ -1,24 +1,25 @@
 // SPDX-License-Identifier: MIT
 pragma solidity 0.8.26;
 
-import {console} from "@forge-std/console.sol";
-import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
-import {SafeTransferLib} from "@solmate/utils/SafeTransferLib.sol";
 import {Ownable2StepUpgradeable} from "@openzeppelin-upgradeable/contracts/access/Ownable2StepUpgradeable.sol";
+import {IRewardsDistributor} from "./interfaces/IRewardsDistributor.sol";
 
-interface IRewardsDistributor {
-    function distributeRewards() external;
-}
-
-// TODO should be pausable?
-contract RewardsDistributor is Ownable2StepUpgradeable {
+contract RewardsDistributor is Ownable2StepUpgradeable, IRewardsDistributor {
     /*//////////////////////////////////////////////////////////////
                                LIBRARIES
     //////////////////////////////////////////////////////////////*/
     using SafeERC20 for IERC20;
 
+    /*//////////////////////////////////////////////////////////////
+                                 VARIABLES
+    //////////////////////////////////////////////////////////////*/
+
+    /// @notice the reward token, i.e. SHU
+    /// @dev set in initialize, can't be changed
+    IERC20 public rewardToken;
+
     /*//////////////////////////////////////////////////////////////
                                  STRUCTS
     //////////////////////////////////////////////////////////////*/
@@ -33,125 +34,104 @@
                              MAPPINGS/ARRAYS
     //////////////////////////////////////////////////////////////*/
 
-    mapping(address receiver => mapping(address token => RewardConfiguration configuration))
+    mapping(address receiver => RewardConfiguration configuration)
         public rewardConfigurations;
 
-    mapping(address receiver => address[] rewardsTokens) public rewardTokens;
-
     /*//////////////////////////////////////////////////////////////
                                  EVENTS
     //////////////////////////////////////////////////////////////*/
 
     event RewardConfigurationSet(
         address indexed receiver,
-        address indexed token,
         uint256 emissionRate
     );
 
-    event RewardDistributed(
-        address indexed receiver,
-        address indexed token,
-        uint256 reward
-    );
+    event RewardCollected(address indexed receiver, uint256 reward);
 
-    /// @notice Ensure logic contract is unusable
-    constructor() {
-        _disableInitializers();
-    }
+    event RewardTokenSet(address indexed rewardToken);
+
+    /*//////////////////////////////////////////////////////////////
+                                 ERRORS
+    //////////////////////////////////////////////////////////////*/
+
+    /// @notice Thrown when address is zero
+    error ZeroAddress();
 
     /// @notice Initialize the contract
     /// @param newOwner The owner of the contract, i.e. the DAO contract address
-    function initialize(address newOwner) public initializer {
-        __Ownable2Step_init();
-
+    constructor(address newOwner, address _rewardToken) {
         // Transfer ownership to the DAO contract
         _transferOwnership(newOwner);
-    }
-
-    /// @notice Add a reward configuration
-    /// @param receiver The receiver of the rewards
-    /// @param token The reward token
-    /// @param emissionRate The emission rate
-    function setRewardConfiguration(
-        address receiver,
-        address token,
-        uint256 emissionRate
-    ) external onlyOwner {
-        require(receiver != address(0), "Invalid receiver");
-        require(token != address(0), "No native rewards allowed");
-
-        if (rewardConfigurations[receiver][token].emissionRate == 0) {
-            rewardTokens[receiver].push(token);
-        }
-
-        rewardConfigurations[receiver][token] = RewardConfiguration(
-            emissionRate,
-            block.timestamp
-        );
-
-        if (emissionRate == 0) {
-            // remove the token
-            address[] storage tokens = rewardTokens[receiver];
-            for (uint256 i = 0; i < tokens.length; i++) {
-                if (tokens[i] == token) {
-                    tokens[i] = tokens[tokens.length - 1];
-                    tokens.pop();
-                    break;
-                }
-            }
-        }
 
-        emit RewardConfigurationSet(receiver, token, emissionRate);
+        // Set the reward token
+        rewardToken = IERC20(_rewardToken);
     }
 
     /// @notice Distribute rewards to receiver
-    /// @param token The reward token
-    function distributeReward(address token) external {
-        distributeRewardInternal(msg.sender, token);
-    }
-
-    /// @notice Distribute rewards to all tokens
-    function distributeRewards() external {
+    function collectRewards() external override returns (uint256 rewards) {
         address receiver = msg.sender;
 
-        for (uint256 i = 0; i < rewardTokens[receiver].length; i++) {
-            distributeRewardInternal(receiver, rewardTokens[receiver][i]);
-        }
-    }
-
-    /// @notice Distribute rewards to token
-    /// @param receiver The receiver of the rewards
-    /// @param token The reward token
-    function distributeRewardInternal(
-        address receiver,
-        address token
-    ) internal {
         RewardConfiguration storage rewardConfiguration = rewardConfigurations[
             receiver
-        ][token];
+        ];
 
         // difference in time since last update
         uint256 timeDelta = block.timestamp - rewardConfiguration.lastUpdate;
 
-        if (timeDelta == 0) {
+        if (rewardConfiguration.emissionRate == 0 || timeDelta == 0) {
             // nothing to do
-            return;
+            return 0;
         }
 
-        uint256 reward = rewardConfiguration.emissionRate * timeDelta;
+        rewards = rewardConfiguration.emissionRate * timeDelta;
 
         // update the last update timestamp
         rewardConfiguration.lastUpdate = block.timestamp;
 
         // transfer the reward
-        IERC20(token).safeTransfer(receiver, reward);
+        rewardToken.safeTransfer(receiver, rewards);
+
+        emit RewardCollected(receiver, rewards);
+    }
+
+    /// @notice Add a reward configuration
+    /// @param receiver The receiver of the rewards
+    /// @param emissionRate The emission rate
+    function setRewardConfiguration(
+        address receiver,
+        uint256 emissionRate
+    ) external override onlyOwner {
+        require(receiver != address(0), ZeroAddress());
+
+        rewardConfigurations[receiver] = RewardConfiguration(
+            emissionRate,
+            block.timestamp
+        );
+
+        emit RewardConfigurationSet(receiver, emissionRate);
+    }
 
-        emit RewardDistributed(receiver, token, reward);
+    /// @notice Withdraw funds from the contract
+    /// @param to The address to withdraw to
+    /// @param amount The amount to withdraw
+    function withdrawFunds(
+        address to,
+        uint256 amount
+    ) public override onlyOwner {
+        rewardToken.safeTransfer(to, amount);
     }
 
-    function getRewardTokens(
-        address receiver
-    ) external view returns (address[] memory) {
-        return rewardTokens[receiver];
+    /// @notice Set the reward token
+    /// @param _rewardToken The reward token
+    function setRewardToken(address _rewardToken) external onlyOwner {
+        require(_rewardToken != address(0), ZeroAddress());
+
+        // withdraw remaining old reward token
+        withdrawFunds(msg.sender, rewardToken.balanceOf(address(this)));
+
+        // set the new reward token
+        rewardToken = IERC20(_rewardToken);
+
+        emit RewardTokenSet(_rewardToken);
     }
 }

src/Staking.sol

@@ -119,7 +119,7 @@
     /// amount is less than the minimum stake set by the DAO
     error FirstStakeLessThanMinStake();
 
-    /// @notice Trownn when amount is zero
+    /// @notice Trown when amount is zero
     error ZeroAmount();
 
     /// @notice Thrown when someone try to unstake a stake that doesn't belong
@@ -146,14 +146,14 @@
 
     /// @notice Ensure only keypers can stake
     modifier onlyKeyper() {
         require(keypers[msg.sender], OnlyKeyper());
         _;
     }
 
     /// @notice Update rewards for a keyper
     modifier updateRewards() {
         // Distribute rewards
-        rewardsDistributor.distributeReward(address(stakingToken));
+        rewardsDistributor.collectRewards();
 
         _;
     }
@@ -201,7 +201,7 @@
         uint256 amount
     ) external onlyKeyper updateRewards returns (uint256) {
         /////////////////////////// CHECKS ///////////////////////////////
         require(amount > 0, ZeroAmount());
 
         address keyper = msg.sender;
 
@@ -210,7 +210,7 @@
 
         // If the keyper has no stakes, the first stake must be at least the minimum stake
         if (stakesIds.length() == 0) {
             require(amount >= minStake, FirstStakeLessThanMinStake());
         }
 
         /////////////////////////// EFFECTS ///////////////////////////////
@@ -269,20 +269,20 @@
         uint256 stakeId,
         uint256 amount
     ) external updateRewards {
         require(
             keyperStakes[keyper].contains(stakeId),
             StakeDoesNotBelongToKeyper()
         );
         Stake memory keyperStake = stakes[stakeId];
 
         require(keyperStake.amount > 0, StakeDoesNotExist());
 
         // If caller doesn't specify the amount, the contract will transfer the
         // stake amount for the stakeId
         if (amount == 0) {
             amount = keyperStake.amount;
         } else {
             require(amount <= keyperStake.amount, WithdrawAmountTooHigh());
         }
 
         // Checks below only apply if keyper is still a keyper
@@ -290,12 +290,12 @@
         // ignored and minStake is not enforced
         if (keypers[keyper]) {
             // Only the keyper can unstake
             require(msg.sender == keyper, OnlyKeyper());
 
             // If the lock period is less than the global lock period, the stake
             // must be locked for the lock period
             if (lockPeriod < keyperStake.lockPeriod) {
                 require(
                     block.timestamp > keyperStake.timestamp + lockPeriod,
                     StakeIsStillLocked()
                 );

src/interfaces/IRewardsDistributor.sol

@@ -2,22 +2,14 @@
 pragma solidity 0.8.26;
 
 interface IRewardsDistributor {
+    function collectRewards() external returns (uint256);
+
+    function withdrawFunds(address to, uint256 amount) external;
+
     function setRewardConfiguration(
         address receiver,
-        address token,
         uint256 emissionRate
     ) external;
 
-    function distributeReward(address token) external;
-
-    function distributeRewards() external;
-
-    function getRewardTokens(
-        address receiver
-    ) external view returns (address[] memory);
-
-    function rewardConfigurations(
-        address receiver,
-        address token
-    ) external view returns (uint256 emissionRate, uint256 lastUpdate);
+    function setRewardToken(address _rewardToken) external;
 }