review with gui

script/Deploy.s.sol

@@ -4,13 +4,18 @@ pragma solidity 0.8.26;
 import "@forge-std/Script.sol";
 import {TransparentUpgradeableProxy} from "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
 import {RewardsDistributor} from "src/RewardsDistributor.sol";
+import {DelegateStaking} from "src/DelegateStaking.sol";
 import {Staking} from "src/Staking.sol";
 import "./Constants.sol";
 
 contract Deploy is Script {
     function run()
         public
-        returns (Staking stakingProxy, RewardsDistributor rewardsDistributor)
+        returns (
+            Staking stakingProxy,
+            RewardsDistributor rewardsDistributor,
+            DelegateStaking delegateProxy
+        )
     {
         vm.startBroadcast();
 
@@ -37,6 +42,25 @@ contract Deploy is Script {
             MIN_STAKE
         );
 
+        DelegateStaking delegate = new DelegateStaking();
+        delegateProxy = DelegateStaking(
+            address(
+                new TransparentUpgradeableProxy(
+                    address(delegate),
+                    address(CONTRACT_OWNER),
+                    ""
+                )
+            )
+        );
+
+        delegateProxy.initialize(
+            CONTRACT_OWNER,
+            STAKING_TOKEN,
+            address(rewardsDistributor),
+            address(stakingProxy),
+            LOCK_PERIOD
+        );
+
         vm.stopBroadcast();
     }
 }

script/testnet/DeployTestnet.s.sol

@@ -5,14 +5,19 @@ import "@forge-std/Script.sol";
 import {TransparentUpgradeableProxy} from "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
 import {RewardsDistributor} from "src/RewardsDistributor.sol";
 import {Staking} from "src/Staking.sol";
+import {DelegateStaking} from "src/DelegateStaking.sol";
 import {MockGovToken} from "test/mocks/MockGovToken.sol";
 import "./Constants.sol";
 
 // forge script script/testnet/DeployTestnet.s.sol --rpc-url testnet -vvvvv --slow --always-use-create-2-factory --account test --etherscan-api-key testnet --verify --chain 11155111
 contract DeployTestnet is Script {
     function run()
         public
-        returns (Staking stakingProxy, RewardsDistributor rewardsDistributor)
+        returns (
+            Staking stakingProxy,
+            RewardsDistributor rewardsDistributor,
+            DelegateStaking delegateProxy
+        )
     {
         vm.startBroadcast();
 
@@ -42,6 +47,25 @@ contract DeployTestnet is Script {
             MIN_STAKE
         );
 
+        DelegateStaking delegate = new DelegateStaking();
+        delegateProxy = DelegateStaking(
+            address(
+                new TransparentUpgradeableProxy(
+                    address(delegate),
+                    address(CONTRACT_OWNER),
+                    ""
+                )
+            )
+        );
+
+        delegateProxy.initialize(
+            CONTRACT_OWNER,
+            MOCKED_SHU,
+            address(rewardsDistributor),
+            address(stakingProxy),
+            LOCK_PERIOD
+        );
+
         vm.stopBroadcast();
     }
 }

src/BaseStaking.sol

@@ -180,19 +180,15 @@ abstract contract BaseStaking is OwnableUpgradeable, ERC20VotesUpgradeable {
     function convertToShares(
         uint256 assets
     ) public view virtual returns (uint256) {
-        // sum + 1 on both sides to prevent donation attack
-        // this is the same as OZ ERC4626 prevetion to inflation attack with decimal offset = 0
-        return assets.mulDivDown(totalSupply() + 1, _totalAssets() + 1);
+        return assets.mulDivDown(totalSupply(), _totalAssets());
     }
 
     /// @notice Get the total amount of assets the shares are worth
     /// @param shares The amount of shares
     function convertToAssets(
         uint256 shares
     ) public view virtual returns (uint256) {
-        // sum + 1 on both sides to prevent donation attack
-        // this is the same as OZ ERC4626 prevetion to inflation attack with decimal offset = 0
-        return shares.mulDivDown(_totalAssets() + 1, totalSupply() + 1);
+        return shares.mulDivDown(_totalAssets(), totalSupply());
     }
 
     /// @notice Get the stake ids belonging to a user
@@ -213,24 +209,17 @@ abstract contract BaseStaking is OwnableUpgradeable, ERC20VotesUpgradeable {
     /// @notice Deposit SHU into the contract
     /// @param user The user address
     /// @param amount The amount of SHU to deposit
-    function _deposit(address user, uint256 amount) internal {
+    function _deposit(uint256 amount) internal {
         // Calculate the amount of shares to mint
         uint256 shares = convertToShares(amount);
 
-        // A first deposit donation attack may result in shares being 0 if the
-        // contract has very high assets balance but a very low total supply.
-        // Although this attack is not profitable for the attacker, as they will
-        // spend more tokens than they will receive, it can still be used to perform a DDOS attack
-        // against a specific user. The targeted user can still withdraw their SHU,
-        // but this is only guaranteed if someone mints to increase the total supply of shares,
-        // because previewWithdraw rounds up and their shares will be less than the burn amount.
-        require(shares > 0, SharesMustBeGreaterThanZero());
-
         // Update the total locked amount
-        totalLocked[user] += amount;
+        unchecked {
+            totalLocked[msg.sender] += amount;
+        }
 
         // Mint the shares
-        _mint(user, shares);
+        _mint(msg.sender, shares);
 
         // Lock the SHU in the contract
         stakingToken.safeTransferFrom(msg.sender, address(this), amount);
@@ -263,9 +252,7 @@ abstract contract BaseStaking is OwnableUpgradeable, ERC20VotesUpgradeable {
     /// @notice Get the amount of shares that will be burned
     /// @param assets The amount of assets
     function _previewWithdraw(uint256 assets) internal view returns (uint256) {
-        // sum + 1 on both sides to prevent donation attack
-        // this is the same as OZ ERC4626 prevetion to inflation attack with decimal offset = 0
-        return assets.mulDivUp(totalSupply() + 1, _totalAssets() + 1);
+        return assets.mulDivUp(totalSupply(), _totalAssets());
     }
 
     /// @notice Calculates the amount to withdraw

src/DelegateStaking.sol

@@ -172,9 +172,11 @@ contract DelegateStaking is BaseStaking {
         stakes[stakeId].lockPeriod = lockPeriod;
 
         // Increase the keyper total delegated amount
-        totalDelegated[keyper] += amount;
+        unchecked {
+            totalDelegated[keyper] += amount;
+        }
 
-        _deposit(user, amount);
+        _deposit(amount);
 
         emit Staked(user, keyper, amount, lockPeriod);
     }

src/RewardsDistributor.sol

@@ -75,7 +75,7 @@ contract RewardsDistributor is Ownable, IRewardsDistributor {
 
     /// @notice Distribute rewards to receiver
     /// Caller must be the receiver
-    function collectRewards() external override returns (uint256 rewards) {
+    function collectRewards() public override returns (uint256 rewards) {
         address receiver = msg.sender;
 
         RewardConfiguration storage rewardConfiguration = rewardConfigurations[
@@ -85,13 +85,11 @@ contract RewardsDistributor is Ownable, IRewardsDistributor {
         // difference in time since last update
         uint256 timeDelta = block.timestamp - rewardConfiguration.lastUpdate;
 
-        uint256 funds = rewardToken.balanceOf(address(this));
-
         rewards = rewardConfiguration.emissionRate * timeDelta;
 
         // the contract must have enough funds to distribute
         // we don't want to revert in case its zero to not block the staking contract
-        if (rewards == 0 || funds < rewards) {
+        if (rewards == 0 || rewardToken.balanceOf(address(this)) < rewards) {
             return 0;
         }
 
@@ -108,7 +106,7 @@ contract RewardsDistributor is Ownable, IRewardsDistributor {
     /// @param receiver The receiver of the rewards
     function collectRewardsTo(
         address receiver
-    ) external override returns (uint256 rewards) {
+    ) public override returns (uint256 rewards) {
         RewardConfiguration storage rewardConfiguration = rewardConfigurations[
             receiver
         ];
@@ -120,12 +118,14 @@ contract RewardsDistributor is Ownable, IRewardsDistributor {
 
         require(timeDelta > 0, TimeDeltaZero());
 
-        uint256 funds = rewardToken.balanceOf(address(this));
-
         rewards = rewardConfiguration.emissionRate * timeDelta;
 
         // the contract must have enough funds to distribute
-        require(funds >= rewards, NotEnoughFunds());
+        // and the rewards must be greater than zero
+        require(
+            rewards > 0 && rewardToken.balanceOf(address(this)) >= rewards,
+            NotEnoughFunds()
+        );
 
         // update the last update timestamp
         rewardConfiguration.lastUpdate = block.timestamp;
@@ -142,7 +142,7 @@ contract RewardsDistributor is Ownable, IRewardsDistributor {
     function setRewardConfiguration(
         address receiver,
         uint256 emissionRate
-    ) external override onlyOwner {
+    ) public override onlyOwner {
         require(receiver != address(0), ZeroAddress());
 
         // to remove a rewards, it should call removeRewardConfiguration
@@ -151,41 +151,52 @@ contract RewardsDistributor is Ownable, IRewardsDistributor {
         // only update last update if it's the first time
         if (rewardConfigurations[receiver].lastUpdate == 0) {
             rewardConfigurations[receiver].lastUpdate = block.timestamp;
+        } else {
+            // claim the rewards before updating the emission rate
+            collectRewardsTo(receiver);
         }
+
         rewardConfigurations[receiver].emissionRate = emissionRate;
 
         emit RewardConfigurationSet(receiver, emissionRate);
     }
 
     /// @notice Remove a reward configuration
     /// @param receiver The receiver of the rewards
-    function removeRewardConfiguration(address receiver) external onlyOwner {
-        delete rewardConfigurations[receiver];
+    function removeRewardConfiguration(address receiver) public onlyOwner {
+        rewardConfigurations[receiver].lastUpdate = 0;
+        rewardConfigurations[receiver].emissionRate = 0;
 
         emit RewardConfigurationSet(receiver, 0);
     }
 
-    /// @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);
-    }
-
     /// @notice Set the reward token
     /// @param _rewardToken The reward token
-    function setRewardToken(address _rewardToken) external onlyOwner {
+    function setRewardToken(address _rewardToken) public onlyOwner {
         require(_rewardToken != address(0), ZeroAddress());
 
         // withdraw remaining old reward token
-        withdrawFunds(msg.sender, rewardToken.balanceOf(address(this)));
+        withdrawFunds(
+            address(rewardToken),
+            msg.sender,
+            rewardToken.balanceOf(address(this))
+        );
 
         // set the new reward token
         rewardToken = IERC20(_rewardToken);
 
         emit RewardTokenSet(_rewardToken);
     }
+
+    /// @notice Withdraw funds from the contract
+    /// @param to The address to withdraw to
+    /// @param amount The amount to withdraw
+    function withdrawFunds(
+        address token,
+        address to,
+        uint256 amount
+    ) public onlyOwner {
+        require(to != address(0), ZeroAddress());
+        IERC20(token).safeTransfer(to, amount);
+    }
 }