From 31b41915e2af3436708c47007912928092e17937 Mon Sep 17 00:00:00 2001
From: Pavel Menshikov <menshikov.pv@mail.ru>
Date: Sat, 13 Jul 2024 23:38:29 +0500
Subject: [PATCH 1/5] first version of fast-withdrawal feature is ready

---
 src/contracts/IStakeToken.sol |  56 +++++++++--
 src/contracts/StakeToken.sol  | 172 ++++++++++++++++++++++++++++------
 tests/Cooldown.t.sol          | 166 ++++++++++++++++++++++++++------
 tests/ERC20Std.t.sol          |   2 +-
 tests/StkTestUtils.t.sol      |  14 ++-
 5 files changed, 343 insertions(+), 67 deletions(-)

diff --git a/src/contracts/IStakeToken.sol b/src/contracts/IStakeToken.sol
index 4378e59..ed60ce9 100644
--- a/src/contracts/IStakeToken.sol
+++ b/src/contracts/IStakeToken.sol
@@ -2,8 +2,11 @@
 pragma solidity ^0.8.0;
 
 interface IStakeToken {
-  struct CooldownSnapshot {
-    uint40 timestamp;
+  struct CooldownSetup {
+    // make more sense to display time from which can be withdrawn, not activation time
+    /// @notice The time after which funds can be redeemed
+    uint32 timestamp;
+    /// @notice The amount of tokens which can be redeemed
     uint216 amount;
   }
 
@@ -11,13 +14,22 @@ interface IStakeToken {
     /// @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;
+    uint32 defaultCooldownSeconds;
     /// @notice The address of the underlying asset
     address stakedToken;
-    // reserved for future use
+    /// @notice The maximum time available for reduction cooldown period
+    uint32 maxReductionSeconds;
+    /// @notice The maximum fee in BIPS that will be taken when the cooldown period is reduced by maxReductionTime
+    uint216 maxFee;
+    /// @notice The address of treasury
+    address treasury;
   }
 
-  event Cooldown(address indexed user, uint256 amount);
+  event Cooldown(address indexed user, uint256 amount, uint32 timeToRedeem);
+  event FeesSentToTreasury(uint256 amount);
+  event TreasuryChanged(address treasury);
+  event MaxFeeChanged(uint256 maxFee);
+  event MaxReductionSecondsChanged(uint256 maxReductionSeconds);
 
   event Staked(address indexed from, address indexed to, uint256 assets, uint256 shares);
   event Redeem(address indexed from, address indexed to, uint256 assets, uint256 shares);
@@ -51,6 +63,12 @@ interface IStakeToken {
    */
   function cooldown() external;
 
+  /**
+   * @dev Activates the cooldown period and reduces it
+   * @param reducedTime Time by which the cooldown will be reduced
+   */
+  function reducedCooldown(uint256 reducedTime) external;
+
   /**
    * @dev Allows staking a certain amount of STAKED_TOKEN with gasless approvals (permit)
    * @param amount The amount to be staked
@@ -88,9 +106,15 @@ interface IStakeToken {
 
   /**
    * @dev Getter of the cooldown seconds
-   * @return cooldownSeconds the amount of seconds between starting the cooldown and being able to redeem
+   * @return defaultCooldownSeconds The amount of seconds between starting the cooldown and being able to redeem by default
+   */
+  function getDefaultCooldownSeconds() external view returns (uint256);
+
+  /**
+   * @dev Getter of the maximum reduction cooldown seconds
+   * @return maxReductionSeconds The maximum time available for reduction cooldown period
    */
-  function getCooldownSeconds() external view returns (uint256);
+  function getMaxReductionSeconds() external view returns (uint256);
 
   /**
    * @dev Setter of cooldown seconds
@@ -99,6 +123,24 @@ interface IStakeToken {
    */
   function setCooldownSeconds(uint256 cooldownSeconds) external;
 
+  /**
+   * @dev Setter of treasury address
+   * @param treasury new treasury address to set for collecting fast-withdrawal fees
+   */
+  function setTreasury(address treasury) external;
+
+  /**
+   * @dev Setter of max fee
+   * @param maxFee Amount of fees in BPS, which should be paid for max cooldown reduction
+   */
+  function setMaxFee(uint256 maxFee) external;
+
+  /**
+   * @dev Setter of max cooldown reduction time in seconds
+   * @param newMaxReductionTime number of seconds by which the cooldown can be reduced with the payment of fees
+   */
+  function setMaxReductionSeconds(uint256 newMaxReductionTime) external;
+
   /**
    * @dev returns the exact amount of shares that would be received for the provided number of assets
    * @param assets the number of assets to stake
diff --git a/src/contracts/StakeToken.sol b/src/contracts/StakeToken.sol
index a1b71d9..eb6b5e6 100644
--- a/src/contracts/StakeToken.sol
+++ b/src/contracts/StakeToken.sol
@@ -23,12 +23,13 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
 
   uint216 public constant INITIAL_EXCHANGE_RATE = 1e18;
   uint256 public constant EXCHANGE_RATE_UNIT = 1e18;
+  uint216 public constant HUNDRED_PERCENT = 10_000;
 
   IRewardsController public immutable REWARDS_CONTROLLER;
 
   /// @custom:storage-location erc7201:aave.storage.StakeToken
   struct StakeTokenStorage {
-    mapping(address => CooldownSnapshot) _stakersCooldowns;
+    mapping(address => CooldownSetup) _stakersCooldowns;
     SmConfig _smConfig;
     /// @notice Current exchangeRate of the stk
     uint216 _currentExchangeRate;
@@ -54,7 +55,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
     }
   }
 
-  function stakersCooldowns(address user) public view returns (CooldownSnapshot memory) {
+  function stakersCooldowns(address user) public view returns (CooldownSetup memory) {
     StakeTokenStorage storage $ = _getStakeTokenStorage();
     return $._stakersCooldowns[user];
   }
@@ -70,18 +71,29 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
     string calldata symbol,
     address newSlashingAdmin,
     uint256 cooldownSeconds,
-    uint256 unstakeWindow
+    uint256 unstakeWindow,
+    address treasury,
+    uint256 maxFee,
+    uint256 maxReductionSeconds
   ) external virtual initializer {
-    StakeTokenStorage storage $ = _getStakeTokenStorage();
-    $._smConfig.stakedToken = stakedToken;
+    // @pavelvm5 made like this, due to stack-too-deep error, will fix in future, if any args will be added could be useful to rewrite to assembly
+    _getStakeTokenStorage()._smConfig.stakedToken = stakedToken;
+    _getStakeTokenStorage()._minAssetsRemaining = 10 ** decimals();
+
     __ERC20_init(name, symbol); // TODO: should naming be inherited from underlying or not?
     __Ownable_init(newSlashingAdmin);
     __EIP712_init(string(abi.encodePacked('stk', name)), '1');
+
     _setSlashingAdmin(newSlashingAdmin);
+    _setTreasury(treasury);
+
+    _setMaxFee(maxFee);
+
     _setCooldownSeconds(cooldownSeconds);
     _setUnstakeWindow(unstakeWindow);
+    _setMaxReductionSeconds(maxReductionSeconds);
+
     _updateExchangeRate(INITIAL_EXCHANGE_RATE);
-    $._minAssetsRemaining = 10 ** decimals();
   }
 
   function decimals() public view override returns (uint8) {
@@ -104,8 +116,11 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
   }
 
   function _setUnstakeWindow(uint256 newUnstakeWindow) internal {
+    require(newUnstakeWindow >= 1 hours, 'TOO_LOW_UNSTAKE_WINDOW');
+
     StakeTokenStorage storage $ = _getStakeTokenStorage();
     $._smConfig.unstakeWindowSeconds = newUnstakeWindow.toUint32();
+
     emit UnstakeWindowChanged(newUnstakeWindow);
   }
 
@@ -169,9 +184,15 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
 
   /// @inheritdoc IStakeToken
   function cooldownOnBehalfOf(address from) external onlyOwner {
+    // @audit-info same modif here
     _cooldown(from);
   }
 
+  /// @inheritdoc IStakeToken
+  function reducedCooldown(uint256 reductionTime) external {
+    _reducedCooldown(msg.sender, reductionTime.toUint32());
+  }
+
   /// @inheritdoc IStakeToken
   function redeem(address to, uint256 amount) external {
     _redeem(msg.sender, to, amount.toUint104());
@@ -179,6 +200,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
 
   /// @inheritdoc IStakeToken
   function redeemOnBehalf(address from, address to, uint256 amount) external onlyOwner {
+    // @audit-info why we have onlyOwner here, it looks like rug, we should set different role here
     _redeem(from, to, amount.toUint104());
   }
 
@@ -232,22 +254,115 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
     _setCooldownSeconds(cooldownSeconds);
   }
 
+  // TODO: @pavelvm5 wouldn't sort functions here cause it can cause conflicts, but standard is external - first, public ..,  internal, makes easier to understand everything
   /// @inheritdoc IStakeToken
-  function getCooldownSeconds() external view returns (uint256) {
+  function setTreasury(address treasury) external onlyOwner {
+    _setTreasury(treasury);
+  }
+
+  function _setTreasury(address newTreasury) internal {
+    require(newTreasury != address(0), 'INVALID_TREASURY_ADDRESS'); // @audit-info-gas/bytecode-optimization @pavelvm5 propose to change from require(true) to if(false) revert ERROR_NAME(), cause its more optimized
+
+    _getStakeTokenStorage()._smConfig.treasury = newTreasury;
+
+    emit TreasuryChanged(newTreasury);
+  }
+
+  /// @inheritdoc IStakeToken
+  function setMaxFee(uint256 maxFee) external onlyOwner {
+    _setMaxFee(maxFee);
+  }
+
+  function _setMaxFee(uint256 newMaxFee) internal {
+    require(newMaxFee.toUint216() <= HUNDRED_PERCENT && newMaxFee > 0, 'INVALID_MAX_FEE_PARAMETER');
+
+    _getStakeTokenStorage()._smConfig.maxFee = newMaxFee.toUint216();
+
+    emit MaxFeeChanged(newMaxFee);
+  }
+
+  /// @inheritdoc IStakeToken
+  function setMaxReductionSeconds(uint256 newMaxReductionTime) external onlyOwner {
+    _setMaxReductionSeconds(newMaxReductionTime);
+  }
+
+  function _setMaxReductionSeconds(uint256 newMaxReductionSeconds) internal {
     StakeTokenStorage storage $ = _getStakeTokenStorage();
-    return $._smConfig.cooldownSeconds;
+    SmConfig memory smConfig = $._smConfig;
+
+    require(
+      newMaxReductionSeconds <= smConfig.defaultCooldownSeconds,
+      'INVALID_MAX_REDUCTION_TIME'
+    );
+
+    _getStakeTokenStorage()._smConfig.maxReductionSeconds = newMaxReductionSeconds.toUint32();
+
+    emit MaxReductionSecondsChanged(newMaxReductionSeconds);
+  }
+
+  /// @inheritdoc IStakeToken
+  function getDefaultCooldownSeconds() external view returns (uint256) {
+    StakeTokenStorage storage $ = _getStakeTokenStorage();
+    return $._smConfig.defaultCooldownSeconds;
+  }
+
+  /// @inheritdoc IStakeToken
+  function getMaxReductionSeconds() public view returns (uint256) {
+    StakeTokenStorage storage $ = _getStakeTokenStorage();
+
+    return $._smConfig.maxReductionSeconds;
   }
 
   function _cooldown(address from) internal {
     uint256 amount = balanceOf(from);
+
     require(amount != 0, 'INVALID_BALANCE_ON_COOLDOWN');
+
+    StakeTokenStorage storage $ = _getStakeTokenStorage();
+
+    uint32 timeToRedeem = (block.timestamp + $._smConfig.defaultCooldownSeconds).toUint32();
+
+    $._stakersCooldowns[from] = CooldownSetup({
+      timestamp: timeToRedeem,
+      amount: amount.toUint216()
+    });
+
+    emit Cooldown(from, amount, timeToRedeem);
+  }
+
+  function _reducedCooldown(address from, uint32 reductionTime) internal {
     StakeTokenStorage storage $ = _getStakeTokenStorage();
-    $._stakersCooldowns[from] = CooldownSnapshot({
-      timestamp: uint40(block.timestamp),
-      amount: uint216(amount)
+    SmConfig memory smConfig = $._smConfig;
+
+    require(reductionTime <= smConfig.maxReductionSeconds, 'MAX_REDUCTION_TIME_EXCEEDED');
+
+    uint216 amount = balanceOf(from).toUint216();
+    require(amount != 0, 'INVALID_BALANCE_ON_COOLDOWN');
+
+    uint216 feeAmount = (amount * smConfig.maxFee * reductionTime) /
+      smConfig.maxReductionSeconds /
+      HUNDRED_PERCENT;
+
+    uint32 timeToRedeem = (block.timestamp + smConfig.defaultCooldownSeconds - reductionTime)
+      .toUint32();
+
+    // @pavelvm5 move to internal block, cause it's logically responsible for transferring fees to treasury, don't want to make another internal function
+    {
+      uint256 underlyingForTreasury = previewRedeem(feeAmount);
+
+      _burn(from, feeAmount);
+
+      IERC20(smConfig.stakedToken).safeTransfer(smConfig.treasury, underlyingForTreasury);
+
+      emit FeesSentToTreasury(underlyingForTreasury);
+    }
+
+    $._stakersCooldowns[from] = CooldownSetup({
+      timestamp: timeToRedeem,
+      amount: amount - feeAmount
     });
 
-    emit Cooldown(from, amount);
+    emit Cooldown(from, amount, timeToRedeem);
   }
 
   /**
@@ -256,7 +371,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
    */
   function _setCooldownSeconds(uint256 cooldownSeconds) internal {
     StakeTokenStorage storage $ = _getStakeTokenStorage();
-    $._smConfig.cooldownSeconds = cooldownSeconds.toUint32();
+    $._smConfig.defaultCooldownSeconds = cooldownSeconds.toUint32();
     emit CooldownSecondsChanged(cooldownSeconds);
   }
 
@@ -271,7 +386,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
     uint256 sharesToMint = previewStake(amount);
     require(sharesToMint != 0, 'INVALID_ZERO_AMOUNT_AFTER_CONVERSION');
 
-    _mint(to, sharesToMint.toUint104());
+    _mint(to, sharesToMint);
 
     StakeTokenStorage storage $ = _getStakeTokenStorage();
     IERC20($._smConfig.stakedToken).safeTransferFrom(from, address(this), amount);
@@ -289,20 +404,17 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
     require(amount != 0, 'INVALID_ZERO_AMOUNT');
 
     StakeTokenStorage storage $ = _getStakeTokenStorage();
-    CooldownSnapshot memory cooldownSnapshot = $._stakersCooldowns[from];
+    CooldownSetup memory cooldownSetup = $._stakersCooldowns[from];
     SmConfig memory cachedSmConfig = $._smConfig;
+
+    require(block.timestamp >= cooldownSetup.timestamp, 'INSUFFICIENT_COOLDOWN');
     require(
-      (block.timestamp >= cooldownSnapshot.timestamp + cachedSmConfig.cooldownSeconds),
-      'INSUFFICIENT_COOLDOWN'
-    );
-    require(
-      (block.timestamp - (cooldownSnapshot.timestamp + cachedSmConfig.cooldownSeconds) <=
-        cachedSmConfig.unstakeWindowSeconds),
+      block.timestamp - cooldownSetup.timestamp <= cachedSmConfig.unstakeWindowSeconds,
       'UNSTAKE_WINDOW_FINISHED'
     );
 
-    uint256 maxRedeemable = cooldownSnapshot.amount;
-    require(maxRedeemable != 0, 'INVALID_ZERO_MAX_REDEEMABLE');
+    uint256 maxRedeemable = cooldownSetup.amount;
+    require(maxRedeemable != 0, 'INVALID_ZERO_MAX_REDEEMABLE'); // @audit-info @pavelvm5 need to check this, I think it's unreachable
 
     uint256 amountToRedeem = (amount > maxRedeemable) ? maxRedeemable : amount;
 
@@ -329,15 +441,15 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
   /**
    * @dev calculates the exchange rate based on totalAssets and totalShares
    * @dev always rounds up to ensure 100% backing of shares by rounding in favor of the contract
-   * @param totalAssets The total amount of assets staked
-   * @param totalShares The total amount of shares
+   * @param _totalAssets The total amount of assets staked
+   * @param _totalShares The total amount of shares
    * @return exchangeRate as 18 decimal precision uint216
    */
   function _getExchangeRate(
-    uint256 totalAssets,
-    uint256 totalShares
+    uint256 _totalAssets,
+    uint256 _totalShares
   ) internal pure returns (uint216) {
-    return (((totalShares * EXCHANGE_RATE_UNIT) + totalAssets - 1) / totalAssets).toUint216();
+    return (((_totalShares * EXCHANGE_RATE_UNIT) + _totalAssets - 1) / _totalAssets).toUint216(); // @audit-info @pavelvm5 shadow declaration warning here, changed names
   }
 
   function _update(address from, address to, uint256 amount) internal override {
@@ -353,12 +465,14 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
       // Sender
       REWARDS_CONTROLLER.handleAction(from, cachedTotalSupply, balanceOfFrom);
       StakeTokenStorage storage $ = _getStakeTokenStorage();
-      CooldownSnapshot memory previousSenderCooldown = $._stakersCooldowns[from];
+      CooldownSetup memory previousSenderCooldown = $._stakersCooldowns[from];
       if (previousSenderCooldown.timestamp != 0) {
         // update to 0 means redeem
         // this is based on the assumption that erc20 forbids transfer to 0
         if (to == address(0)) {
+          // @audit-info  @pavelvm5 can modify this place for optimization, cause sending fees to treasure and burning inside reducedCooldown will go here and this place will be overwritten in any way
           if (previousSenderCooldown.amount <= amount) {
+            // @audit-info @pavelvm5 should be == here, cause we shouldn't be able to redeem more than we have in cooldown, we can add this check
             delete $._stakersCooldowns[from];
           } else {
             $._stakersCooldowns[from].amount = uint216(previousSenderCooldown.amount - amount);
diff --git a/tests/Cooldown.t.sol b/tests/Cooldown.t.sol
index d09fddc..663d1c4 100644
--- a/tests/Cooldown.t.sol
+++ b/tests/Cooldown.t.sol
@@ -22,15 +22,16 @@ contract Cooldown is StkTestUtils {
 
     vm.startPrank(user);
     stakeToken.cooldown();
-    IStakeToken.CooldownSnapshot memory snapshotBefore = stakeToken.stakersCooldowns(user);
-    assertEq(snapshotBefore.timestamp, block.timestamp);
-    assertEq(snapshotBefore.amount, amountToStake);
+    IStakeToken.CooldownSetup memory cooldownBefore = stakeToken.stakersCooldowns(user);
+    // @pavelvm5 rewrite here, cause now we set time for redeem unlock, not when cooldown was called
+    assertEq(cooldownBefore.timestamp, block.timestamp + stakeToken.getDefaultCooldownSeconds());
+    assertEq(cooldownBefore.amount, amountToStake);
 
-    vm.warp(block.timestamp + stakeToken.getCooldownSeconds());
+    vm.warp(block.timestamp + stakeToken.getDefaultCooldownSeconds());
     _redeem(amountToRedeem, user, user);
 
-    IStakeToken.CooldownSnapshot memory snapshotAfter = stakeToken.stakersCooldowns(user);
-    assertEq(snapshotAfter.amount, amountToStake - amountToRedeem);
+    IStakeToken.CooldownSetup memory cooldownAfter = stakeToken.stakersCooldowns(user);
+    assertEq(cooldownAfter.amount, amountToStake - amountToRedeem);
   }
 
   function test_cooldownNoIncreaseInAmount(
@@ -47,16 +48,16 @@ contract Cooldown is StkTestUtils {
     vm.startPrank(user);
     stakeToken.cooldown();
 
-    IStakeToken.CooldownSnapshot memory snapshotBefore = stakeToken.stakersCooldowns(user);
+    IStakeToken.CooldownSetup memory cooldownBefore = stakeToken.stakersCooldowns(user);
 
     // increase amount
     _stake(amountToTopUp, user);
 
-    IStakeToken.CooldownSnapshot memory snapshotAfter = stakeToken.stakersCooldowns(user);
-    assertEq(snapshotBefore.timestamp, snapshotAfter.timestamp);
-    assertEq(snapshotBefore.amount, snapshotAfter.amount);
-    assertEq(snapshotAfter.timestamp, block.timestamp);
-    assertEq(snapshotAfter.amount, amountToStake);
+    IStakeToken.CooldownSetup memory cooldownAfter = stakeToken.stakersCooldowns(user);
+    assertEq(cooldownBefore.timestamp, cooldownAfter.timestamp);
+    assertEq(cooldownBefore.amount, cooldownAfter.amount);
+    assertEq(cooldownAfter.timestamp, block.timestamp + stakeToken.getDefaultCooldownSeconds());
+    assertEq(cooldownAfter.amount, amountToStake);
   }
 
   function test_cooldownOnTransfer(
@@ -77,30 +78,31 @@ contract Cooldown is StkTestUtils {
     vm.startPrank(user);
     stakeToken.cooldown();
 
-    IStakeToken.CooldownSnapshot memory snapshot0 = stakeToken.stakersCooldowns(user);
+    IStakeToken.CooldownSetup memory cooldown0 = stakeToken.stakersCooldowns(user);
 
     // Receiving token should not affect the amount
     _stake(amountToStakeOther, otherUser);
     vm.prank(otherUser);
     stakeToken.transfer(user, amountToStakeOther);
-    IStakeToken.CooldownSnapshot memory snapshot1 = stakeToken.stakersCooldowns(user);
-    assertEq(snapshot0.timestamp, snapshot1.timestamp, 'MISMATCH_BEFORE_COOLDOWN');
-    assertEq(snapshot0.amount, snapshot1.amount, 'MISMATCH_BEFORE_COOLDOWN_AMOUNT');
+    IStakeToken.CooldownSetup memory cooldown1 = stakeToken.stakersCooldowns(user);
+
+    assertEq(cooldown0.timestamp, cooldown1.timestamp, 'MISMATCH_BEFORE_COOLDOWN');
+    assertEq(cooldown0.amount, cooldown1.amount, 'MISMATCH_BEFORE_COOLDOWN_AMOUNT');
 
     // Sending token should not affect the amount as long as balance > amount
     vm.prank(user);
     stakeToken.transfer(otherUser, amountToStakeOther);
-    IStakeToken.CooldownSnapshot memory snapshot2 = stakeToken.stakersCooldowns(user);
-    assertEq(snapshot0.timestamp, snapshot2.timestamp, 'MISMATCH_COOLDOWN');
-    assertEq(snapshot0.amount, snapshot2.amount, 'MISMATCH_COOLDOWN_AMOUNT');
+    IStakeToken.CooldownSetup memory cooldown2 = stakeToken.stakersCooldowns(user);
+    assertEq(cooldown1.timestamp, cooldown2.timestamp, 'MISMATCH_COOLDOWN');
+    assertEq(cooldown0.amount, cooldown2.amount, 'MISMATCH_COOLDOWN_AMOUNT');
 
     // Sending token should decrease the cooldown amount when balance <= amount
     vm.startPrank(user);
     stakeToken.transfer(otherUser, amountToStake);
     vm.stopPrank();
-    IStakeToken.CooldownSnapshot memory snapshot3 = stakeToken.stakersCooldowns(user);
-    assertEq(snapshot3.timestamp, 0, 'MISMATCH_AFTER_COOLDOWN');
-    assertEq(snapshot3.amount, 0, 'MISMATCH_AFTER_COOLDOWN_AMOUNT');
+    IStakeToken.CooldownSetup memory cooldown3 = stakeToken.stakersCooldowns(user);
+    assertEq(cooldown3.timestamp, 0, 'MISMATCH_AFTER_COOLDOWN');
+    assertEq(cooldown3.amount, 0, 'MISMATCH_AFTER_COOLDOWN_AMOUNT');
   }
 
   function test_cooldownInsufficient_shouldRevert(
@@ -111,7 +113,7 @@ contract Cooldown is StkTestUtils {
     address destination
   ) public {
     vm.assume(amountToUnstake != 0 && amountToStake >= amountToUnstake);
-    vm.assume(secondsAfterCooldownActivation < stakeToken.getCooldownSeconds());
+    vm.assume(secondsAfterCooldownActivation < stakeToken.getDefaultCooldownSeconds());
     vm.assume(user != address(proxyAdmin) && user != address(0) && destination != address(0));
 
     _stake(amountToStake, user);
@@ -134,7 +136,7 @@ contract Cooldown is StkTestUtils {
     vm.assume(amountToUnstake != 0 && amountToStake >= amountToUnstake);
     vm.assume(
       secondsAfterCooldownActivation >
-        stakeToken.getCooldownSeconds() + stakeToken.getUnstakeWindow()
+        stakeToken.getDefaultCooldownSeconds() + stakeToken.getUnstakeWindow()
     );
     vm.assume(user != address(proxyAdmin) && user != address(0) && destination != address(0));
 
@@ -169,14 +171,120 @@ contract Cooldown is StkTestUtils {
     vm.prank(user);
     stakeToken.cooldown();
     _stake(amountToTopUp, user);
-    IStakeToken.CooldownSnapshot memory snapshotAfterSecondStake = stakeToken.stakersCooldowns(
-      user
-    );
-    assertEq(snapshotAfterSecondStake.amount, amountToStake, 'STAKE_SHOULD_NOT_ALTER_COOLDOWN');
-    vm.warp(block.timestamp + stakeToken.getCooldownSeconds());
+    IStakeToken.CooldownSetup memory cooldownAfterSecondStake = stakeToken.stakersCooldowns(user);
+    assertEq(cooldownAfterSecondStake.amount, amountToStake, 'STAKE_SHOULD_NOT_ALTER_COOLDOWN');
+
+    vm.warp(block.timestamp + stakeToken.getDefaultCooldownSeconds());
     _redeem(amountToUnstake, user, destination);
 
     assertEq(underlyingToken.balanceOf(destination), amountToStake, 'WRONG_AMOUNT_REDEEMED');
     assertEq(stakeToken.balanceOf(user), amountToTopUp, 'WRONG_AMOUNT_LEFT');
   }
+
+  // @pavelvm5 fast-withdrawal tests start here, check only impact using reducedCooldown function, redeem/transfer process is the same
+  function test_reducedCooldown(
+    uint104 amountToStake,
+    uint104 amountToRedeem,
+    address user
+  ) public {
+    uint256 fee = (amountToStake * maxFee) / 10_000;
+
+    vm.assume(amountToStake >= amountToRedeem + fee && amountToRedeem > 0);
+    vm.assume(user != address(proxyAdmin) && user != address(0) && user != treasury);
+
+    _stake(amountToStake, user);
+
+    vm.startPrank(user);
+    stakeToken.reducedCooldown(maxReductionSeconds);
+
+    IStakeToken.CooldownSetup memory cooldownBefore = stakeToken.stakersCooldowns(user);
+
+    uint256 cooldownSeconds = stakeToken.getDefaultCooldownSeconds() - maxReductionSeconds;
+
+    assertEq(cooldownBefore.timestamp, block.timestamp + cooldownSeconds);
+
+    assertEq(cooldownBefore.amount, amountToStake - fee);
+
+    vm.warp(block.timestamp + cooldownSeconds);
+    _redeem(amountToRedeem, user, user);
+
+    IStakeToken.CooldownSetup memory cooldownAfter = stakeToken.stakersCooldowns(user);
+
+    console.log(cooldownAfter.amount);
+    console.log(amountToStake);
+    console.log(fee);
+    console.log(amountToRedeem);
+
+    assertEq(cooldownAfter.amount, amountToStake - fee - amountToRedeem);
+  }
+
+  function test_reducedCooldownTreasuryFees(
+    uint104 amountToStake,
+    uint104 amountToRedeem,
+    address user
+  ) public {
+    uint256 fee = (amountToStake * maxFee) / 10_000;
+
+    vm.assume(amountToStake >= amountToRedeem + fee && amountToRedeem > 0);
+    vm.assume(user != address(proxyAdmin) && user != address(0) && user != treasury);
+
+    _stake(amountToStake, user);
+
+    vm.startPrank(user);
+
+    uint256 totalSupplyBefore = stakeToken.totalSupply();
+
+    uint256 balanceTreasuryBefore = underlyingToken.balanceOf(treasury);
+    uint256 underlyingFeesToTreasury = stakeToken.previewRedeem(fee);
+
+    stakeToken.reducedCooldown(maxReductionSeconds);
+
+    uint256 totalSupplyAfter = stakeToken.totalSupply();
+
+    assertEq(totalSupplyBefore - totalSupplyAfter, fee);
+
+    uint256 balanceTreasuryAfter = underlyingToken.balanceOf(treasury);
+
+    assertEq(balanceTreasuryAfter - balanceTreasuryBefore, underlyingFeesToTreasury);
+  }
+
+  function test_reducedCooldownDifferentTime(
+    uint104 amountToStake,
+    uint104 amountToRedeem,
+    address user,
+    uint32 reducedTime
+  ) public {
+    uint256 fee = (amountToStake * maxFee * reducedTime) / maxReductionSeconds / 10_000;
+
+    vm.assume(amountToStake >= amountToRedeem + fee && amountToRedeem > 0);
+    vm.assume(user != address(proxyAdmin) && user != address(0) && user != treasury);
+    vm.assume(reducedTime <= maxReductionSeconds);
+
+    _stake(amountToStake, user);
+
+    vm.startPrank(user);
+
+    uint256 balanceTreasuryBefore = underlyingToken.balanceOf(treasury);
+    uint256 underlyingFeesToTreasury = stakeToken.previewRedeem(fee);
+
+    stakeToken.reducedCooldown(reducedTime);
+
+    uint256 balanceTreasuryAfter = underlyingToken.balanceOf(treasury);
+
+    assertEq(balanceTreasuryAfter - balanceTreasuryBefore, underlyingFeesToTreasury);
+
+    IStakeToken.CooldownSetup memory cooldownBefore = stakeToken.stakersCooldowns(user);
+
+    uint256 cooldownSeconds = stakeToken.getDefaultCooldownSeconds() - reducedTime;
+
+    assertEq(cooldownBefore.timestamp, block.timestamp + cooldownSeconds);
+
+    assertEq(cooldownBefore.amount, amountToStake - fee);
+
+    vm.warp(block.timestamp + cooldownSeconds);
+    _redeem(amountToRedeem, user, user);
+
+    IStakeToken.CooldownSetup memory cooldownAfter = stakeToken.stakersCooldowns(user);
+    assertEq(cooldownAfter.amount, amountToStake - fee - amountToRedeem);
+  }
 }
diff --git a/tests/ERC20Std.t.sol b/tests/ERC20Std.t.sol
index 5fd394b..bb18acd 100644
--- a/tests/ERC20Std.t.sol
+++ b/tests/ERC20Std.t.sol
@@ -37,7 +37,7 @@ contract ERC20Std is StkTestUtils {
 
     vm.prank(USER);
     stakeToken.cooldown();
-    vm.warp(block.timestamp + stakeToken.getCooldownSeconds());
+    vm.warp(block.timestamp + stakeToken.getDefaultCooldownSeconds());
     _redeem(amountRedeemed, USER, destination);
 
     assertEq(stakeToken.totalSupply(), amountStaked - amountRedeemed);
diff --git a/tests/StkTestUtils.t.sol b/tests/StkTestUtils.t.sol
index 4d4857e..63562b5 100644
--- a/tests/StkTestUtils.t.sol
+++ b/tests/StkTestUtils.t.sol
@@ -25,6 +25,10 @@ contract StkTestUtils is Test {
   ProxyAdmin public proxyAdmin;
   StakeToken public stakeToken;
 
+  address public treasury;
+  uint256 public maxFee;
+  uint256 public maxReductionSeconds;
+
   function setUp() public virtual {
     underlyingToken = new MockERC20('TestToken', 'TEST');
     rewardToken = new MockERC20('TestReward', 'REWARD');
@@ -32,6 +36,11 @@ contract StkTestUtils is Test {
     EmissionManager manager = new EmissionManager(address(controller), admin);
     stakeTokenImpl = new StakeToken(IRewardsController(address(controller)));
     proxyAdmin = new ProxyAdmin(admin);
+
+    treasury = address(uint160(uint256(keccak256('treasury'))));
+    maxFee = 200; // 2% in BIPS
+    maxReductionSeconds = 10 days;
+
     stakeToken = StakeToken(
       address(
         new TransparentUpgradeableProxy(
@@ -44,7 +53,10 @@ contract StkTestUtils is Test {
             'stkTest',
             admin,
             15 days,
-            2 days
+            2 days,
+            treasury,
+            maxFee,
+            maxReductionSeconds
           )
         )
       )

From 7d105a48ccb3fcda08dd6751aba58723875a783e Mon Sep 17 00:00:00 2001
From: Pavel Menshikov <menshikov.pv@mail.ru>
Date: Mon, 15 Jul 2024 17:13:42 +0500
Subject: [PATCH 2/5] changed maxReductionTime to minCooldownSeconds, changed
 uint216 maxFee to uint16

---
 src/contracts/IStakeToken.sol | 20 +++++++++------
 src/contracts/StakeToken.sol  | 47 +++++++++++++++++++++--------------
 tests/Cooldown.t.sol          |  4 +++
 tests/StkTestUtils.t.sol      |  6 ++---
 4 files changed, 48 insertions(+), 29 deletions(-)

diff --git a/src/contracts/IStakeToken.sol b/src/contracts/IStakeToken.sol
index ed60ce9..55554fe 100644
--- a/src/contracts/IStakeToken.sol
+++ b/src/contracts/IStakeToken.sol
@@ -17,10 +17,10 @@ interface IStakeToken {
     uint32 defaultCooldownSeconds;
     /// @notice The address of the underlying asset
     address stakedToken;
-    /// @notice The maximum time available for reduction cooldown period
-    uint32 maxReductionSeconds;
+    /// @notice The minimum cooldown time available for redeeming assets
+    uint32 minCooldownSeconds;
     /// @notice The maximum fee in BIPS that will be taken when the cooldown period is reduced by maxReductionTime
-    uint216 maxFee;
+    uint16 maxFee;
     /// @notice The address of treasury
     address treasury;
   }
@@ -29,7 +29,7 @@ interface IStakeToken {
   event FeesSentToTreasury(uint256 amount);
   event TreasuryChanged(address treasury);
   event MaxFeeChanged(uint256 maxFee);
-  event MaxReductionSecondsChanged(uint256 maxReductionSeconds);
+  event MinCooldownSecondsChanged(uint256 maxReductionSeconds);
 
   event Staked(address indexed from, address indexed to, uint256 assets, uint256 shares);
   event Redeem(address indexed from, address indexed to, uint256 assets, uint256 shares);
@@ -116,6 +116,12 @@ interface IStakeToken {
    */
   function getMaxReductionSeconds() external view returns (uint256);
 
+  /**
+   * @dev Getter of the minimum cooldown seconds possible
+   * @return minCooldownSeconds The minimum cooldown time available
+   */
+  function getMinCooldownSeconds() external view returns (uint256);
+
   /**
    * @dev Setter of cooldown seconds
    * Can only be called by the cooldown admin
@@ -136,10 +142,10 @@ interface IStakeToken {
   function setMaxFee(uint256 maxFee) external;
 
   /**
-   * @dev Setter of max cooldown reduction time in seconds
-   * @param newMaxReductionTime number of seconds by which the cooldown can be reduced with the payment of fees
+   * @dev Setter of min cooldown time in seconds
+   * @param newMinCooldownSeconds number of seconds the cooldown can be reduced to with the payment of fees
    */
-  function setMaxReductionSeconds(uint256 newMaxReductionTime) external;
+  function setMinCooldownSeconds(uint256 newMinCooldownSeconds) external;
 
   /**
    * @dev returns the exact amount of shares that would be received for the provided number of assets
diff --git a/src/contracts/StakeToken.sol b/src/contracts/StakeToken.sol
index eb6b5e6..ef43bd1 100644
--- a/src/contracts/StakeToken.sol
+++ b/src/contracts/StakeToken.sol
@@ -21,9 +21,9 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
   using SafeCast for uint256;
   using SafeCast for uint104;
 
+  uint16 public constant HUNDRED_PERCENT = 10_000;
   uint216 public constant INITIAL_EXCHANGE_RATE = 1e18;
   uint256 public constant EXCHANGE_RATE_UNIT = 1e18;
-  uint216 public constant HUNDRED_PERCENT = 10_000;
 
   IRewardsController public immutable REWARDS_CONTROLLER;
 
@@ -74,7 +74,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
     uint256 unstakeWindow,
     address treasury,
     uint256 maxFee,
-    uint256 maxReductionSeconds
+    uint256 minCooldownSeconds
   ) external virtual initializer {
     // @pavelvm5 made like this, due to stack-too-deep error, will fix in future, if any args will be added could be useful to rewrite to assembly
     _getStakeTokenStorage()._smConfig.stakedToken = stakedToken;
@@ -91,7 +91,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
 
     _setCooldownSeconds(cooldownSeconds);
     _setUnstakeWindow(unstakeWindow);
-    _setMaxReductionSeconds(maxReductionSeconds);
+    _setMinCooldownSeconds(minCooldownSeconds);
 
     _updateExchangeRate(INITIAL_EXCHANGE_RATE);
   }
@@ -274,30 +274,27 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
   }
 
   function _setMaxFee(uint256 newMaxFee) internal {
-    require(newMaxFee.toUint216() <= HUNDRED_PERCENT && newMaxFee > 0, 'INVALID_MAX_FEE_PARAMETER');
+    require(newMaxFee.toUint16() <= HUNDRED_PERCENT && newMaxFee > 0, 'INVALID_MAX_FEE_PARAMETER');
 
-    _getStakeTokenStorage()._smConfig.maxFee = newMaxFee.toUint216();
+    _getStakeTokenStorage()._smConfig.maxFee = newMaxFee.toUint16();
 
     emit MaxFeeChanged(newMaxFee);
   }
 
   /// @inheritdoc IStakeToken
-  function setMaxReductionSeconds(uint256 newMaxReductionTime) external onlyOwner {
-    _setMaxReductionSeconds(newMaxReductionTime);
+  function setMinCooldownSeconds(uint256 newMinCooldownSeconds) external onlyOwner {
+    _setMinCooldownSeconds(newMinCooldownSeconds);
   }
 
-  function _setMaxReductionSeconds(uint256 newMaxReductionSeconds) internal {
-    StakeTokenStorage storage $ = _getStakeTokenStorage();
-    SmConfig memory smConfig = $._smConfig;
-
+  function _setMinCooldownSeconds(uint256 newMinCooldownSeconds) internal {
     require(
-      newMaxReductionSeconds <= smConfig.defaultCooldownSeconds,
-      'INVALID_MAX_REDUCTION_TIME'
+      newMinCooldownSeconds <= _getStakeTokenStorage()._smConfig.defaultCooldownSeconds,
+      'INVALID_MIN_COOLDOWN_SECONDS'
     );
 
-    _getStakeTokenStorage()._smConfig.maxReductionSeconds = newMaxReductionSeconds.toUint32();
+    _getStakeTokenStorage()._smConfig.minCooldownSeconds = newMinCooldownSeconds.toUint32();
 
-    emit MaxReductionSecondsChanged(newMaxReductionSeconds);
+    emit MinCooldownSecondsChanged(newMinCooldownSeconds);
   }
 
   /// @inheritdoc IStakeToken
@@ -310,7 +307,13 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
   function getMaxReductionSeconds() public view returns (uint256) {
     StakeTokenStorage storage $ = _getStakeTokenStorage();
 
-    return $._smConfig.maxReductionSeconds;
+    return $._smConfig.defaultCooldownSeconds - $._smConfig.minCooldownSeconds;
+  }
+
+  function getMinCooldownSeconds() public view returns (uint256) {
+    StakeTokenStorage storage $ = _getStakeTokenStorage();
+
+    return $._smConfig.minCooldownSeconds;
   }
 
   function _cooldown(address from) internal {
@@ -334,13 +337,16 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
     StakeTokenStorage storage $ = _getStakeTokenStorage();
     SmConfig memory smConfig = $._smConfig;
 
-    require(reductionTime <= smConfig.maxReductionSeconds, 'MAX_REDUCTION_TIME_EXCEEDED');
+    require(
+      reductionTime <= smConfig.defaultCooldownSeconds - smConfig.minCooldownSeconds,
+      'MAX_REDUCTION_TIME_EXCEEDED'
+    );
 
     uint216 amount = balanceOf(from).toUint216();
     require(amount != 0, 'INVALID_BALANCE_ON_COOLDOWN');
 
     uint216 feeAmount = (amount * smConfig.maxFee * reductionTime) /
-      smConfig.maxReductionSeconds /
+      (smConfig.defaultCooldownSeconds - smConfig.minCooldownSeconds) /
       HUNDRED_PERCENT;
 
     uint32 timeToRedeem = (block.timestamp + smConfig.defaultCooldownSeconds - reductionTime)
@@ -371,7 +377,10 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
    */
   function _setCooldownSeconds(uint256 cooldownSeconds) internal {
     StakeTokenStorage storage $ = _getStakeTokenStorage();
+    require(cooldownSeconds >= $._smConfig.minCooldownSeconds, 'INVALID_COOLDOWN_SECONDS');
+
     $._smConfig.defaultCooldownSeconds = cooldownSeconds.toUint32();
+
     emit CooldownSecondsChanged(cooldownSeconds);
   }
 
@@ -405,7 +414,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
 
     StakeTokenStorage storage $ = _getStakeTokenStorage();
     CooldownSetup memory cooldownSetup = $._stakersCooldowns[from];
-    SmConfig memory cachedSmConfig = $._smConfig;
+    SmConfig memory cachedSmConfig = $._smConfig; // @audit I think this copying is less optimized too, we read/copy 2 slots here, instead of double-time reading the same slot
 
     require(block.timestamp >= cooldownSetup.timestamp, 'INSUFFICIENT_COOLDOWN');
     require(
diff --git a/tests/Cooldown.t.sol b/tests/Cooldown.t.sol
index 663d1c4..c4cbdff 100644
--- a/tests/Cooldown.t.sol
+++ b/tests/Cooldown.t.sol
@@ -195,6 +195,8 @@ contract Cooldown is StkTestUtils {
     _stake(amountToStake, user);
 
     vm.startPrank(user);
+
+    uint256 maxReductionSeconds = stakeToken.getMaxReductionSeconds();
     stakeToken.reducedCooldown(maxReductionSeconds);
 
     IStakeToken.CooldownSetup memory cooldownBefore = stakeToken.stakersCooldowns(user);
@@ -237,6 +239,7 @@ contract Cooldown is StkTestUtils {
     uint256 balanceTreasuryBefore = underlyingToken.balanceOf(treasury);
     uint256 underlyingFeesToTreasury = stakeToken.previewRedeem(fee);
 
+    uint256 maxReductionSeconds = stakeToken.getMaxReductionSeconds();
     stakeToken.reducedCooldown(maxReductionSeconds);
 
     uint256 totalSupplyAfter = stakeToken.totalSupply();
@@ -254,6 +257,7 @@ contract Cooldown is StkTestUtils {
     address user,
     uint32 reducedTime
   ) public {
+    uint256 maxReductionSeconds = stakeToken.getMaxReductionSeconds();
     uint256 fee = (amountToStake * maxFee * reducedTime) / maxReductionSeconds / 10_000;
 
     vm.assume(amountToStake >= amountToRedeem + fee && amountToRedeem > 0);
diff --git a/tests/StkTestUtils.t.sol b/tests/StkTestUtils.t.sol
index 63562b5..94ee1f6 100644
--- a/tests/StkTestUtils.t.sol
+++ b/tests/StkTestUtils.t.sol
@@ -27,7 +27,7 @@ contract StkTestUtils is Test {
 
   address public treasury;
   uint256 public maxFee;
-  uint256 public maxReductionSeconds;
+  uint256 public minCooldownSeconds;
 
   function setUp() public virtual {
     underlyingToken = new MockERC20('TestToken', 'TEST');
@@ -39,7 +39,7 @@ contract StkTestUtils is Test {
 
     treasury = address(uint160(uint256(keccak256('treasury'))));
     maxFee = 200; // 2% in BIPS
-    maxReductionSeconds = 10 days;
+    minCooldownSeconds = 5 days;
 
     stakeToken = StakeToken(
       address(
@@ -56,7 +56,7 @@ contract StkTestUtils is Test {
             2 days,
             treasury,
             maxFee,
-            maxReductionSeconds
+            minCooldownSeconds
           )
         )
       )

From 97669726af4166501ce3d241b56a4b66a7ed8b0a Mon Sep 17 00:00:00 2001
From: Pavel Menshikov <menshikov.pv@mail.ru>
Date: Mon, 15 Jul 2024 17:21:27 +0500
Subject: [PATCH 3/5] reduced some unnecessary code in new functions

---
 src/contracts/StakeToken.sol | 17 ++++++-----------
 1 file changed, 6 insertions(+), 11 deletions(-)

diff --git a/src/contracts/StakeToken.sol b/src/contracts/StakeToken.sol
index ef43bd1..1ca1415 100644
--- a/src/contracts/StakeToken.sol
+++ b/src/contracts/StakeToken.sol
@@ -118,15 +118,13 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
   function _setUnstakeWindow(uint256 newUnstakeWindow) internal {
     require(newUnstakeWindow >= 1 hours, 'TOO_LOW_UNSTAKE_WINDOW');
 
-    StakeTokenStorage storage $ = _getStakeTokenStorage();
-    $._smConfig.unstakeWindowSeconds = newUnstakeWindow.toUint32();
+    _getStakeTokenStorage()._smConfig.unstakeWindowSeconds = newUnstakeWindow.toUint32();
 
     emit UnstakeWindowChanged(newUnstakeWindow);
   }
 
   function getUnstakeWindow() external view returns (uint256) {
-    StakeTokenStorage storage $ = _getStakeTokenStorage();
-    return $._smConfig.unstakeWindowSeconds;
+    return _getStakeTokenStorage()._smConfig.unstakeWindowSeconds;
   }
 
   function setSlashingAdmin(address newSlashingAdmin) external onlyOwner {
@@ -184,7 +182,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
 
   /// @inheritdoc IStakeToken
   function cooldownOnBehalfOf(address from) external onlyOwner {
-    // @audit-info same modif here
+    // @audit-info same modif onlyOwner here
     _cooldown(from);
   }
 
@@ -299,8 +297,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
 
   /// @inheritdoc IStakeToken
   function getDefaultCooldownSeconds() external view returns (uint256) {
-    StakeTokenStorage storage $ = _getStakeTokenStorage();
-    return $._smConfig.defaultCooldownSeconds;
+    return _getStakeTokenStorage()._smConfig.defaultCooldownSeconds;
   }
 
   /// @inheritdoc IStakeToken
@@ -311,9 +308,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
   }
 
   function getMinCooldownSeconds() public view returns (uint256) {
-    StakeTokenStorage storage $ = _getStakeTokenStorage();
-
-    return $._smConfig.minCooldownSeconds;
+    return _getStakeTokenStorage()._smConfig.minCooldownSeconds;
   }
 
   function _cooldown(address from) internal {
@@ -414,7 +409,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
 
     StakeTokenStorage storage $ = _getStakeTokenStorage();
     CooldownSetup memory cooldownSetup = $._stakersCooldowns[from];
-    SmConfig memory cachedSmConfig = $._smConfig; // @audit I think this copying is less optimized too, we read/copy 2 slots here, instead of double-time reading the same slot
+    SmConfig memory cachedSmConfig = $._smConfig; // @audit I think this copying is less optimized, we read&copy 2 slots here, instead of double-time reading the same slot
 
     require(block.timestamp >= cooldownSetup.timestamp, 'INSUFFICIENT_COOLDOWN');
     require(

From 48c395091def3d346db8f5b1d58e652c157d419d Mon Sep 17 00:00:00 2001
From: Pavel Menshikov <menshikov.pv@mail.ru>
Date: Thu, 18 Jul 2024 03:19:20 +0500
Subject: [PATCH 4/5] added functionality to sum reduction for the same
 cooldown

---
 src/contracts/IStakeToken.sol |  13 +-
 src/contracts/StakeToken.sol  |  72 +++---
 tests/Cooldown.t.sol          | 427 ++++++++++++++++++++--------------
 3 files changed, 300 insertions(+), 212 deletions(-)

diff --git a/src/contracts/IStakeToken.sol b/src/contracts/IStakeToken.sol
index 55554fe..55ecc12 100644
--- a/src/contracts/IStakeToken.sol
+++ b/src/contracts/IStakeToken.sol
@@ -2,30 +2,29 @@
 pragma solidity ^0.8.0;
 
 interface IStakeToken {
-  struct CooldownSetup {
-    // make more sense to display time from which can be withdrawn, not activation time
+  struct CooldownSnapshot {
     /// @notice The time after which funds can be redeemed
-    uint32 timestamp;
+    uint40 timestamp;
     /// @notice The amount of tokens which can be redeemed
     uint216 amount;
   }
 
   struct SmConfig {
     /// @notice Seconds available to redeem once the cooldown period is fulfilled
-    uint32 unstakeWindowSeconds;
+    uint40 unstakeWindowSeconds;
     /// @notice Seconds between starting cooldown and being able to withdraw
-    uint32 defaultCooldownSeconds;
+    uint40 defaultCooldownSeconds;
     /// @notice The address of the underlying asset
     address stakedToken;
     /// @notice The minimum cooldown time available for redeeming assets
-    uint32 minCooldownSeconds;
+    uint40 minCooldownSeconds;
     /// @notice The maximum fee in BIPS that will be taken when the cooldown period is reduced by maxReductionTime
     uint16 maxFee;
     /// @notice The address of treasury
     address treasury;
   }
 
-  event Cooldown(address indexed user, uint256 amount, uint32 timeToRedeem);
+  event Cooldown(address indexed user, uint256 amount, uint256 timeToRedeem);
   event FeesSentToTreasury(uint256 amount);
   event TreasuryChanged(address treasury);
   event MaxFeeChanged(uint256 maxFee);
diff --git a/src/contracts/StakeToken.sol b/src/contracts/StakeToken.sol
index 1ca1415..e529bf7 100644
--- a/src/contracts/StakeToken.sol
+++ b/src/contracts/StakeToken.sol
@@ -29,7 +29,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
 
   /// @custom:storage-location erc7201:aave.storage.StakeToken
   struct StakeTokenStorage {
-    mapping(address => CooldownSetup) _stakersCooldowns;
+    mapping(address => CooldownSnapshot) _stakersCooldowns;
     SmConfig _smConfig;
     /// @notice Current exchangeRate of the stk
     uint216 _currentExchangeRate;
@@ -55,7 +55,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
     }
   }
 
-  function stakersCooldowns(address user) public view returns (CooldownSetup memory) {
+  function stakersCooldowns(address user) public view returns (CooldownSnapshot memory) {
     StakeTokenStorage storage $ = _getStakeTokenStorage();
     return $._stakersCooldowns[user];
   }
@@ -76,7 +76,6 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
     uint256 maxFee,
     uint256 minCooldownSeconds
   ) external virtual initializer {
-    // @pavelvm5 made like this, due to stack-too-deep error, will fix in future, if any args will be added could be useful to rewrite to assembly
     _getStakeTokenStorage()._smConfig.stakedToken = stakedToken;
     _getStakeTokenStorage()._minAssetsRemaining = 10 ** decimals();
 
@@ -118,7 +117,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
   function _setUnstakeWindow(uint256 newUnstakeWindow) internal {
     require(newUnstakeWindow >= 1 hours, 'TOO_LOW_UNSTAKE_WINDOW');
 
-    _getStakeTokenStorage()._smConfig.unstakeWindowSeconds = newUnstakeWindow.toUint32();
+    _getStakeTokenStorage()._smConfig.unstakeWindowSeconds = newUnstakeWindow.toUint40();
 
     emit UnstakeWindowChanged(newUnstakeWindow);
   }
@@ -188,18 +187,18 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
 
   /// @inheritdoc IStakeToken
   function reducedCooldown(uint256 reductionTime) external {
-    _reducedCooldown(msg.sender, reductionTime.toUint32());
+    _reducedCooldown(msg.sender, uint40(reductionTime)); // @pavelvm no need safecast here, cause type(uint256).max should be okay, I guess
   }
 
   /// @inheritdoc IStakeToken
   function redeem(address to, uint256 amount) external {
-    _redeem(msg.sender, to, amount.toUint104());
+    _redeem(msg.sender, to, amount.toUint104()); // @pavelvm5 I think here we can remove safeCast too, so type(uint256).max will be ok and get max here
   }
 
   /// @inheritdoc IStakeToken
   function redeemOnBehalf(address from, address to, uint256 amount) external onlyOwner {
     // @audit-info why we have onlyOwner here, it looks like rug, we should set different role here
-    _redeem(from, to, amount.toUint104());
+    _redeem(from, to, amount.toUint104()); // @pavelvm5 I think here we can remove safeCast too, so type(uint256).max will be ok and get max here
   }
 
   /// @inheritdoc IStakeToken
@@ -290,7 +289,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
       'INVALID_MIN_COOLDOWN_SECONDS'
     );
 
-    _getStakeTokenStorage()._smConfig.minCooldownSeconds = newMinCooldownSeconds.toUint32();
+    _getStakeTokenStorage()._smConfig.minCooldownSeconds = newMinCooldownSeconds.toUint40();
 
     emit MinCooldownSecondsChanged(newMinCooldownSeconds);
   }
@@ -318,9 +317,9 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
 
     StakeTokenStorage storage $ = _getStakeTokenStorage();
 
-    uint32 timeToRedeem = (block.timestamp + $._smConfig.defaultCooldownSeconds).toUint32();
+    uint40 timeToRedeem = (block.timestamp + $._smConfig.defaultCooldownSeconds).toUint40();
 
-    $._stakersCooldowns[from] = CooldownSetup({
+    $._stakersCooldowns[from] = CooldownSnapshot({
       timestamp: timeToRedeem,
       amount: amount.toUint216()
     });
@@ -328,26 +327,41 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
     emit Cooldown(from, amount, timeToRedeem);
   }
 
-  function _reducedCooldown(address from, uint32 reductionTime) internal {
+  function _reducedCooldown(address from, uint40 reductionTime) internal {
     StakeTokenStorage storage $ = _getStakeTokenStorage();
+
+    CooldownSnapshot memory cooldownSnapshot = $._stakersCooldowns[from];
     SmConfig memory smConfig = $._smConfig;
 
-    require(
-      reductionTime <= smConfig.defaultCooldownSeconds - smConfig.minCooldownSeconds,
-      'MAX_REDUCTION_TIME_EXCEEDED'
-    );
+    uint40 maxReductionSeconds = smConfig.defaultCooldownSeconds - smConfig.minCooldownSeconds;
 
-    uint216 amount = balanceOf(from).toUint216();
-    require(amount != 0, 'INVALID_BALANCE_ON_COOLDOWN');
+    uint40 timeToRedeem;
+    uint216 amount;
+
+    if (block.timestamp + smConfig.minCooldownSeconds < cooldownSnapshot.timestamp) {
+      uint40 maxReductionAvailable = cooldownSnapshot.timestamp - smConfig.minCooldownSeconds;
+
+      reductionTime = reductionTime > maxReductionAvailable ? maxReductionAvailable : reductionTime;
+
+      timeToRedeem = cooldownSnapshot.timestamp - reductionTime;
+
+      amount = cooldownSnapshot.amount;
+    } else if (block.timestamp >= cooldownSnapshot.timestamp) {
+      reductionTime = reductionTime > maxReductionSeconds ? maxReductionSeconds : reductionTime;
+
+      timeToRedeem = (block.timestamp + smConfig.defaultCooldownSeconds - reductionTime).toUint40();
+
+      amount = balanceOf(from).toUint216();
+
+      require(amount != 0, 'INVALID_BALANCE_ON_COOLDOWN');
+    } else {
+      revert('MAX_REDUCTION_TIME_EXCEEDED');
+    }
 
     uint216 feeAmount = (amount * smConfig.maxFee * reductionTime) /
-      (smConfig.defaultCooldownSeconds - smConfig.minCooldownSeconds) /
+      maxReductionSeconds /
       HUNDRED_PERCENT;
 
-    uint32 timeToRedeem = (block.timestamp + smConfig.defaultCooldownSeconds - reductionTime)
-      .toUint32();
-
-    // @pavelvm5 move to internal block, cause it's logically responsible for transferring fees to treasury, don't want to make another internal function
     {
       uint256 underlyingForTreasury = previewRedeem(feeAmount);
 
@@ -358,7 +372,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
       emit FeesSentToTreasury(underlyingForTreasury);
     }
 
-    $._stakersCooldowns[from] = CooldownSetup({
+    $._stakersCooldowns[from] = CooldownSnapshot({
       timestamp: timeToRedeem,
       amount: amount - feeAmount
     });
@@ -374,7 +388,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
     StakeTokenStorage storage $ = _getStakeTokenStorage();
     require(cooldownSeconds >= $._smConfig.minCooldownSeconds, 'INVALID_COOLDOWN_SECONDS');
 
-    $._smConfig.defaultCooldownSeconds = cooldownSeconds.toUint32();
+    $._smConfig.defaultCooldownSeconds = cooldownSeconds.toUint40();
 
     emit CooldownSecondsChanged(cooldownSeconds);
   }
@@ -408,16 +422,16 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
     require(amount != 0, 'INVALID_ZERO_AMOUNT');
 
     StakeTokenStorage storage $ = _getStakeTokenStorage();
-    CooldownSetup memory cooldownSetup = $._stakersCooldowns[from];
+    CooldownSnapshot memory cooldownSnapshot = $._stakersCooldowns[from];
     SmConfig memory cachedSmConfig = $._smConfig; // @audit I think this copying is less optimized, we read&copy 2 slots here, instead of double-time reading the same slot
 
-    require(block.timestamp >= cooldownSetup.timestamp, 'INSUFFICIENT_COOLDOWN');
+    require(block.timestamp >= cooldownSnapshot.timestamp, 'INSUFFICIENT_COOLDOWN');
     require(
-      block.timestamp - cooldownSetup.timestamp <= cachedSmConfig.unstakeWindowSeconds,
+      block.timestamp - cooldownSnapshot.timestamp <= cachedSmConfig.unstakeWindowSeconds,
       'UNSTAKE_WINDOW_FINISHED'
     );
 
-    uint256 maxRedeemable = cooldownSetup.amount;
+    uint256 maxRedeemable = cooldownSnapshot.amount;
     require(maxRedeemable != 0, 'INVALID_ZERO_MAX_REDEEMABLE'); // @audit-info @pavelvm5 need to check this, I think it's unreachable
 
     uint256 amountToRedeem = (amount > maxRedeemable) ? maxRedeemable : amount;
@@ -469,7 +483,7 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
       // Sender
       REWARDS_CONTROLLER.handleAction(from, cachedTotalSupply, balanceOfFrom);
       StakeTokenStorage storage $ = _getStakeTokenStorage();
-      CooldownSetup memory previousSenderCooldown = $._stakersCooldowns[from];
+      CooldownSnapshot memory previousSenderCooldown = $._stakersCooldowns[from];
       if (previousSenderCooldown.timestamp != 0) {
         // update to 0 means redeem
         // this is based on the assumption that erc20 forbids transfer to 0
diff --git a/tests/Cooldown.t.sol b/tests/Cooldown.t.sol
index c4cbdff..a9f4f9b 100644
--- a/tests/Cooldown.t.sol
+++ b/tests/Cooldown.t.sol
@@ -14,172 +14,174 @@ contract Cooldown is StkTestUtils {
   /**
    * cooldown should activate for a given block.timestamp and cooldown the currently held amount
    */
-  function test_cooldown(uint104 amountToStake, uint104 amountToRedeem, address user) public {
-    vm.assume(amountToStake >= amountToRedeem && amountToRedeem > 0);
-    vm.assume(user != address(proxyAdmin) && user != address(0));
-
-    _stake(amountToStake, user);
-
-    vm.startPrank(user);
-    stakeToken.cooldown();
-    IStakeToken.CooldownSetup memory cooldownBefore = stakeToken.stakersCooldowns(user);
-    // @pavelvm5 rewrite here, cause now we set time for redeem unlock, not when cooldown was called
-    assertEq(cooldownBefore.timestamp, block.timestamp + stakeToken.getDefaultCooldownSeconds());
-    assertEq(cooldownBefore.amount, amountToStake);
-
-    vm.warp(block.timestamp + stakeToken.getDefaultCooldownSeconds());
-    _redeem(amountToRedeem, user, user);
-
-    IStakeToken.CooldownSetup memory cooldownAfter = stakeToken.stakersCooldowns(user);
-    assertEq(cooldownAfter.amount, amountToStake - amountToRedeem);
-  }
-
-  function test_cooldownNoIncreaseInAmount(
-    uint104 amountToStake,
-    uint104 amountToTopUp,
-    address user
-  ) public {
-    vm.assume(
-      amountToStake > 0 && amountToTopUp > 0 && type(uint104).max - amountToStake > amountToTopUp
-    );
-    vm.assume(user != address(proxyAdmin) && user != address(0));
-
-    _stake(amountToStake, user);
-    vm.startPrank(user);
-    stakeToken.cooldown();
-
-    IStakeToken.CooldownSetup memory cooldownBefore = stakeToken.stakersCooldowns(user);
-
-    // increase amount
-    _stake(amountToTopUp, user);
-
-    IStakeToken.CooldownSetup memory cooldownAfter = stakeToken.stakersCooldowns(user);
-    assertEq(cooldownBefore.timestamp, cooldownAfter.timestamp);
-    assertEq(cooldownBefore.amount, cooldownAfter.amount);
-    assertEq(cooldownAfter.timestamp, block.timestamp + stakeToken.getDefaultCooldownSeconds());
-    assertEq(cooldownAfter.amount, amountToStake);
-  }
-
-  function test_cooldownOnTransfer(
-    uint104 amountToStake,
-    uint104 amountToStakeOther,
-    address user,
-    address otherUser
-  ) public {
-    vm.assume(
-      amountToStake > 0 &&
-        amountToStakeOther > 0 &&
-        type(uint104).max - amountToStake > amountToStakeOther
-    );
-    vm.assume(user != address(proxyAdmin) && user != address(0) && user != otherUser);
-    vm.assume(otherUser != address(0) && otherUser != address(proxyAdmin));
-
-    _stake(amountToStake, user);
-    vm.startPrank(user);
-    stakeToken.cooldown();
-
-    IStakeToken.CooldownSetup memory cooldown0 = stakeToken.stakersCooldowns(user);
-
-    // Receiving token should not affect the amount
-    _stake(amountToStakeOther, otherUser);
-    vm.prank(otherUser);
-    stakeToken.transfer(user, amountToStakeOther);
-    IStakeToken.CooldownSetup memory cooldown1 = stakeToken.stakersCooldowns(user);
-
-    assertEq(cooldown0.timestamp, cooldown1.timestamp, 'MISMATCH_BEFORE_COOLDOWN');
-    assertEq(cooldown0.amount, cooldown1.amount, 'MISMATCH_BEFORE_COOLDOWN_AMOUNT');
-
-    // Sending token should not affect the amount as long as balance > amount
-    vm.prank(user);
-    stakeToken.transfer(otherUser, amountToStakeOther);
-    IStakeToken.CooldownSetup memory cooldown2 = stakeToken.stakersCooldowns(user);
-    assertEq(cooldown1.timestamp, cooldown2.timestamp, 'MISMATCH_COOLDOWN');
-    assertEq(cooldown0.amount, cooldown2.amount, 'MISMATCH_COOLDOWN_AMOUNT');
-
-    // Sending token should decrease the cooldown amount when balance <= amount
-    vm.startPrank(user);
-    stakeToken.transfer(otherUser, amountToStake);
-    vm.stopPrank();
-    IStakeToken.CooldownSetup memory cooldown3 = stakeToken.stakersCooldowns(user);
-    assertEq(cooldown3.timestamp, 0, 'MISMATCH_AFTER_COOLDOWN');
-    assertEq(cooldown3.amount, 0, 'MISMATCH_AFTER_COOLDOWN_AMOUNT');
-  }
-
-  function test_cooldownInsufficient_shouldRevert(
-    uint104 amountToStake,
-    uint104 amountToUnstake,
-    uint40 secondsAfterCooldownActivation,
-    address user,
-    address destination
-  ) public {
-    vm.assume(amountToUnstake != 0 && amountToStake >= amountToUnstake);
-    vm.assume(secondsAfterCooldownActivation < stakeToken.getDefaultCooldownSeconds());
-    vm.assume(user != address(proxyAdmin) && user != address(0) && destination != address(0));
-
-    _stake(amountToStake, user);
-    vm.prank(user);
-    stakeToken.cooldown();
-
-    vm.warp(block.timestamp + secondsAfterCooldownActivation);
-    vm.prank(user);
-    vm.expectRevert('INSUFFICIENT_COOLDOWN');
-    stakeToken.redeem(destination, amountToUnstake);
-  }
-
-  function test_cooldownWindowClosed_shouldRevert(
-    uint104 amountToStake,
-    uint104 amountToUnstake,
-    uint40 secondsAfterCooldownActivation,
-    address user,
-    address destination
-  ) public {
-    vm.assume(amountToUnstake != 0 && amountToStake >= amountToUnstake);
-    vm.assume(
-      secondsAfterCooldownActivation >
-        stakeToken.getDefaultCooldownSeconds() + stakeToken.getUnstakeWindow()
-    );
-    vm.assume(user != address(proxyAdmin) && user != address(0) && destination != address(0));
-
-    _stake(amountToStake, user);
-    vm.prank(user);
-    stakeToken.cooldown();
-
-    vm.warp(block.timestamp + secondsAfterCooldownActivation);
-    vm.prank(user);
-    vm.expectRevert('UNSTAKE_WINDOW_FINISHED');
-    stakeToken.redeem(destination, amountToUnstake);
-  }
-
-  function test_redeemMoreThenCooldown_shouldRedeemMax(
-    uint104 amountToStake,
-    uint104 amountToTopUp,
-    uint104 amountToUnstake,
-    address user,
-    address destination
-  ) public {
-    vm.assume(amountToStake != 0 && amountToUnstake >= amountToStake);
-    vm.assume(amountToTopUp != 0 && type(uint104).max - amountToTopUp >= amountToStake);
-    vm.assume(
-      user != address(proxyAdmin) &&
-        user != address(0) &&
-        user != address(stakeToken) &&
-        destination != address(0) &&
-        destination != address(stakeToken)
-    );
-
-    _stake(amountToStake, user);
-    vm.prank(user);
-    stakeToken.cooldown();
-    _stake(amountToTopUp, user);
-    IStakeToken.CooldownSetup memory cooldownAfterSecondStake = stakeToken.stakersCooldowns(user);
-    assertEq(cooldownAfterSecondStake.amount, amountToStake, 'STAKE_SHOULD_NOT_ALTER_COOLDOWN');
-
-    vm.warp(block.timestamp + stakeToken.getDefaultCooldownSeconds());
-    _redeem(amountToUnstake, user, destination);
-
-    assertEq(underlyingToken.balanceOf(destination), amountToStake, 'WRONG_AMOUNT_REDEEMED');
-    assertEq(stakeToken.balanceOf(user), amountToTopUp, 'WRONG_AMOUNT_LEFT');
-  }
+  // function test_cooldown(uint104 amountToStake, uint104 amountToRedeem, address user) public {
+  //   vm.assume(amountToStake >= amountToRedeem && amountToRedeem > 0);
+  //   vm.assume(user != address(proxyAdmin) && user != address(0));
+
+  //   _stake(amountToStake, user);
+
+  //   vm.startPrank(user);
+  //   stakeToken.cooldown();
+  //   IStakeToken.CooldownSnapshot memory cooldownBefore = stakeToken.stakersCooldowns(user);
+  //   // @pavelvm5 rewrite here, cause now we set time for redeem unlock, not when cooldown was called
+  //   assertEq(cooldownBefore.timestamp, block.timestamp + stakeToken.getDefaultCooldownSeconds());
+  //   assertEq(cooldownBefore.amount, amountToStake);
+
+  //   vm.warp(block.timestamp + stakeToken.getDefaultCooldownSeconds());
+  //   _redeem(amountToRedeem, user, user);
+
+  //   IStakeToken.CooldownSnapshot memory cooldownAfter = stakeToken.stakersCooldowns(user);
+  //   assertEq(cooldownAfter.amount, amountToStake - amountToRedeem);
+  // }
+
+  // function test_cooldownNoIncreaseInAmount(
+  //   uint104 amountToStake,
+  //   uint104 amountToTopUp,
+  //   address user
+  // ) public {
+  //   vm.assume(
+  //     amountToStake > 0 && amountToTopUp > 0 && type(uint104).max - amountToStake > amountToTopUp
+  //   );
+  //   vm.assume(user != address(proxyAdmin) && user != address(0));
+
+  //   _stake(amountToStake, user);
+  //   vm.startPrank(user);
+  //   stakeToken.cooldown();
+
+  //   IStakeToken.CooldownSnapshot memory cooldownBefore = stakeToken.stakersCooldowns(user);
+
+  //   // increase amount
+  //   _stake(amountToTopUp, user);
+
+  //   IStakeToken.CooldownSnapshot memory cooldownAfter = stakeToken.stakersCooldowns(user);
+  //   assertEq(cooldownBefore.timestamp, cooldownAfter.timestamp);
+  //   assertEq(cooldownBefore.amount, cooldownAfter.amount);
+  //   assertEq(cooldownAfter.timestamp, block.timestamp + stakeToken.getDefaultCooldownSeconds());
+  //   assertEq(cooldownAfter.amount, amountToStake);
+  // }
+
+  // function test_cooldownOnTransfer(
+  //   uint104 amountToStake,
+  //   uint104 amountToStakeOther,
+  //   address user,
+  //   address otherUser
+  // ) public {
+  //   vm.assume(
+  //     amountToStake > 0 &&
+  //       amountToStakeOther > 0 &&
+  //       type(uint104).max - amountToStake > amountToStakeOther
+  //   );
+  //   vm.assume(user != address(proxyAdmin) && user != address(0) && user != otherUser);
+  //   vm.assume(otherUser != address(0) && otherUser != address(proxyAdmin));
+
+  //   _stake(amountToStake, user);
+  //   vm.startPrank(user);
+  //   stakeToken.cooldown();
+
+  //   IStakeToken.CooldownSnapshot memory cooldown0 = stakeToken.stakersCooldowns(user);
+
+  //   // Receiving token should not affect the amount
+  //   _stake(amountToStakeOther, otherUser);
+  //   vm.prank(otherUser);
+  //   stakeToken.transfer(user, amountToStakeOther);
+  //   IStakeToken.CooldownSnapshot memory cooldown1 = stakeToken.stakersCooldowns(user);
+
+  //   assertEq(cooldown0.timestamp, cooldown1.timestamp, 'MISMATCH_BEFORE_COOLDOWN');
+  //   assertEq(cooldown0.amount, cooldown1.amount, 'MISMATCH_BEFORE_COOLDOWN_AMOUNT');
+
+  //   // Sending token should not affect the amount as long as balance > amount
+  //   vm.prank(user);
+  //   stakeToken.transfer(otherUser, amountToStakeOther);
+  //   IStakeToken.CooldownSnapshot memory cooldown2 = stakeToken.stakersCooldowns(user);
+  //   assertEq(cooldown1.timestamp, cooldown2.timestamp, 'MISMATCH_COOLDOWN');
+  //   assertEq(cooldown0.amount, cooldown2.amount, 'MISMATCH_COOLDOWN_AMOUNT');
+
+  //   // Sending token should decrease the cooldown amount when balance <= amount
+  //   vm.startPrank(user);
+  //   stakeToken.transfer(otherUser, amountToStake);
+  //   vm.stopPrank();
+  //   IStakeToken.CooldownSnapshot memory cooldown3 = stakeToken.stakersCooldowns(user);
+  //   assertEq(cooldown3.timestamp, 0, 'MISMATCH_AFTER_COOLDOWN');
+  //   assertEq(cooldown3.amount, 0, 'MISMATCH_AFTER_COOLDOWN_AMOUNT');
+  // }
+
+  // function test_cooldownInsufficient_shouldRevert(
+  //   uint104 amountToStake,
+  //   uint104 amountToUnstake,
+  //   uint40 secondsAfterCooldownActivation,
+  //   address user,
+  //   address destination
+  // ) public {
+  //   vm.assume(amountToUnstake != 0 && amountToStake >= amountToUnstake);
+  //   vm.assume(secondsAfterCooldownActivation < stakeToken.getDefaultCooldownSeconds());
+  //   vm.assume(user != address(proxyAdmin) && user != address(0) && destination != address(0));
+
+  //   _stake(amountToStake, user);
+  //   vm.prank(user);
+  //   stakeToken.cooldown();
+
+  //   vm.warp(block.timestamp + secondsAfterCooldownActivation);
+  //   vm.prank(user);
+  //   vm.expectRevert('INSUFFICIENT_COOLDOWN');
+  //   stakeToken.redeem(destination, amountToUnstake);
+  // }
+
+  // function test_cooldownWindowClosed_shouldRevert(
+  //   uint104 amountToStake,
+  //   uint104 amountToUnstake,
+  //   uint40 secondsAfterCooldownActivation,
+  //   address user,
+  //   address destination
+  // ) public {
+  //   vm.assume(amountToUnstake != 0 && amountToStake >= amountToUnstake);
+  //   vm.assume(
+  //     secondsAfterCooldownActivation >
+  //       stakeToken.getDefaultCooldownSeconds() + stakeToken.getUnstakeWindow()
+  //   );
+  //   vm.assume(user != address(proxyAdmin) && user != address(0) && destination != address(0));
+
+  //   _stake(amountToStake, user);
+  //   vm.prank(user);
+  //   stakeToken.cooldown();
+
+  //   vm.warp(block.timestamp + secondsAfterCooldownActivation);
+  //   vm.prank(user);
+  //   vm.expectRevert('UNSTAKE_WINDOW_FINISHED');
+  //   stakeToken.redeem(destination, amountToUnstake);
+  // }
+
+  // function test_redeemMoreThenCooldown_shouldRedeemMax(
+  //   uint104 amountToStake,
+  //   uint104 amountToTopUp,
+  //   uint104 amountToUnstake,
+  //   address user,
+  //   address destination
+  // ) public {
+  //   vm.assume(amountToStake != 0 && amountToUnstake >= amountToStake);
+  //   vm.assume(amountToTopUp != 0 && type(uint104).max - amountToTopUp >= amountToStake);
+  //   vm.assume(
+  //     user != address(proxyAdmin) &&
+  //       user != address(0) &&
+  //       user != address(stakeToken) &&
+  //       destination != address(0) &&
+  //       destination != address(stakeToken)
+  //   );
+
+  //   _stake(amountToStake, user);
+  //   vm.prank(user);
+  //   stakeToken.cooldown();
+  //   _stake(amountToTopUp, user);
+  //   IStakeToken.CooldownSnapshot memory cooldownAfterSecondStake = stakeToken.stakersCooldowns(
+  //     user
+  //   );
+  //   assertEq(cooldownAfterSecondStake.amount, amountToStake, 'STAKE_SHOULD_NOT_ALTER_COOLDOWN');
+
+  //   vm.warp(block.timestamp + stakeToken.getDefaultCooldownSeconds());
+  //   _redeem(amountToUnstake, user, destination);
+
+  //   assertEq(underlyingToken.balanceOf(destination), amountToStake, 'WRONG_AMOUNT_REDEEMED');
+  //   assertEq(stakeToken.balanceOf(user), amountToTopUp, 'WRONG_AMOUNT_LEFT');
+  // }
 
   // @pavelvm5 fast-withdrawal tests start here, check only impact using reducedCooldown function, redeem/transfer process is the same
   function test_reducedCooldown(
@@ -199,7 +201,7 @@ contract Cooldown is StkTestUtils {
     uint256 maxReductionSeconds = stakeToken.getMaxReductionSeconds();
     stakeToken.reducedCooldown(maxReductionSeconds);
 
-    IStakeToken.CooldownSetup memory cooldownBefore = stakeToken.stakersCooldowns(user);
+    IStakeToken.CooldownSnapshot memory cooldownBefore = stakeToken.stakersCooldowns(user);
 
     uint256 cooldownSeconds = stakeToken.getDefaultCooldownSeconds() - maxReductionSeconds;
 
@@ -210,12 +212,7 @@ contract Cooldown is StkTestUtils {
     vm.warp(block.timestamp + cooldownSeconds);
     _redeem(amountToRedeem, user, user);
 
-    IStakeToken.CooldownSetup memory cooldownAfter = stakeToken.stakersCooldowns(user);
-
-    console.log(cooldownAfter.amount);
-    console.log(amountToStake);
-    console.log(fee);
-    console.log(amountToRedeem);
+    IStakeToken.CooldownSnapshot memory cooldownAfter = stakeToken.stakersCooldowns(user);
 
     assertEq(cooldownAfter.amount, amountToStake - fee - amountToRedeem);
   }
@@ -277,18 +274,96 @@ contract Cooldown is StkTestUtils {
 
     assertEq(balanceTreasuryAfter - balanceTreasuryBefore, underlyingFeesToTreasury);
 
-    IStakeToken.CooldownSetup memory cooldownBefore = stakeToken.stakersCooldowns(user);
+    IStakeToken.CooldownSnapshot memory cooldownBefore = stakeToken.stakersCooldowns(user);
 
     uint256 cooldownSeconds = stakeToken.getDefaultCooldownSeconds() - reducedTime;
 
     assertEq(cooldownBefore.timestamp, block.timestamp + cooldownSeconds);
-
     assertEq(cooldownBefore.amount, amountToStake - fee);
 
     vm.warp(block.timestamp + cooldownSeconds);
     _redeem(amountToRedeem, user, user);
 
-    IStakeToken.CooldownSetup memory cooldownAfter = stakeToken.stakersCooldowns(user);
+    IStakeToken.CooldownSnapshot memory cooldownAfter = stakeToken.stakersCooldowns(user);
     assertEq(cooldownAfter.amount, amountToStake - fee - amountToRedeem);
   }
+
+  function test_reducedCooldownSecondTime(
+    uint104 amountToStake,
+    uint104 amountToRedeem,
+    address user,
+    uint32 reducedTime
+  ) public {
+    uint256 maxReductionSeconds = stakeToken.getMaxReductionSeconds();
+    uint256 fee = (amountToStake * maxFee * reducedTime) / maxReductionSeconds / 10_000;
+
+    vm.assume(amountToStake >= amountToRedeem + fee && amountToRedeem > 0);
+    vm.assume(user != address(proxyAdmin) && user != address(0) && user != treasury);
+    vm.assume(reducedTime <= maxReductionSeconds / 2);
+
+    _stake(amountToStake, user);
+
+    vm.startPrank(user);
+
+    stakeToken.cooldown();
+
+    vm.warp(block.timestamp + maxReductionSeconds / 2);
+
+    uint256 balanceTreasuryBefore = underlyingToken.balanceOf(treasury);
+    uint256 underlyingFeesToTreasury = stakeToken.previewRedeem(fee);
+
+    IStakeToken.CooldownSnapshot memory cooldownBefore = stakeToken.stakersCooldowns(user);
+    assertEq(cooldownBefore.amount, amountToStake);
+
+    stakeToken.reducedCooldown(reducedTime);
+
+    uint256 balanceTreasuryAfter = underlyingToken.balanceOf(treasury);
+    assertEq(balanceTreasuryAfter - balanceTreasuryBefore, underlyingFeesToTreasury);
+
+    IStakeToken.CooldownSnapshot memory cooldownAfter = stakeToken.stakersCooldowns(user);
+    assertEq(cooldownBefore.timestamp - reducedTime, cooldownAfter.timestamp);
+    assertEq(cooldownAfter.amount, amountToStake - fee);
+
+    vm.warp(stakeToken.stakersCooldowns(user).timestamp);
+    _redeem(amountToRedeem, user, user);
+
+    IStakeToken.CooldownSnapshot memory cooldownAfterRedeem = stakeToken.stakersCooldowns(user);
+    assertEq(cooldownAfterRedeem.amount, amountToStake - fee - amountToRedeem);
+  }
+
+  function test_reducedCooldownRevert(uint104 amountToStake, address user) public {
+    vm.assume(amountToStake > 0);
+    vm.assume(user != address(proxyAdmin) && user != address(0) && user != treasury);
+
+    _stake(amountToStake, user);
+
+    vm.startPrank(user);
+
+    stakeToken.reducedCooldown(stakeToken.getMaxReductionSeconds());
+
+    vm.expectRevert('MAX_REDUCTION_TIME_EXCEEDED');
+    stakeToken.reducedCooldown(1);
+  }
+
+  function test_reducedCooldownRevertWithWarp(
+    uint104 amountToStake,
+    address user,
+    uint40 reducedTime
+  ) public {
+    vm.assume(amountToStake > 0);
+    vm.assume(user != address(proxyAdmin) && user != address(0) && user != treasury);
+
+    _stake(amountToStake, user);
+
+    vm.startPrank(user);
+
+    stakeToken.reducedCooldown(reducedTime);
+
+    IStakeToken.CooldownSnapshot memory cooldownAfter = stakeToken.stakersCooldowns(user);
+
+    vm.warp(cooldownAfter.timestamp - minCooldownSeconds);
+
+    vm.expectRevert('MAX_REDUCTION_TIME_EXCEEDED');
+    stakeToken.reducedCooldown(1);
+  }
 }

From 8904d1af79f1bab6aa1bb5db022acd62e1d02fa0 Mon Sep 17 00:00:00 2001
From: Pavel Menshikov <menshikov.pv@mail.ru>
Date: Thu, 18 Jul 2024 03:24:24 +0500
Subject: [PATCH 5/5] removed safety check

---
 src/contracts/StakeToken.sol |   2 -
 tests/Cooldown.t.sol         | 336 +++++++++++++++++------------------
 2 files changed, 168 insertions(+), 170 deletions(-)

diff --git a/src/contracts/StakeToken.sol b/src/contracts/StakeToken.sol
index e529bf7..5ae918b 100644
--- a/src/contracts/StakeToken.sol
+++ b/src/contracts/StakeToken.sol
@@ -115,8 +115,6 @@ contract StakeToken is ERC20PermitUpgradeable, IStakeToken, Rescuable {
   }
 
   function _setUnstakeWindow(uint256 newUnstakeWindow) internal {
-    require(newUnstakeWindow >= 1 hours, 'TOO_LOW_UNSTAKE_WINDOW');
-
     _getStakeTokenStorage()._smConfig.unstakeWindowSeconds = newUnstakeWindow.toUint40();
 
     emit UnstakeWindowChanged(newUnstakeWindow);
diff --git a/tests/Cooldown.t.sol b/tests/Cooldown.t.sol
index a9f4f9b..b4d7473 100644
--- a/tests/Cooldown.t.sol
+++ b/tests/Cooldown.t.sol
@@ -14,174 +14,174 @@ contract Cooldown is StkTestUtils {
   /**
    * cooldown should activate for a given block.timestamp and cooldown the currently held amount
    */
-  // function test_cooldown(uint104 amountToStake, uint104 amountToRedeem, address user) public {
-  //   vm.assume(amountToStake >= amountToRedeem && amountToRedeem > 0);
-  //   vm.assume(user != address(proxyAdmin) && user != address(0));
-
-  //   _stake(amountToStake, user);
-
-  //   vm.startPrank(user);
-  //   stakeToken.cooldown();
-  //   IStakeToken.CooldownSnapshot memory cooldownBefore = stakeToken.stakersCooldowns(user);
-  //   // @pavelvm5 rewrite here, cause now we set time for redeem unlock, not when cooldown was called
-  //   assertEq(cooldownBefore.timestamp, block.timestamp + stakeToken.getDefaultCooldownSeconds());
-  //   assertEq(cooldownBefore.amount, amountToStake);
-
-  //   vm.warp(block.timestamp + stakeToken.getDefaultCooldownSeconds());
-  //   _redeem(amountToRedeem, user, user);
-
-  //   IStakeToken.CooldownSnapshot memory cooldownAfter = stakeToken.stakersCooldowns(user);
-  //   assertEq(cooldownAfter.amount, amountToStake - amountToRedeem);
-  // }
-
-  // function test_cooldownNoIncreaseInAmount(
-  //   uint104 amountToStake,
-  //   uint104 amountToTopUp,
-  //   address user
-  // ) public {
-  //   vm.assume(
-  //     amountToStake > 0 && amountToTopUp > 0 && type(uint104).max - amountToStake > amountToTopUp
-  //   );
-  //   vm.assume(user != address(proxyAdmin) && user != address(0));
-
-  //   _stake(amountToStake, user);
-  //   vm.startPrank(user);
-  //   stakeToken.cooldown();
-
-  //   IStakeToken.CooldownSnapshot memory cooldownBefore = stakeToken.stakersCooldowns(user);
-
-  //   // increase amount
-  //   _stake(amountToTopUp, user);
-
-  //   IStakeToken.CooldownSnapshot memory cooldownAfter = stakeToken.stakersCooldowns(user);
-  //   assertEq(cooldownBefore.timestamp, cooldownAfter.timestamp);
-  //   assertEq(cooldownBefore.amount, cooldownAfter.amount);
-  //   assertEq(cooldownAfter.timestamp, block.timestamp + stakeToken.getDefaultCooldownSeconds());
-  //   assertEq(cooldownAfter.amount, amountToStake);
-  // }
-
-  // function test_cooldownOnTransfer(
-  //   uint104 amountToStake,
-  //   uint104 amountToStakeOther,
-  //   address user,
-  //   address otherUser
-  // ) public {
-  //   vm.assume(
-  //     amountToStake > 0 &&
-  //       amountToStakeOther > 0 &&
-  //       type(uint104).max - amountToStake > amountToStakeOther
-  //   );
-  //   vm.assume(user != address(proxyAdmin) && user != address(0) && user != otherUser);
-  //   vm.assume(otherUser != address(0) && otherUser != address(proxyAdmin));
-
-  //   _stake(amountToStake, user);
-  //   vm.startPrank(user);
-  //   stakeToken.cooldown();
-
-  //   IStakeToken.CooldownSnapshot memory cooldown0 = stakeToken.stakersCooldowns(user);
-
-  //   // Receiving token should not affect the amount
-  //   _stake(amountToStakeOther, otherUser);
-  //   vm.prank(otherUser);
-  //   stakeToken.transfer(user, amountToStakeOther);
-  //   IStakeToken.CooldownSnapshot memory cooldown1 = stakeToken.stakersCooldowns(user);
-
-  //   assertEq(cooldown0.timestamp, cooldown1.timestamp, 'MISMATCH_BEFORE_COOLDOWN');
-  //   assertEq(cooldown0.amount, cooldown1.amount, 'MISMATCH_BEFORE_COOLDOWN_AMOUNT');
-
-  //   // Sending token should not affect the amount as long as balance > amount
-  //   vm.prank(user);
-  //   stakeToken.transfer(otherUser, amountToStakeOther);
-  //   IStakeToken.CooldownSnapshot memory cooldown2 = stakeToken.stakersCooldowns(user);
-  //   assertEq(cooldown1.timestamp, cooldown2.timestamp, 'MISMATCH_COOLDOWN');
-  //   assertEq(cooldown0.amount, cooldown2.amount, 'MISMATCH_COOLDOWN_AMOUNT');
-
-  //   // Sending token should decrease the cooldown amount when balance <= amount
-  //   vm.startPrank(user);
-  //   stakeToken.transfer(otherUser, amountToStake);
-  //   vm.stopPrank();
-  //   IStakeToken.CooldownSnapshot memory cooldown3 = stakeToken.stakersCooldowns(user);
-  //   assertEq(cooldown3.timestamp, 0, 'MISMATCH_AFTER_COOLDOWN');
-  //   assertEq(cooldown3.amount, 0, 'MISMATCH_AFTER_COOLDOWN_AMOUNT');
-  // }
-
-  // function test_cooldownInsufficient_shouldRevert(
-  //   uint104 amountToStake,
-  //   uint104 amountToUnstake,
-  //   uint40 secondsAfterCooldownActivation,
-  //   address user,
-  //   address destination
-  // ) public {
-  //   vm.assume(amountToUnstake != 0 && amountToStake >= amountToUnstake);
-  //   vm.assume(secondsAfterCooldownActivation < stakeToken.getDefaultCooldownSeconds());
-  //   vm.assume(user != address(proxyAdmin) && user != address(0) && destination != address(0));
-
-  //   _stake(amountToStake, user);
-  //   vm.prank(user);
-  //   stakeToken.cooldown();
-
-  //   vm.warp(block.timestamp + secondsAfterCooldownActivation);
-  //   vm.prank(user);
-  //   vm.expectRevert('INSUFFICIENT_COOLDOWN');
-  //   stakeToken.redeem(destination, amountToUnstake);
-  // }
-
-  // function test_cooldownWindowClosed_shouldRevert(
-  //   uint104 amountToStake,
-  //   uint104 amountToUnstake,
-  //   uint40 secondsAfterCooldownActivation,
-  //   address user,
-  //   address destination
-  // ) public {
-  //   vm.assume(amountToUnstake != 0 && amountToStake >= amountToUnstake);
-  //   vm.assume(
-  //     secondsAfterCooldownActivation >
-  //       stakeToken.getDefaultCooldownSeconds() + stakeToken.getUnstakeWindow()
-  //   );
-  //   vm.assume(user != address(proxyAdmin) && user != address(0) && destination != address(0));
-
-  //   _stake(amountToStake, user);
-  //   vm.prank(user);
-  //   stakeToken.cooldown();
-
-  //   vm.warp(block.timestamp + secondsAfterCooldownActivation);
-  //   vm.prank(user);
-  //   vm.expectRevert('UNSTAKE_WINDOW_FINISHED');
-  //   stakeToken.redeem(destination, amountToUnstake);
-  // }
-
-  // function test_redeemMoreThenCooldown_shouldRedeemMax(
-  //   uint104 amountToStake,
-  //   uint104 amountToTopUp,
-  //   uint104 amountToUnstake,
-  //   address user,
-  //   address destination
-  // ) public {
-  //   vm.assume(amountToStake != 0 && amountToUnstake >= amountToStake);
-  //   vm.assume(amountToTopUp != 0 && type(uint104).max - amountToTopUp >= amountToStake);
-  //   vm.assume(
-  //     user != address(proxyAdmin) &&
-  //       user != address(0) &&
-  //       user != address(stakeToken) &&
-  //       destination != address(0) &&
-  //       destination != address(stakeToken)
-  //   );
-
-  //   _stake(amountToStake, user);
-  //   vm.prank(user);
-  //   stakeToken.cooldown();
-  //   _stake(amountToTopUp, user);
-  //   IStakeToken.CooldownSnapshot memory cooldownAfterSecondStake = stakeToken.stakersCooldowns(
-  //     user
-  //   );
-  //   assertEq(cooldownAfterSecondStake.amount, amountToStake, 'STAKE_SHOULD_NOT_ALTER_COOLDOWN');
-
-  //   vm.warp(block.timestamp + stakeToken.getDefaultCooldownSeconds());
-  //   _redeem(amountToUnstake, user, destination);
-
-  //   assertEq(underlyingToken.balanceOf(destination), amountToStake, 'WRONG_AMOUNT_REDEEMED');
-  //   assertEq(stakeToken.balanceOf(user), amountToTopUp, 'WRONG_AMOUNT_LEFT');
-  // }
+  function test_cooldown(uint104 amountToStake, uint104 amountToRedeem, address user) public {
+    vm.assume(amountToStake >= amountToRedeem && amountToRedeem > 0);
+    vm.assume(user != address(proxyAdmin) && user != address(0));
+
+    _stake(amountToStake, user);
+
+    vm.startPrank(user);
+    stakeToken.cooldown();
+    IStakeToken.CooldownSnapshot memory cooldownBefore = stakeToken.stakersCooldowns(user);
+    // @pavelvm5 rewrite here, cause now we set time for redeem unlock, not when cooldown was called
+    assertEq(cooldownBefore.timestamp, block.timestamp + stakeToken.getDefaultCooldownSeconds());
+    assertEq(cooldownBefore.amount, amountToStake);
+
+    vm.warp(block.timestamp + stakeToken.getDefaultCooldownSeconds());
+    _redeem(amountToRedeem, user, user);
+
+    IStakeToken.CooldownSnapshot memory cooldownAfter = stakeToken.stakersCooldowns(user);
+    assertEq(cooldownAfter.amount, amountToStake - amountToRedeem);
+  }
+
+  function test_cooldownNoIncreaseInAmount(
+    uint104 amountToStake,
+    uint104 amountToTopUp,
+    address user
+  ) public {
+    vm.assume(
+      amountToStake > 0 && amountToTopUp > 0 && type(uint104).max - amountToStake > amountToTopUp
+    );
+    vm.assume(user != address(proxyAdmin) && user != address(0));
+
+    _stake(amountToStake, user);
+    vm.startPrank(user);
+    stakeToken.cooldown();
+
+    IStakeToken.CooldownSnapshot memory cooldownBefore = stakeToken.stakersCooldowns(user);
+
+    // increase amount
+    _stake(amountToTopUp, user);
+
+    IStakeToken.CooldownSnapshot memory cooldownAfter = stakeToken.stakersCooldowns(user);
+    assertEq(cooldownBefore.timestamp, cooldownAfter.timestamp);
+    assertEq(cooldownBefore.amount, cooldownAfter.amount);
+    assertEq(cooldownAfter.timestamp, block.timestamp + stakeToken.getDefaultCooldownSeconds());
+    assertEq(cooldownAfter.amount, amountToStake);
+  }
+
+  function test_cooldownOnTransfer(
+    uint104 amountToStake,
+    uint104 amountToStakeOther,
+    address user,
+    address otherUser
+  ) public {
+    vm.assume(
+      amountToStake > 0 &&
+        amountToStakeOther > 0 &&
+        type(uint104).max - amountToStake > amountToStakeOther
+    );
+    vm.assume(user != address(proxyAdmin) && user != address(0) && user != otherUser);
+    vm.assume(otherUser != address(0) && otherUser != address(proxyAdmin));
+
+    _stake(amountToStake, user);
+    vm.startPrank(user);
+    stakeToken.cooldown();
+
+    IStakeToken.CooldownSnapshot memory cooldown0 = stakeToken.stakersCooldowns(user);
+
+    // Receiving token should not affect the amount
+    _stake(amountToStakeOther, otherUser);
+    vm.prank(otherUser);
+    stakeToken.transfer(user, amountToStakeOther);
+    IStakeToken.CooldownSnapshot memory cooldown1 = stakeToken.stakersCooldowns(user);
+
+    assertEq(cooldown0.timestamp, cooldown1.timestamp, 'MISMATCH_BEFORE_COOLDOWN');
+    assertEq(cooldown0.amount, cooldown1.amount, 'MISMATCH_BEFORE_COOLDOWN_AMOUNT');
+
+    // Sending token should not affect the amount as long as balance > amount
+    vm.prank(user);
+    stakeToken.transfer(otherUser, amountToStakeOther);
+    IStakeToken.CooldownSnapshot memory cooldown2 = stakeToken.stakersCooldowns(user);
+    assertEq(cooldown1.timestamp, cooldown2.timestamp, 'MISMATCH_COOLDOWN');
+    assertEq(cooldown0.amount, cooldown2.amount, 'MISMATCH_COOLDOWN_AMOUNT');
+
+    // Sending token should decrease the cooldown amount when balance <= amount
+    vm.startPrank(user);
+    stakeToken.transfer(otherUser, amountToStake);
+    vm.stopPrank();
+    IStakeToken.CooldownSnapshot memory cooldown3 = stakeToken.stakersCooldowns(user);
+    assertEq(cooldown3.timestamp, 0, 'MISMATCH_AFTER_COOLDOWN');
+    assertEq(cooldown3.amount, 0, 'MISMATCH_AFTER_COOLDOWN_AMOUNT');
+  }
+
+  function test_cooldownInsufficient_shouldRevert(
+    uint104 amountToStake,
+    uint104 amountToUnstake,
+    uint40 secondsAfterCooldownActivation,
+    address user,
+    address destination
+  ) public {
+    vm.assume(amountToUnstake != 0 && amountToStake >= amountToUnstake);
+    vm.assume(secondsAfterCooldownActivation < stakeToken.getDefaultCooldownSeconds());
+    vm.assume(user != address(proxyAdmin) && user != address(0) && destination != address(0));
+
+    _stake(amountToStake, user);
+    vm.prank(user);
+    stakeToken.cooldown();
+
+    vm.warp(block.timestamp + secondsAfterCooldownActivation);
+    vm.prank(user);
+    vm.expectRevert('INSUFFICIENT_COOLDOWN');
+    stakeToken.redeem(destination, amountToUnstake);
+  }
+
+  function test_cooldownWindowClosed_shouldRevert(
+    uint104 amountToStake,
+    uint104 amountToUnstake,
+    uint40 secondsAfterCooldownActivation,
+    address user,
+    address destination
+  ) public {
+    vm.assume(amountToUnstake != 0 && amountToStake >= amountToUnstake);
+    vm.assume(
+      secondsAfterCooldownActivation >
+        stakeToken.getDefaultCooldownSeconds() + stakeToken.getUnstakeWindow()
+    );
+    vm.assume(user != address(proxyAdmin) && user != address(0) && destination != address(0));
+
+    _stake(amountToStake, user);
+    vm.prank(user);
+    stakeToken.cooldown();
+
+    vm.warp(block.timestamp + secondsAfterCooldownActivation);
+    vm.prank(user);
+    vm.expectRevert('UNSTAKE_WINDOW_FINISHED');
+    stakeToken.redeem(destination, amountToUnstake);
+  }
+
+  function test_redeemMoreThenCooldown_shouldRedeemMax(
+    uint104 amountToStake,
+    uint104 amountToTopUp,
+    uint104 amountToUnstake,
+    address user,
+    address destination
+  ) public {
+    vm.assume(amountToStake != 0 && amountToUnstake >= amountToStake);
+    vm.assume(amountToTopUp != 0 && type(uint104).max - amountToTopUp >= amountToStake);
+    vm.assume(
+      user != address(proxyAdmin) &&
+        user != address(0) &&
+        user != address(stakeToken) &&
+        destination != address(0) &&
+        destination != address(stakeToken)
+    );
+
+    _stake(amountToStake, user);
+    vm.prank(user);
+    stakeToken.cooldown();
+    _stake(amountToTopUp, user);
+    IStakeToken.CooldownSnapshot memory cooldownAfterSecondStake = stakeToken.stakersCooldowns(
+      user
+    );
+    assertEq(cooldownAfterSecondStake.amount, amountToStake, 'STAKE_SHOULD_NOT_ALTER_COOLDOWN');
+
+    vm.warp(block.timestamp + stakeToken.getDefaultCooldownSeconds());
+    _redeem(amountToUnstake, user, destination);
+
+    assertEq(underlyingToken.balanceOf(destination), amountToStake, 'WRONG_AMOUNT_REDEEMED');
+    assertEq(stakeToken.balanceOf(user), amountToTopUp, 'WRONG_AMOUNT_LEFT');
+  }
 
   // @pavelvm5 fast-withdrawal tests start here, check only impact using reducedCooldown function, redeem/transfer process is the same
   function test_reducedCooldown(