Delegate fuzz

src/DelegateStaking.sol

@@ -16,150 +16,165 @@
 
 /// @notice Shutter Delegate Staking Contract
 ///         Allows users to stake SHU and earn rewards in exchange.
-contract Delegate is ERC20VotesUpgradeable, OwnableUpgradeable {
+contract DelegateStaking is ERC20VotesUpgradeable, OwnableUpgradeable {
     /*//////////////////////////////////////////////////////////////
                                LIBRARIES
     //////////////////////////////////////////////////////////////*/
     using EnumerableSet for EnumerableSet.UintSet;
 
     using SafeTransferLib for IERC20;
 
     using FixedPointMathLib for uint256;
 
     /*//////////////////////////////////////////////////////////////
                                  VARIABLES
     //////////////////////////////////////////////////////////////*/
 
     /// @notice the staking token, i.e. SHU
     /// @dev set in initialize, can't be changed
     IERC20 public stakingToken;
 
     /// @notice the rewards distributor contract
     /// @dev only owner can change
     IRewardsDistributor public rewardsDistributor;
 
     /// @notice the staking contract
     /// @dev only owner can change
     IStaking public staking;
 
     /// @notice Unique identifier that will be used for the next stake.
     uint256 internal nextStakeId;
 
     /// @notice the lock period in seconds
     /// @dev only owner can change
     uint256 public lockPeriod;
 
     /*//////////////////////////////////////////////////////////////
                                  STRUCTS
     //////////////////////////////////////////////////////////////*/
 
     /// @notice the stake struct
     /// @dev timestamp is the time the stake was made
     struct Stake {
         uint256 keyper;
         uint256 amount;
         uint256 timestamp;
         uint256 lockPeriod;
     }
 
     /*//////////////////////////////////////////////////////////////
                                 MAPPINGS
     //////////////////////////////////////////////////////////////*/
 
     /// @notice stores the metadata associated with a given stake
     mapping(uint256 id => Stake _stake) public stakes;
 
     // @notice stake ids belonging to a user
     mapping(address user => EnumerableSet.UintSet stakeIds) private userStakes;
 
+    /// @notice how many SHU a user has locked
+    mapping(address keyper => uint256 totalLocked) public totalLocked;
+
     /*//////////////////////////////////////////////////////////////
                                  EVENTS
     //////////////////////////////////////////////////////////////*/
 
     /// @notice Emitted when a keyper stakes SHU
-    event Staked(address indexed user, address indexed keyper, uint256 amount);
+    event Staked(
+        address indexed user,
+        address indexed keyper,
+        uint256 amount,
+        uint256 lockPeriod
+    );
 
     /// @notice Emitted when a keyper unstakes SHU
     event Unstaked(address indexed user, uint256 amount, uint256 shares);
 
     /// @notice Emitted when a keyper claims rewards
     event RewardsClaimed(address indexed user, uint256 rewards);
 
     /// @notice Emitted when the rewards distributor is changed
     event NewRewardsDistributor(address indexed rewardsDistributor);
 
     /// @notice Emitted when a new staking contract is set
     event NewStakingContract(address indexed stakingContract);
 
+    /// @notice Emitted when the lock period is changed
+    event NewLockPeriod(uint256 indexed lockPeriod);
+
     /*//////////////////////////////////////////////////////////////
                                  ERRORS
     //////////////////////////////////////////////////////////////*/
 
     /// @notice  Thrown when transfer/tranferFrom is called
     error TransferDisabled();
 
     /// @notice Thrown when a keyper has no shares
     error UserHasNoShares();
 
     /// @notice Trown when amount is zero
     error ZeroAmount();
 
     /// @notice Thrown when someone try to unstake a stake that doesn't belong
     /// to them
     error StakeDoesNotBelongToUser();
 
     /// @notice Thrown when someone try to unstake a stake that doesn't exist
     error StakeDoesNotExist();
 
     /// @notice Thrown when someone try to unstake a amount that is greater than
     /// the stake amount belonging to the stake id
     error WithdrawAmountTooHigh();
 
+    /// @notice Thrown when someone try to unstake a stake that is still locked
+    error StakeIsStillLocked();
+
     /// @notice Thrown when a keyper try to claim rewards but has no rewards to
     /// claim
     error NoRewardsToClaim();
 
     /// @notice Thrown when the argument is the zero address
     error AddressZero();
 
     /*//////////////////////////////////////////////////////////////
                                  MODIFIERS
     //////////////////////////////////////////////////////////////*/
 
     /// @notice Update rewards for a keyper
     modifier updateRewards() {
         // Distribute rewards
         rewardsDistributor.collectRewards();
 
         _;
     }
 
     /// @notice Ensure logic contract is unusable
     constructor() {
         _disableInitializers();
     }
 
     /// @notice Initialize the contract
     /// @param _owner The owner of the contract, i.e. the DAO contract address
     /// @param _stakingToken The address of the staking token, i.e. SHU
     /// @param _rewardsDistributor The address of the rewards distributor
     /// contract
-    /// @param _staking The address of the staking contract
+    /// @param _stakingContract The address of the staking contract
     /// @param _lockPeriod The lock period in seconds
     function initialize(
         address _owner,
         address _stakingToken,
         address _rewardsDistributor,
-        address _staking,
+        address _stakingContract,
         uint256 _lockPeriod
     ) public initializer {
         __ERC20_init("Delegated Staking SHU", "dSHU");
 
         // Transfer ownership to the DAO contract
         _transferOwnership(_owner);
 
-        stakingToken = IERC20(_staking);
+        stakingToken = IERC20(_stakingToken);
         rewardsDistributor = IRewardsDistributor(_rewardsDistributor);
+        staking = IStaking(_stakingContract);
         lockPeriod = _lockPeriod;
 
         nextStakeId = 1;
@@ -393,12 +408,14 @@
     ///            locked amount, the function will return 0
     /// @param user The user address
     /// @return amount The maximum amount of assets that a user can withdraw
-    function maxWithdraw(address user) public view virtual returns (uint256) {
+    function maxWithdraw(
+        address user
+    ) public view virtual returns (uint256 amount) {
         uint256 shares = balanceOf(user);
         require(shares > 0, UserHasNoShares());
 
         uint256 assets = convertToAssets(shares);
-        uint256 locked = totalLocked[keyper] - unlockedAmount;
+        uint256 locked = totalLocked[user];
 
         // need the first branch as convertToAssets rounds down
         amount = locked >= assets ? 0 : assets - locked;

test/DelegateStaking.t.sol

@@ -0,0 +1,184 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.26;
+
+import "@forge-std/Test.sol";
+import {TransparentUpgradeableProxy, ITransparentUpgradeableProxy} from "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
+import {IERC20Metadata} from "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";
+
+import {Staking} from "src/Staking.sol";
+import {DelegateStaking} from "src/DelegateStaking.sol";
+import {RewardsDistributor} from "src/RewardsDistributor.sol";
+import {IRewardsDistributor} from "src/interfaces/IRewardsDistributor.sol";
+import {MockGovToken} from "test/mocks/MockGovToken.sol";
+import {ProxyUtils} from "test/helpers/ProxyUtils.sol";
+import {DelegateStakingHarness} from "test/helpers/DelegateStakingHarness.sol";
+
+contract DelegateStakingTest is Test {
+    DelegateStakingHarness public delegate;
+    IRewardsDistributor public rewardsDistributor;
+    Staking public staking;
+    MockGovToken public govToken;
+
+    uint256 constant LOCK_PERIOD = 182 days; // 6 months
+    uint256 constant REWARD_RATE = 0.1e18;
+
+    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.
+        _jumpAhead(1234);
+
+        govToken = new MockGovToken();
+        _mintGovToken(address(this), 100_000_000e18);
+        vm.label(address(govToken), "govToken");
+
+        // deploy rewards distributor
+        rewardsDistributor = IRewardsDistributor(
+            new RewardsDistributor(address(this), address(govToken))
+        );
+
+        // deploy staking
+        address stakingImpl = address(new Staking());
+
+        staking = Staking(
+            address(
+                new TransparentUpgradeableProxy(stakingImpl, address(this), "")
+            )
+        );
+        vm.label(address(staking), "staking");
+
+        staking.initialize(
+            address(this), // owner
+            address(govToken),
+            address(rewardsDistributor),
+            0,
+            0
+        );
+
+        address delegateImpl = address(new DelegateStakingHarness());
+
+        delegate = DelegateStakingHarness(
+            address(
+                new TransparentUpgradeableProxy(delegateImpl, address(this), "")
+            )
+        );
+        vm.label(address(delegate), "staking");
+
+        delegate.initialize(
+            address(this), // owner
+            address(govToken),
+            address(rewardsDistributor),
+            address(staking),
+            LOCK_PERIOD
+        );
+
+        rewardsDistributor.setRewardConfiguration(
+            address(staking),
+            REWARD_RATE
+        );
+
+        // fund reward distribution
+        govToken.transfer(address(rewardsDistributor), 100_000_000e18);
+    }
+
+    function _jumpAhead(uint256 _seconds) public {
+        vm.warp(vm.getBlockTimestamp() + _seconds);
+    }
+
+    function _boundMintAmount(uint96 _amount) internal pure returns (uint256) {
+        return bound(_amount, 0, 10_000_000e18);
+    }
+
+    function _boundRealisticTimeAhead(
+        uint256 _time
+    ) internal pure returns (uint256) {
+        return bound(_time, 1, 105 weeks); // two years
+    }
+
+    function _boundUnlockedTime(uint256 _time) internal view returns (uint256) {
+        return bound(_time, vm.getBlockTimestamp() + LOCK_PERIOD, 105 weeks);
+    }
+
+    function _mintGovToken(address _to, uint256 _amount) internal {
+        vm.assume(
+            _to != address(0) &&
+                _to != address(delegate) &&
+                _to != ProxyUtils.getAdminAddress(address(delegate))
+        );
+
+        govToken.mint(_to, _amount);
+    }
+
+    function _boundToRealisticStake(
+        uint256 _stakeAmount
+    ) public pure returns (uint256 _boundedStakeAmount) {
+        _boundedStakeAmount = uint256(
+            bound(_stakeAmount, 100e18, 5_000_000e18)
+        );
+    }
+
+    function _stake(
+        address _user,
+        address _keyper,
+        uint256 _amount
+    ) internal returns (uint256 stakeId) {
+        vm.assume(
+            _keyper != address(0) &&
+                _keyper != ProxyUtils.getAdminAddress(address(staking))
+        );
+
+        vm.assume(
+            _user != address(0) &&
+                _user != address(this) &&
+                _user != address(delegate) &&
+                _user != ProxyUtils.getAdminAddress(address(delegate)) &&
+                _user != address(rewardsDistributor)
+        );
+
+        vm.startPrank(_user);
+        govToken.approve(address(delegate), _amount);
+        stakeId = delegate.stake(_keyper, _amount);
+        vm.stopPrank();
+    }
+}
+
+contract Initializer is DelegateStakingTest {
+    function test_Initialize() public view {
+        assertEq(
+            IERC20Metadata(address(delegate)).name(),
+            "Delegated Staking SHU"
+        );
+        assertEq(IERC20Metadata(address(delegate)).symbol(), "dSHU");
+        assertEq(delegate.owner(), address(this), "Wrong owner");
+        assertEq(
+            address(delegate.stakingToken()),
+            address(govToken),
+            "Wrong staking token"
+        );
+        assertEq(
+            address(delegate.rewardsDistributor()),
+            address(rewardsDistributor),
+            "Wrong rewards distributor"
+        );
+        assertEq(delegate.lockPeriod(), LOCK_PERIOD, "Wrong lock period");
+        assertEq(
+            address(delegate.staking()),
+            address(staking),
+            "Wrong staking"
+        );
+
+        assertEq(delegate.exposed_nextStakeId(), 1);
+    }
+
+    function test_RevertIf_InitializeImplementation() public {
+        DelegateStaking delegateImpl = new DelegateStaking();
+
+        vm.expectRevert();
+        delegateImpl.initialize(
+            address(this),
+            address(govToken),
+            address(rewardsDistributor),
+            address(staking),
+            LOCK_PERIOD
+        );
+    }
+}

test/Staking.t.sol

@@ -155,18 +155,6 @@ contract StakingTest is Test {
 
         return _amount.mulDivDown(supply + 1, assets + 1);
     }
-
-    function _assertMinRelativeLoss(
-        uint256 spent,
-        uint256 received,
-        uint256 minRelLoss,
-        string memory errorMessage
-    ) internal pure {
-        assertGt(spent, received, "Spent should be greater than received");
-
-        uint256 relativeLoss = ((spent - received) * 1e18) / spent;
-        assertGe(relativeLoss, minRelLoss, errorMessage);
-    }
 }
 
 contract Initializer is StakingTest {

test/helpers/DelegateStakingHarness.sol

@@ -0,0 +1,10 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.26;
+
+import {DelegateStaking} from "src/DelegateStaking.sol";
+
+contract DelegateStakingHarness is DelegateStaking {
+    function exposed_nextStakeId() external view returns (uint256) {
+        return nextStakeId;
+    }
+}