compile

src/RewardsDistribution.sol

@@ -1,9 +1,11 @@
 // SPDX-License-Identifier: MIT
-pragma solidity 0.8.25;
+pragma solidity 0.8.20;
 
 import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
-import {Ownable2StepUpgradeable} from "@openzeppelin-upgradeable/contracts/access/Ownable2StepUpgradeable.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";
 
 interface IRewardsDistributor {
     function distributeRewards() external;
@@ -14,15 +16,15 @@ contract RewardsDistributor is Ownable2StepUpgradeable {
     /*//////////////////////////////////////////////////////////////
                                LIBRARIES
     //////////////////////////////////////////////////////////////*/
-    using SafeTransferLib for ERC20;
+    using SafeERC20 for IERC20;
 
     /*//////////////////////////////////////////////////////////////
                                  STRUCTS
     //////////////////////////////////////////////////////////////*/
 
     /// @notice the reward configuration
     struct RewardConfiguration {
-        ERC20 token; // the reward token
+        address token; // the reward token
         uint256 emissionRate; // emission per second
         uint256 lastUpdate; // last update timestamp
     }
@@ -31,7 +33,7 @@ contract RewardsDistributor is Ownable2StepUpgradeable {
     mapping(address receiver => RewardConfiguration[])
         public rewardConfigurations;
 
-    mapping(address receiver => mapping(ERC20 => uint256 id))
+    mapping(address receiver => mapping(address token => uint256 id))
         public rewardConfigurationsIds;
 
     /// @notice Ensure logic contract is unusable
@@ -54,9 +56,12 @@ contract RewardsDistributor is Ownable2StepUpgradeable {
     /// @param emissionRate The emission rate
     function addRewardConfiguration(
         address receiver,
-        ERC20 token,
+        address token,
         uint256 emissionRate
     ) external onlyOwner {
+        require(token != address(0), "No native rewards allowed");
+        require(emissionRate > 0, "Emission rate must be greater than 0");
+
         rewardConfigurations[receiver].push(
             RewardConfiguration(token, emissionRate, block.timestamp)
         );
@@ -73,7 +78,7 @@ contract RewardsDistributor is Ownable2StepUpgradeable {
     /// @dev set the emission rate to 0 to stop the rewards
     function updateEmissonRate(
         address receiver,
-        ERC20 token,
+        address token,
         uint256 emissionRate
     ) external onlyOwner {
         uint256 id = rewardConfigurationsIds[receiver][token];
@@ -88,16 +93,18 @@ contract RewardsDistributor is Ownable2StepUpgradeable {
 
     /// @notice Distribute rewards to receiver
     /// @param token The reward token
-    function distributeReward(ERC20 token) external {
-        uint256 id = rewardConfigurationsIds[msg.sender][token];
+    function distributeReward(address token) external {
+        address receiver = msg.sender;
+
+        uint256 id = rewardConfigurationsIds[receiver][token];
 
         require(
-            rewardConfigurations[msg.sender].length > 0 && id > 0,
+            rewardConfigurations[receiver].length > 0 && id > 0,
             "No reward configuration found"
         );
 
         RewardConfiguration storage rewardConfiguration = rewardConfigurations[
-            msg.sender
+            receiver
         ][id - 1];
 
         // difference in time since last update
@@ -114,6 +121,7 @@ contract RewardsDistributor is Ownable2StepUpgradeable {
         rewardConfiguration.lastUpdate = block.timestamp;
 
         // transfer the reward
-        token.safeTransfer(msg.sender, reward);
+        // TODO change to safeTransfer
+        IERC20(token).safeTransfer(receiver, reward);
     }
 }

src/Staking.sol

@@ -1,15 +1,16 @@
 // SPDX-License-Identifier: MIT
-pragma solidity 0.8.25;
+pragma solidity 0.8.20;
 
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import {Ownable2StepUpgradeable} from "@openzeppelin-upgradeable/contracts/access/Ownable2StepUpgradeable.sol";
 import {ERC20VotesUpgradeable} from "@openzeppelin-upgradeable/contracts/token/ERC20/extensions/ERC20VotesUpgradeable.sol";
 import {SafeTransferLib} from "@solmate/utils/SafeTransferLib.sol";
 import {FixedPointMathLib} from "@solmate/utils/FixedPointMathLib.sol";
 
 interface IRewardsDistributor {
-    function distributeRewards() external;
+    function distributeReward(address token) external;
 }
 
 // TODO should be pausable?
@@ -18,7 +19,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
     /*//////////////////////////////////////////////////////////////
                                LIBRARIES
     //////////////////////////////////////////////////////////////*/
-    using SafeTransferLib for ERC20;
+    using SafeERC20 for IERC20;
     using FixedPointMathLib for uint256;
 
     /*//////////////////////////////////////////////////////////////
@@ -151,7 +152,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         }
 
         // Before doing anything, get the unclaimed rewards first
-        rewardsDistributor.distributeRewards();
+        rewardsDistributor.distributeReward(address(stakingToken));
 
         uint256 sharesToMint = convertToShares(amount);
 
@@ -271,7 +272,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         /////////////////////////// EFFECTS ///////////////////////////////
 
         // Get the unclaimed rewards
-        rewardsDistributor.distributeReward(stakingToken);
+        rewardsDistributor.distributeReward(address(stakingToken));
 
         // Calculates the amounf of shares to burn
         uint256 shares = convertToShares(amount);
@@ -321,14 +322,14 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         require(address(rewardToken) != address(0), "No native token rewards");
 
         // Before doing anything, get the unclaimed rewards first
-        rewardsDistributor.distributeRewards();
+        rewardsDistributor.distributeReward(address(stakingToken));
 
-        address sender = msg.sender;
+        address keyper = msg.sender;
 
         // If the reward token is the staking token, the user is claimingthe staking rewards
         if (rewardToken == stakingToken) {
             // Prevents the keyper from claiming more than they should
-            uint256 maxWithdrawAmount = maxWithdraw(sender);
+            uint256 maxWithdrawAmount = maxWithdraw(keyper);
 
             require(maxWithdrawAmount > 0, "No rewards to claim");
 
@@ -340,12 +341,12 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
 
             // the keyper assets after claiming the rewards must be greater or
             // equal the total locked amount
-            uint256 lockedInShares = convertToShares(totalLocked[sender]);
+            uint256 lockedInShares = convertToShares(totalLocked[keyper]);
 
             // Calculates the amount of shares to burn
             uint256 shares = convertToShares(amount);
 
-            uint256 balance = balanceOf(sender);
+            uint256 balance = balanceOf(keyper);
 
             // If the balance minus the shares is less than the locked in shares
             // the shares will be the balance minus the locked in shares
@@ -357,40 +358,25 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
             _burn(keyper, shares);
 
             // Transfer the SHU to the keyper
-            stakingToken.transfer(sender, amount);
-
-            emit ClaimRewards(sender, address(rewardToken), amount);
-
-            return;
+            rewardToken.safeTransfer(keyper, amount);
         }
 
-        // Calculate the user's total rewards for the specified reward token
-        // TODO check this
-        uint256 totalRewards = (balanceOf(sender) * totalSupply()) /
-            rewardToken.balanceOf(address(this));
-
-        if (amount > totalRewards) {
-            amount = totalRewards;
-        }
-
-        // Transfer the specified amount of rewards to the user
-        rewardToken.transfer(sender, amount);
-
-        emit ClaimRewards(sender, address(rewardToken), amount);
+        emit ClaimRewards(keyper, address(rewardToken), amount);
     }
 
     /// @notice Harvest rewards for a keyper
     /// @param keyper The keyper address
     function harvest(address keyper) external {
-        rewardsDistributor.distributeRewards();
+        rewardsDistributor.distributeReward(address(stakingToken));
 
         uint256 assets = convertToAssets(balanceOf(keyper));
 
         // If the keyper has no assets, there are no rewards to claim
         require(assets > 0, "No rewards to claim");
 
-        // Restake the rewards
-        stake(assets);
+        _burn(keyper, balanceOf(keyper));
+
+        _mint(keyper, convertToShares(assets));
     }
 
     /*//////////////////////////////////////////////////////////////