Skip to content

Commit

Permalink
moved to extension
Browse files Browse the repository at this point in the history
  • Loading branch information
Pavel Menshikov committed Aug 19, 2024
1 parent 58eb40a commit 3696658
Show file tree
Hide file tree
Showing 10 changed files with 211 additions and 147 deletions.
169 changes: 50 additions & 119 deletions src/contracts/StakeToken.sol
Original file line number Diff line number Diff line change
@@ -1,63 +1,48 @@
// SPDX-License-Identifier: BUSL-1.1
pragma solidity ^0.8.0;

import {UpgradeableOwnableWithGuardian} from 'solidity-utils/contracts/access-control/UpgradeableOwnableWithGuardian.sol';
import {IERC20} from 'openzeppelin-contracts/contracts/token/ERC20/IERC20.sol';
import {IERC20Permit} from 'openzeppelin-contracts/contracts/token/ERC20/extensions/IERC20Permit.sol';
import {IERC20Metadata} from 'openzeppelin-contracts/contracts/token/ERC20/extensions/IERC20Metadata.sol';

import {Initializable} from 'openzeppelin-contracts-upgradeable/contracts/proxy/utils/Initializable.sol';
import {ERC20PermitUpgradeable, ERC20Upgradeable, IERC20Permit} from 'openzeppelin-contracts-upgradeable/contracts/token/ERC20/extensions/ERC20PermitUpgradeable.sol';
import {ERC4626Upgradeable, IERC20Metadata, IERC20, Math, IERC4626} from 'openzeppelin-contracts-upgradeable/contracts/token/ERC20/extensions/ERC4626Upgradeable.sol';
import {ERC20PausableUpgradeable} from 'openzeppelin-contracts-upgradeable/contracts/token/ERC20/extensions/ERC20PausableUpgradeable.sol';
import {IERC4626} from 'openzeppelin-contracts/contracts/interfaces/IERC4626.sol';

import {IAccessControl} from 'openzeppelin-contracts/contracts/access/IAccessControl.sol';
import {IERC20 as SafeIERC20} from 'openzeppelin-contracts/contracts/token/ERC20/IERC20.sol';
import {SafeERC20} from 'openzeppelin-contracts/contracts/token/ERC20/utils/SafeERC20.sol';
import {SafeCast} from 'openzeppelin-contracts/contracts/utils/math/SafeCast.sol';

import {IPoolAddressesProvider} from './interfaces/IPoolAddressesProvider.sol';
import {IRewardsController} from './interfaces/IRewardsController.sol';

import {IStakeToken} from './interfaces/IStakeToken.sol';
import {Initializable} from 'openzeppelin-contracts-upgradeable/contracts/proxy/utils/Initializable.sol';
import {ERC20Upgradeable} from 'openzeppelin-contracts-upgradeable/contracts/token/ERC20/ERC20Upgradeable.sol';
import {ERC20PausableUpgradeable} from 'openzeppelin-contracts-upgradeable/contracts/token/ERC20/extensions/ERC20PausableUpgradeable.sol';
import {ERC20PermitUpgradeable} from 'openzeppelin-contracts-upgradeable/contracts/token/ERC20/extensions/ERC20PermitUpgradeable.sol';
import {ERC4626Upgradeable} from 'openzeppelin-contracts-upgradeable/contracts/token/ERC20/extensions/ERC4626Upgradeable.sol';

import {UpgradeableOwnableWithGuardian} from 'solidity-utils/contracts/access-control/UpgradeableOwnableWithGuardian.sol';

import {SafeERC20} from 'openzeppelin-contracts/contracts/token/ERC20/utils/SafeERC20.sol';
import {SafeCast} from 'openzeppelin-contracts/contracts/utils/math/SafeCast.sol';
import {Math} from 'openzeppelin-contracts/contracts/utils/math/Math.sol';

import {StakeTokenUpgradeable} from './extension/StakeTokenUpgradeable.sol';

contract StakeToken is
Initializable,
ERC20PermitUpgradeable,
ERC20PausableUpgradeable,
ERC20PermitUpgradeable,
ERC4626Upgradeable,
UpgradeableOwnableWithGuardian,
IStakeToken
StakeTokenUpgradeable
{
using SafeERC20 for SafeIERC20;
using SafeERC20 for IERC20;
using SafeCast for uint256;
using Math for uint256;

uint256 public constant MIN_ASSETS_REMAINING = 1e6;

uint216 public constant INITIAL_EXCHANGE_RATE = 1e18;
uint256 public constant EXCHANGE_RATE_UNIT = 1e18;

IRewardsController public immutable REWARDS_CONTROLLER;
IPoolAddressesProvider public immutable ADDRESSES_PROVIDER;

/// @custom:storage-location erc7201:aave.storage.StakeToken
struct StakeTokenStorage {
/// @notice User cooldown options
mapping(address => CooldownSnapshot) _stakerCooldown;
/// @notice Cooldown parameters
SmConfig _smConfig;
/// @notice Current exchangeRate of the stk
uint192 _currentExchangeRate;
}

// keccak256(abi.encode(uint256(keccak256("aave.storage.StakeToken")) - 1)) & ~bytes32(uint256(0xff))
bytes32 private constant StakeTokenStorageLocation =
0x570b5e9089e57b3d227cfcd747a97877e3c5f12150099d7b38848c6202ca0a00;

function _getStakeTokenStorage() internal pure returns (StakeTokenStorage storage $) {
assembly {
$.slot := StakeTokenStorageLocation
}
}

modifier onlySlashingAdmin() {
if (
!IAccessControl(ADDRESSES_PROVIDER.getACLManager()).hasRole('SLASHING_ADMIN', _msgSender())
Expand All @@ -84,16 +69,15 @@ contract StakeToken is
uint256 unstakeWindow_
) external virtual initializer {
__ERC20_init(name, symbol);
__ERC20Permit_init(name);
__ERC20Pausable_init();
__ERC20Permit_init(name);

__ERC4626_init(stakedToken);

__Ownable_init(owner);
__Ownable_With_Guardian_init(guardian);

_setCooldown(cooldown_);
_setUnstakeWindow(unstakeWindow_);
_updateExchangeRate(INITIAL_EXCHANGE_RATE);
__StakeTokenUpgradable_init(cooldown_.toUint32(), unstakeWindow_.toUint32());
}

function slash(
Expand All @@ -117,9 +101,9 @@ contract StakeToken is
uint256 currentShares = totalSupply();
uint256 balance = convertToAssets(currentShares);

_updateExchangeRate(_getExchangeRate(balance - amount, currentShares));
_updateExchangeRate(_getExchangeRate(balance - amount, currentShares).toUint192());

SafeIERC20(asset()).safeTransfer(destination, amount);
IERC20(asset()).safeTransfer(destination, amount);

emit Slashed(destination, amount);

Expand All @@ -143,7 +127,7 @@ contract StakeToken is
sig.s
)
{} catch {
if (IERC20Metadata(asset()).allowance(msg.sender, address(this)) < assets) {
if (IERC20(asset()).allowance(msg.sender, address(this)) < assets) {
revert PermitNotSucceded();
}
}
Expand All @@ -164,35 +148,19 @@ contract StakeToken is
}

function setUnstakeWindow(uint256 newUnstakeWindow) external onlyOwner {
_setUnstakeWindow(newUnstakeWindow);
_setUnstakeWindow(newUnstakeWindow.toUint32());
}

function setCooldown(uint256 newCooldown) external onlyOwner {
_setCooldown(newCooldown);
}

function setPause(bool pause) external onlyOwnerOrGuardian {
if (pause) {
_pause();
} else {
_unpause();
}
}

function getExchangeRate() external view returns (uint256) {
return _getStakeTokenStorage()._currentExchangeRate;
}

function getCooldown() external view returns (uint256) {
return _getStakeTokenStorage()._smConfig.cooldown;
_setCooldown(newCooldown.toUint32());
}

function getUnstakeWindow() external view returns (uint256) {
return _getStakeTokenStorage()._smConfig.unstakeWindow;
function pause() external onlyOwnerOrGuardian {
_pause();
}

function getStakerCooldown(address user) external view returns (CooldownSnapshot memory) {
return _getStakeTokenStorage()._stakerCooldown[user];
function unpause() external onlyOwnerOrGuardian {
_unpause();
}

function maxWithdraw(
Expand All @@ -204,14 +172,14 @@ contract StakeToken is
function maxRedeem(
address owner
) public view override(ERC4626Upgradeable, IERC4626) returns (uint256) {
StakeTokenStorage storage $ = _getStakeTokenStorage();
CooldownSnapshot memory cooldownSnapshot = $._stakerCooldown[owner];
SmConfig memory smConfig = getSmConfig();
CooldownSnapshot memory cooldownSnapshot = getStakerCooldown(owner);

if (
block.timestamp >= cooldownSnapshot.timestamp &&
block.timestamp - cooldownSnapshot.timestamp <= $._smConfig.unstakeWindow
block.timestamp - cooldownSnapshot.timestamp <= smConfig.unstakeWindow
) {
return $._stakerCooldown[owner].amount;
return cooldownSnapshot.amount;
}

return 0;
Expand All @@ -228,7 +196,7 @@ contract StakeToken is
override(ERC20Upgradeable, ERC4626Upgradeable, IERC20Metadata)
returns (uint8)
{
return ERC4626Upgradeable.decimals();
return super.decimals();
}

function _cooldown(address from) internal whenNotPaused {
Expand All @@ -238,38 +206,11 @@ contract StakeToken is
revert ZeroBalanceInStaking();
}

StakeTokenStorage storage $ = _getStakeTokenStorage();
SmConfig memory smConfig = getSmConfig();

uint32 timeForRedemption = (block.timestamp + $._smConfig.cooldown).toUint32();
uint32 timeToUnlock = (block.timestamp + smConfig.cooldown).toUint32();

$._stakerCooldown[from] = CooldownSnapshot({
timestamp: timeForRedemption,
amount: amount.toUint224()
});

emit Cooldown(from, amount, timeForRedemption);
}

function _setUnstakeWindow(uint256 newUnstakeWindow) internal {
_getStakeTokenStorage()._smConfig.unstakeWindow = newUnstakeWindow.toUint32();

emit UnstakeWindowChanged(newUnstakeWindow);
}

function _setCooldown(uint256 newCooldown) internal {
_getStakeTokenStorage()._smConfig.cooldown = newCooldown.toUint32();

emit CooldownChanged(newCooldown);
}

function _updateExchangeRate(uint256 newExchangeRate) internal {
if (newExchangeRate == 0) {
revert ZeroExchangeRate();
}

_getStakeTokenStorage()._currentExchangeRate = newExchangeRate.toUint192();

emit ExchangeRateChanged(newExchangeRate);
_setStakerCooldown(from, amount.toUint224(), timeToUnlock);
}

function _update(
Expand All @@ -289,25 +230,24 @@ contract StakeToken is
uint256 balanceOfFrom = balanceOf(from);
REWARDS_CONTROLLER.handleAction(from, cachedTotalSupply, balanceOfFrom);

StakeTokenStorage storage $ = _getStakeTokenStorage();
CooldownSnapshot memory previousSenderCooldown = $._stakerCooldown[from];
CooldownSnapshot memory cooldownSnapshot = getStakerCooldown(from);

if (previousSenderCooldown.timestamp != 0) {
if (cooldownSnapshot.timestamp != 0) {
if (to == address(0)) {
// redeem
if (previousSenderCooldown.amount == value) {
delete $._stakerCooldown[from];
if (cooldownSnapshot.amount == value) {
_deleteStakerCooldown(from);
} else {
$._stakerCooldown[from].amount = (previousSenderCooldown.amount - value).toUint224();
_setStakerCooldownAmount(from, (cooldownSnapshot.amount - value.toUint224()));
}
} else {
// transfer
uint224 balanceAfter = (balanceOfFrom - value).toUint224();

if (balanceAfter == 0) {
delete $._stakerCooldown[from];
} else if (balanceAfter < previousSenderCooldown.amount) {
$._stakerCooldown[from].amount = balanceAfter;
_deleteStakerCooldown(from);
} else if (balanceAfter < cooldownSnapshot.amount) {
_setStakerCooldownAmount(from, balanceAfter);
}
}
}
Expand All @@ -320,22 +260,13 @@ contract StakeToken is
uint256 assets,
Math.Rounding rounding
) internal view override returns (uint256) {
StakeTokenStorage storage $ = _getStakeTokenStorage();
return assets.mulDiv($._currentExchangeRate, EXCHANGE_RATE_UNIT, rounding);
return assets.mulDiv(getExchangeRate(), EXCHANGE_RATE_UNIT, rounding);
}

function _convertToAssets(
uint256 shares,
Math.Rounding rounding
) internal view override returns (uint256) {
StakeTokenStorage storage $ = _getStakeTokenStorage();
return shares.mulDiv(EXCHANGE_RATE_UNIT, $._currentExchangeRate, rounding);
}

function _getExchangeRate(
uint256 newTotalAssets,
uint256 newTotalShares
) internal pure returns (uint256) {
return newTotalShares.mulDiv(EXCHANGE_RATE_UNIT, newTotalAssets, Math.Rounding.Ceil);
return shares.mulDiv(EXCHANGE_RATE_UNIT, getExchangeRate(), rounding);
}
}
Loading

0 comments on commit 3696658

Please sign in to comment.