Distribute rewards tests

.solhintrc

@@ -7,6 +7,7 @@
     	"no-inline-assembly": "off",
     	"no-empty-blocks": "off",
     	"no-global-import": "off",
-        gas-custom-errors": "off",
+        "gas-custom-errors": "off",
+        "func-name-mixedcase": "off
     }
 }

src/RewardsDistributor.sol

@@ -1,157 +1,135 @@
 // 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 {Ownable2Step} from "@openzeppelin/contracts/access/Ownable2Step.sol";
+import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+import {IRewardsDistributor} from "./interfaces/IRewardsDistributor.sol";
 
-interface IRewardsDistributor {
-    function distributeRewards() external;
-}
-
-// TODO should be pausable?
-contract RewardsDistributor is Ownable2StepUpgradeable {
+contract RewardsDistributor is Ownable2Step, IRewardsDistributor {
     /*//////////////////////////////////////////////////////////////
-                               LIBRARIES
+                                 LIBRARIES
     //////////////////////////////////////////////////////////////*/
     using SafeERC20 for IERC20;
 
+    /*//////////////////////////////////////////////////////////////
+                                 VARIABLES
+    //////////////////////////////////////////////////////////////*/
+
+    /// @notice the reward token, i.e. SHU
+    IERC20 public rewardToken;
+
     /*//////////////////////////////////////////////////////////////
                                  STRUCTS
     //////////////////////////////////////////////////////////////*/
 
     /// @notice the reward configuration
     struct RewardConfiguration {
         uint256 emissionRate; // emission per second
         uint256 lastUpdate; // last update timestamp
     }
 
     /*//////////////////////////////////////////////////////////////
-                             MAPPINGS/ARRAYS
+                                MAPPINGS
     //////////////////////////////////////////////////////////////*/
 
-    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
-    );
-
-    /// @notice Ensure logic contract is unusable
-    constructor() {
-        _disableInitializers();
-    }
+    event RewardCollected(address indexed receiver, uint256 reward);
 
-    /// @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();
-
-        // Transfer ownership to the DAO contract
-        _transferOwnership(newOwner);
-    }
+    event RewardTokenSet(address indexed rewardToken);
 
-    /// @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);
-        }
+    /*//////////////////////////////////////////////////////////////
+                                 ERRORS
+    //////////////////////////////////////////////////////////////*/
 
-        rewardConfigurations[receiver][token] = RewardConfiguration(
-            emissionRate,
-            block.timestamp
-        );
+    /// @notice Thrown when address is zero
+    error ZeroAddress();
 
-        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);
+    /// @notice Initialize the contract
+    /// @param newOwner The owner of the contract, i.e. the DAO contract address
+    /// @param _rewardToken The reward token, i.e. SHU
+    constructor(address newOwner, address _rewardToken) Ownable(newOwner) {
+        // 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

@@ -40,7 +40,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
     uint256 public minStake;
 
     /// @notice Unique identifier that will be used for the next stake.
-    uint256 private nextStakeId;
+    uint256 internal nextStakeId;
 
     /*//////////////////////////////////////////////////////////////
                                  STRUCTS
@@ -119,7 +119,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
     /// 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
@@ -153,7 +153,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
     /// @notice Update rewards for a keyper
     modifier updateRewards() {
         // Distribute rewards
-        rewardsDistributor.distributeReward(address(stakingToken));
+        rewardsDistributor.collectRewards();
 
         _;
     }
@@ -186,6 +186,8 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         rewardsDistributor = IRewardsDistributor(_rewardsDistributor);
         lockPeriod = _lockPeriod;
         minStake = _minStake;
+
+        nextStakeId = 1;
     }
 
     /// @notice Stake SHU
@@ -195,11 +197,11 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
     ///          - The shares are non-transferable
     ///          - Only keypers can stake
     /// @param amount The amount of SHU to stake
-    /// @return The index of the stake
+    /// @return stakeId The index of the stake
     /// TODO slippage protection
     function stake(
         uint256 amount
-    ) external onlyKeyper updateRewards returns (uint256) {
+    ) external onlyKeyper updateRewards returns (uint256 stakeId) {
         /////////////////////////// CHECKS ///////////////////////////////
         require(amount > 0, ZeroAmount());
 
@@ -224,7 +226,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         _mint(keyper, sharesToMint);
 
         // Get next stake id and increment it
-        uint256 stakeId = ++nextStakeId;
+        stakeId = nextStakeId++;
 
         stakes[stakeId] = Stake({
             amount: amount,
@@ -240,8 +242,6 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         stakingToken.safeTransferFrom(keyper, address(this), amount);
 
         emit Staked(keyper, amount, sharesToMint, lockPeriod);
-
-        return stakeId;
     }
 
     /// @notice Unstake SHU

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;
 }