feat: move cooldown to shares uint256 (#20)

src/contracts/IStakeToken.sol

@@ -7,6 +7,14 @@ interface IStakeToken {
     uint216 amount;
   }
 
+  struct CooldownConfig {
+    /// @notice Seconds available to redeem once the cooldown period is fulfilled
+    uint32 unstakeWindowSeconds;
+    /// @notice Seconds between starting cooldown and being able to withdraw
+    uint32 cooldownSeconds;
+    // reserved for future use
+  }
+
   event Cooldown(address indexed user, uint256 amount);
 
   event Staked(address indexed from, address indexed to, uint256 assets, uint256 shares);
@@ -15,6 +23,7 @@ interface IStakeToken {
   event Slashed(address indexed destination, uint256 amount);
   event SlashingExitWindowDurationChanged(uint256 windowSeconds);
   event CooldownSecondsChanged(uint256 cooldownSeconds);
+  event UnstakeWindowChanged(uint256 unstakeWindow);
   event ExchangeRateChanged(uint216 exchangeRate);
   event FundsReturned(uint256 amount);
   event SlashingSettled();

src/contracts/StakeToken.sol

@@ -26,16 +26,12 @@ contract StakeToken is ERC20Permit, IStakeToken, Rescuable {
 
   IERC20 public immutable STAKED_TOKEN;
 
-  /// @notice Seconds available to redeem once the cooldown period is fulfilled
-  uint256 public immutable UNSTAKE_WINDOW;
-
   IRewardsController public immutable REWARDS_CONTROLLER;
 
   mapping(address => uint256) public stakerRewardsToClaim;
   mapping(address => CooldownSnapshot) public stakersCooldowns;
 
-  /// @notice Seconds between starting cooldown and being able to withdraw
-  uint256 internal _cooldownSeconds;
+  CooldownConfig internal _cooldownConfig;
   /// @notice Mirror of latest snapshot value for cheaper access
   uint216 internal _currentExchangeRate;
 
@@ -53,25 +49,25 @@ contract StakeToken is ERC20Permit, IStakeToken, Rescuable {
   constructor(
     string memory name,
     IERC20 stakedToken,
-    uint256 unstakeWindow,
     IRewardsController rewardsController
   ) ERC20Permit(name) {
     uint256 decimals = IERC20Metadata(address(stakedToken)).decimals();
     STAKED_TOKEN = stakedToken;
-    UNSTAKE_WINDOW = unstakeWindow;
     REWARDS_CONTROLLER = rewardsController;
   }
 
   function initialize(
     string calldata name,
     string calldata symbol,
     address newSlashingAdmin,
-    uint256 cooldownSeconds
+    uint256 cooldownSeconds,
+    uint256 unstakeWindow
   ) external virtual initializer {
     _initializeMetadata(name, symbol);
     _transferOwnership(newSlashingAdmin);
     _setSlashingAdmin(newSlashingAdmin);
     _setCooldownSeconds(cooldownSeconds);
+    _setUnstakeWindow(unstakeWindow);
     _updateExchangeRate(INITIAL_EXCHANGE_RATE);
     minAssetsRemaining = 10 ** decimals();
   }
@@ -86,6 +82,19 @@ contract StakeToken is ERC20Permit, IStakeToken, Rescuable {
     return owner();
   }
 
+  function setUnstakeWindow(uint256 newUnstakeWindow) external onlyOwner {
+    _setUnstakeWindow(newUnstakeWindow);
+  }
+
+  function _setUnstakeWindow(uint256 newUnstakeWindow) internal {
+    _cooldownConfig.unstakeWindowSeconds = newUnstakeWindow.toUint32();
+    emit UnstakeWindowChanged(newUnstakeWindow);
+  }
+
+  function getUnstakeWindow() external view returns (uint256) {
+    return _cooldownConfig.unstakeWindowSeconds;
+  }
+
   function setSlashingAdmin(address newSlashingAdmin) external onlyOwner {
     _setSlashingAdmin(newSlashingAdmin);
   }
@@ -199,7 +208,7 @@ contract StakeToken is ERC20Permit, IStakeToken, Rescuable {
 
   /// @inheritdoc IStakeToken
   function getCooldownSeconds() external view returns (uint256) {
-    return _cooldownSeconds;
+    return _cooldownConfig.cooldownSeconds;
   }
 
   function _cooldown(address from) internal {
@@ -218,7 +227,7 @@ contract StakeToken is ERC20Permit, IStakeToken, Rescuable {
    * @param cooldownSeconds the new amount of cooldown seconds
    */
   function _setCooldownSeconds(uint256 cooldownSeconds) internal {
-    _cooldownSeconds = cooldownSeconds;
+    _cooldownConfig.cooldownSeconds = cooldownSeconds.toUint32();
     emit CooldownSecondsChanged(cooldownSeconds);
   }
 
@@ -250,12 +259,14 @@ contract StakeToken is ERC20Permit, IStakeToken, Rescuable {
     require(amount != 0, 'INVALID_ZERO_AMOUNT');
 
     CooldownSnapshot memory cooldownSnapshot = stakersCooldowns[from];
+    CooldownConfig memory cachedCooldownConfig = _cooldownConfig;
     require(
-      (block.timestamp >= cooldownSnapshot.timestamp + _cooldownSeconds),
+      (block.timestamp >= cooldownSnapshot.timestamp + cachedCooldownConfig.cooldownSeconds),
       'INSUFFICIENT_COOLDOWN'
     );
     require(
-      (block.timestamp - (cooldownSnapshot.timestamp + _cooldownSeconds) <= UNSTAKE_WINDOW),
+      (block.timestamp - (cooldownSnapshot.timestamp + cachedCooldownConfig.cooldownSeconds) <=
+        cachedCooldownConfig.unstakeWindowSeconds),
       'UNSTAKE_WINDOW_FINISHED'
     );
 

tests/Cooldown.t.sol

@@ -139,7 +139,8 @@ contract Cooldown is StkTestUtils {
   ) public {
     vm.assume(amountToUnstake != 0 && amountToStake >= amountToUnstake);
     vm.assume(
-      secondsAfterCooldownActivation > stakeToken.getCooldownSeconds() + stakeToken.UNSTAKE_WINDOW()
+      secondsAfterCooldownActivation >
+        stakeToken.getCooldownSeconds() + stakeToken.getUnstakeWindow()
     );
     vm.assume(user != address(proxyAdmin) && user != address(0) && destination != address(0));
 

tests/StkTestUtils.t.sol

@@ -33,7 +33,6 @@ contract StkTestUtils is Test {
     stakeTokenImpl = new StakeToken(
       'stkTest',
       underlyingToken,
-      2 days,
       IRewardsController(address(controller))
     );
     proxyAdmin = new ProxyAdmin(admin);
@@ -47,7 +46,8 @@ contract StkTestUtils is Test {
             'Stake Test',
             'stkTest',
             admin,
-            15 days
+            15 days,
+            2 days
           )
         )
       )