diff --git a/.gas-report b/.gas-report index 46597ab..11b9e23 100644 --- a/.gas-report +++ b/.gas-report @@ -1,30 +1,3 @@ -| lib/openzeppelin-contracts/contracts/proxy/ERC1967/ERC1967Proxy.sol:ERC1967Proxy contract | | | | | | -|-------------------------------------------------------------------------------------------|-----------------|-------|--------|-------|---------| -| Deployment Cost | Deployment Size | | | | | -| 259350 | 1159 | | | | | -| Function Name | min | avg | median | max | # calls | -| MAX_LOCKUP_PERIOD | 644 | 1426 | 644 | 5144 | 23 | -| MAX_MULTIPLIER | 667 | 1567 | 667 | 5167 | 30 | -| MIN_LOCKUP_PERIOD | 646 | 3918 | 5146 | 5146 | 11 | -| MP_RATE_PER_YEAR | 602 | 602 | 602 | 602 | 3 | -| SCALE_FACTOR | 600 | 600 | 600 | 600 | 41 | -| STAKING_TOKEN | 7253 | 7253 | 7253 | 7253 | 182 | -| accountedRewards | 722 | 1289 | 722 | 2722 | 74 | -| emergencyModeEnabled | 7270 | 7270 | 7270 | 7270 | 7 | -| enableEmergencyMode | 28422 | 45323 | 50607 | 50607 | 8 | -| getAccount | 2020 | 2020 | 2020 | 2020 | 71 | -| getStakedBalance | 7464 | 7464 | 7464 | 7464 | 1 | -| isTrustedCodehash | 938 | 1454 | 938 | 2938 | 182 | -| rewardIndex | 766 | 793 | 766 | 2766 | 74 | -| setTrustedCodehash | 52794 | 52794 | 52794 | 52794 | 47 | -| totalMP | 723 | 723 | 723 | 723 | 77 | -| totalMaxMP | 722 | 722 | 722 | 722 | 77 | -| totalStaked | 723 | 723 | 723 | 723 | 78 | -| updateAccountMP | 41679 | 43917 | 44181 | 44181 | 19 | -| updateGlobalState | 36999 | 67243 | 56463 | 89414 | 28 | -| upgradeToAndCall | 29734 | 33540 | 33540 | 37346 | 2 | - - | src/RewardsStreamer.sol:RewardsStreamer contract | | | | | | |--------------------------------------------------|-----------------|--------|--------|--------|---------| | Deployment Cost | Deployment Size | | | | | @@ -42,46 +15,77 @@ | src/RewardsStreamerMP.sol:RewardsStreamerMP contract | | | | | | |------------------------------------------------------|-----------------|--------|--------|--------|---------| | Deployment Cost | Deployment Size | | | | | -| 2084418 | 9588 | | | | | +| 2134697 | 9822 | | | | | | Function Name | min | avg | median | max | # calls | | MAX_LOCKUP_PERIOD | 272 | 272 | 272 | 272 | 23 | -| MAX_MULTIPLIER | 295 | 295 | 295 | 295 | 30 | +| MAX_MULTIPLIER | 273 | 273 | 273 | 273 | 30 | | MIN_LOCKUP_PERIOD | 274 | 274 | 274 | 274 | 11 | | MP_RATE_PER_YEAR | 230 | 230 | 230 | 230 | 3 | -| SCALE_FACTOR | 228 | 228 | 228 | 228 | 41 | -| STAKING_TOKEN | 2381 | 2381 | 2381 | 2381 | 182 | -| accountedRewards | 350 | 917 | 350 | 2350 | 74 | +| SCALE_FACTOR | 294 | 294 | 294 | 294 | 41 | +| STAKING_TOKEN | 2381 | 2381 | 2381 | 2381 | 194 | +| accountedRewards | 373 | 951 | 373 | 2373 | 76 | | emergencyModeEnabled | 2398 | 2398 | 2398 | 2398 | 7 | -| enableEmergencyMode | 2482 | 19389 | 24674 | 24674 | 8 | -| getAccount | 1621 | 1621 | 1621 | 1621 | 71 | -| getStakedBalance | 2589 | 2589 | 2589 | 2589 | 1 | -| initialize | 137795 | 137795 | 137795 | 137795 | 47 | -| isTrustedCodehash | 563 | 1079 | 563 | 2563 | 182 | +| enableEmergencyMode | 2460 | 19367 | 24652 | 24652 | 8 | +| getAccount | 1621 | 1621 | 1621 | 1621 | 72 | +| getStakedBalance | 2567 | 2567 | 2567 | 2567 | 1 | +| initialize | 137773 | 137773 | 137773 | 137773 | 50 | +| isTrustedCodehash | 563 | 1078 | 563 | 2563 | 194 | +| leave | 62061 | 62061 | 62061 | 62061 | 1 | | lock | 9839 | 33584 | 14168 | 76747 | 3 | -| proxiableUUID | 297 | 297 | 297 | 297 | 1 | -| rewardIndex | 394 | 421 | 394 | 2394 | 74 | -| setTrustedCodehash | 26203 | 26203 | 26203 | 26203 | 47 | -| stake | 134612 | 171004 | 176214 | 196693 | 57 | -| totalMP | 351 | 351 | 351 | 351 | 77 | -| totalMaxMP | 350 | 350 | 350 | 350 | 77 | -| totalStaked | 351 | 351 | 351 | 351 | 78 | -| unstake | 63925 | 84721 | 63925 | 120391 | 13 | -| updateAccountMP | 15375 | 17613 | 17877 | 17877 | 19 | +| proxiableUUID | 363 | 363 | 363 | 363 | 3 | +| rewardIndex | 394 | 420 | 394 | 2394 | 76 | +| setTrustedCodehash | 26226 | 26226 | 26226 | 26226 | 50 | +| stake | 134677 | 171329 | 176279 | 196758 | 60 | +| totalMP | 329 | 329 | 329 | 329 | 79 | +| totalMaxMP | 373 | 373 | 373 | 373 | 79 | +| totalStaked | 329 | 329 | 329 | 329 | 80 | +| unstake | 63997 | 84793 | 63997 | 120464 | 13 | +| updateAccountMP | 15353 | 17591 | 17855 | 17855 | 19 | | updateGlobalState | 11066 | 41310 | 30530 | 63481 | 28 | -| upgradeToAndCall | 3146 | 6955 | 6955 | 10765 | 2 | +| upgradeToAndCall | 3124 | 8887 | 10809 | 10809 | 4 | + + +| src/StakeManagerProxy.sol:StakeManagerProxy contract | | | | | | +|------------------------------------------------------|-----------------|-------|--------|-------|---------| +| Deployment Cost | Deployment Size | | | | | +| 278629 | 1263 | | | | | +| Function Name | min | avg | median | max | # calls | +| MAX_LOCKUP_PERIOD | 699 | 1481 | 699 | 5199 | 23 | +| MAX_MULTIPLIER | 700 | 1600 | 700 | 5200 | 30 | +| MIN_LOCKUP_PERIOD | 701 | 3973 | 5201 | 5201 | 11 | +| MP_RATE_PER_YEAR | 657 | 657 | 657 | 657 | 3 | +| SCALE_FACTOR | 721 | 721 | 721 | 721 | 41 | +| STAKING_TOKEN | 7308 | 7308 | 7308 | 7308 | 194 | +| accountedRewards | 800 | 1378 | 800 | 2800 | 76 | +| emergencyModeEnabled | 7325 | 7325 | 7325 | 7325 | 7 | +| enableEmergencyMode | 28455 | 45356 | 50640 | 50640 | 8 | +| getAccount | 2075 | 2075 | 2075 | 2075 | 72 | +| getStakedBalance | 7497 | 7497 | 7497 | 7497 | 1 | +| implementation | 343 | 935 | 343 | 2343 | 277 | +| isTrustedCodehash | 993 | 1508 | 993 | 2993 | 194 | +| rewardIndex | 821 | 847 | 821 | 2821 | 76 | +| setTrustedCodehash | 52872 | 52872 | 52872 | 52872 | 50 | +| totalMP | 756 | 756 | 756 | 756 | 79 | +| totalMaxMP | 800 | 800 | 800 | 800 | 79 | +| totalStaked | 756 | 756 | 756 | 756 | 80 | +| updateAccountMP | 41712 | 43950 | 44214 | 44214 | 19 | +| updateGlobalState | 37054 | 67298 | 56518 | 89469 | 28 | +| upgradeToAndCall | 29767 | 35525 | 37445 | 37445 | 4 | | src/StakeVault.sol:StakeVault contract | | | | | | |----------------------------------------|-----------------|--------|--------|--------|---------| | Deployment Cost | Deployment Size | | | | | -| 1086958 | 5110 | | | | | +| 1374543 | 6483 | | | | | | Function Name | min | avg | median | max | # calls | | STAKING_TOKEN | 216 | 216 | 216 | 216 | 1 | -| emergencyExit | 36298 | 48802 | 48036 | 65136 | 7 | -| lock | 43256 | 66100 | 47633 | 107411 | 3 | -| stake | 206077 | 242480 | 247679 | 268206 | 57 | -| unstake | 90514 | 120704 | 110342 | 153215 | 13 | -| withdraw | 42194 | 42194 | 42194 | 42194 | 1 | +| emergencyExit | 36353 | 48857 | 48091 | 65191 | 7 | +| leave | 33507 | 53144 | 33507 | 92418 | 3 | +| lock | 33245 | 60236 | 48577 | 110544 | 4 | +| stake | 33454 | 242405 | 250826 | 271353 | 61 | +| trustStakeManager | 28997 | 28997 | 28997 | 28997 | 1 | +| unstake | 33260 | 117303 | 113502 | 156376 | 14 | +| withdraw | 42227 | 42227 | 42227 | 42227 | 1 | | src/XPNFTToken.sol:XPNFTToken contract | | | | | | @@ -157,9 +161,9 @@ | Deployment Cost | Deployment Size | | | | | | 625454 | 3260 | | | | | | Function Name | min | avg | median | max | # calls | -| approve | 46330 | 46339 | 46342 | 46342 | 232 | -| balanceOf | 558 | 1380 | 558 | 2558 | 338 | -| mint | 51279 | 58862 | 51279 | 68379 | 248 | +| approve | 46330 | 46339 | 46342 | 46342 | 247 | +| balanceOf | 558 | 1387 | 558 | 2558 | 347 | +| mint | 51279 | 58819 | 51279 | 68379 | 263 | | transfer | 34384 | 48853 | 51484 | 51484 | 13 | diff --git a/.gas-snapshot b/.gas-snapshot index c18236b..b703a12 100644 --- a/.gas-snapshot +++ b/.gas-snapshot @@ -1,15 +1,18 @@ -EmergencyExitTest:test_CannotEnableEmergencyModeTwice() (gas: 89646) -EmergencyExitTest:test_CannotLeaveBeforeEmergencyMode() (gas: 292881) -EmergencyExitTest:test_EmergencyExitBasic() (gas: 395173) -EmergencyExitTest:test_EmergencyExitMultipleUsers() (gas: 823179) -EmergencyExitTest:test_EmergencyExitToAlternateAddress() (gas: 538257) -EmergencyExitTest:test_EmergencyExitWithLock() (gas: 384549) -EmergencyExitTest:test_EmergencyExitWithRewards() (gas: 529680) -EmergencyExitTest:test_OnlyOwnerCanEnableEmergencyMode() (gas: 39344) -IntegrationTest:testStakeFoo() (gas: 1559037) -LockTest:test_LockFailsWithInvalidPeriod() (gas: 299830) -LockTest:test_LockFailsWithNoStake() (gas: 58241) -LockTest:test_LockWithoutPriorLock() (gas: 396542) +EmergencyExitTest:test_CannotEnableEmergencyModeTwice() (gas: 89712) +EmergencyExitTest:test_CannotLeaveBeforeEmergencyMode() (gas: 296084) +EmergencyExitTest:test_EmergencyExitBasic() (gas: 398741) +EmergencyExitTest:test_EmergencyExitMultipleUsers() (gas: 830337) +EmergencyExitTest:test_EmergencyExitToAlternateAddress() (gas: 541602) +EmergencyExitTest:test_EmergencyExitWithLock() (gas: 387785) +EmergencyExitTest:test_EmergencyExitWithRewards() (gas: 533248) +EmergencyExitTest:test_OnlyOwnerCanEnableEmergencyMode() (gas: 39377) +IntegrationTest:testStakeFoo() (gas: 1578393) +LeaveTest:test_LeaveShouldProperlyUpdateAccounting() (gas: 2592444) +LeaveTest:test_RevertWhenStakeManagerIsTrusted() (gas: 293227) +LeaveTest:test_TrustNewStakeManager() (gas: 2642695) +LockTest:test_LockFailsWithInvalidPeriod() (gas: 306175) +LockTest:test_LockFailsWithNoStake() (gas: 61440) +LockTest:test_LockWithoutPriorLock() (gas: 403273) NFTMetadataGeneratorSVGTest:testGenerateMetadata() (gas: 85934) NFTMetadataGeneratorSVGTest:testSetImageStrings() (gas: 58332) NFTMetadataGeneratorSVGTest:testSetImageStringsRevert() (gas: 35804) @@ -17,41 +20,41 @@ NFTMetadataGeneratorURLTest:testGenerateMetadata() (gas: 102512) NFTMetadataGeneratorURLTest:testSetBaseURL() (gas: 49555) NFTMetadataGeneratorURLTest:testSetBaseURLRevert() (gas: 35979) RewardsStreamerTest:testStake() (gas: 869181) -StakeTest:test_StakeMultipleAccounts() (gas: 506030) -StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 661839) -StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 872090) -StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 520069) -StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 541881) -StakeTest:test_StakeOneAccount() (gas: 290297) -StakeTest:test_StakeOneAccountAndRewards() (gas: 446101) -StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 536165) -StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 532401) -StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 310643) -StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 310610) -StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 310721) +StakeTest:test_StakeMultipleAccounts() (gas: 512711) +StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 668852) +StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 879853) +StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 527355) +StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 549222) +StakeTest:test_StakeOneAccount() (gas: 293776) +StakeTest:test_StakeOneAccountAndRewards() (gas: 449912) +StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 540462) +StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 536907) +StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 314452) +StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 314419) +StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 314530) StakingTokenTest:testStakeToken() (gas: 10422) -UnstakeTest:test_StakeMultipleAccounts() (gas: 506074) -UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 661883) -UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 872089) -UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 520068) -UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 541925) -UnstakeTest:test_StakeOneAccount() (gas: 290320) -UnstakeTest:test_StakeOneAccountAndRewards() (gas: 446145) -UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 536209) -UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 532403) -UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 310600) -UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 310610) -UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 310721) -UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 541859) -UnstakeTest:test_UnstakeMultipleAccounts() (gas: 714931) -UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 1057383) -UnstakeTest:test_UnstakeOneAccount() (gas: 499845) -UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 524183) -UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 608646) -UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 555708) -UpgradeTest:test_RevertWhenNotOwner() (gas: 2159991) -UpgradeTest:test_UpgradeStakeManager() (gas: 2445373) -WithdrawTest:test_CannotWithdrawStakedFunds() (gas: 305631) +UnstakeTest:test_StakeMultipleAccounts() (gas: 512733) +UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 668874) +UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 879830) +UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 527332) +UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 549244) +UnstakeTest:test_StakeOneAccount() (gas: 293799) +UnstakeTest:test_StakeOneAccountAndRewards() (gas: 449934) +UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 540484) +UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 536887) +UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 314452) +UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 314419) +UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 314508) +UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 549390) +UnstakeTest:test_UnstakeMultipleAccounts() (gas: 728319) +UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 1073610) +UnstakeTest:test_UnstakeOneAccount() (gas: 508967) +UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 531442) +UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 615950) +UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 564034) +UpgradeTest:test_RevertWhenNotOwner() (gas: 2210367) +UpgradeTest:test_UpgradeStakeManager() (gas: 2499517) +WithdrawTest:test_CannotWithdrawStakedFunds() (gas: 308844) XPNFTTokenTest:testApproveNotAllowed() (gas: 10500) XPNFTTokenTest:testGetApproved() (gas: 10523) XPNFTTokenTest:testIsApprovedForAll() (gas: 10698) diff --git a/certora/specs/EmergencyMode.spec b/certora/specs/EmergencyMode.spec index 4254ddb..2caf88c 100644 --- a/certora/specs/EmergencyMode.spec +++ b/certora/specs/EmergencyMode.spec @@ -58,7 +58,8 @@ rule accountCanOnlyLeaveInEmergencyMode(method f) { f@withrevert(e, args); bool isReverted = lastReverted; - assert !isReverted => isViewFunction(f) || + assert !isReverted => f.selector == sig:streamer.leave().selector || + isViewFunction(f) || isOwnableFunction(f) || isTrustedCodehashAccessFunction(f) || isInitializerFunction(f) || diff --git a/certora/specs/RewardsStreamerMP.spec b/certora/specs/RewardsStreamerMP.spec index ff19673..1c299ec 100644 --- a/certora/specs/RewardsStreamerMP.spec +++ b/certora/specs/RewardsStreamerMP.spec @@ -12,6 +12,7 @@ methods { function updateGlobalState() external; function updateAccountMP(address accountAddress) external; function emergencyModeEnabled() external returns (bool) envfree; + function leave() external; } ghost mathint sumOfBalances { diff --git a/src/RewardsStreamerMP.sol b/src/RewardsStreamerMP.sol index 71b4497..d6da7dd 100644 --- a/src/RewardsStreamerMP.sol +++ b/src/RewardsStreamerMP.sol @@ -164,13 +164,16 @@ contract RewardsStreamerMP is UUPSUpgradeable, IStakeManager, TrustedCodehashAcc if (block.timestamp < account.lockUntil) { revert StakingManager__TokensAreLocked(); } + _unstake(amount, account, msg.sender); + } + function _unstake(uint256 amount, Account storage account, address accountAddress) internal { _updateGlobalState(); - _updateAccountMP(msg.sender); + _updateAccountMP(accountAddress); - uint256 accountRewards = calculateAccountRewards(msg.sender); + uint256 accountRewards = calculateAccountRewards(accountAddress); if (accountRewards > 0) { - distributeRewards(msg.sender, accountRewards); + distributeRewards(accountAddress, accountRewards); } uint256 previousStakedBalance = account.stakedBalance; @@ -181,11 +184,28 @@ contract RewardsStreamerMP is UUPSUpgradeable, IStakeManager, TrustedCodehashAcc account.stakedBalance -= amount; account.accountMP -= mpToReduce; account.maxMP -= maxMPToReduce; + account.accountRewardIndex = rewardIndex; totalMP -= mpToReduce; totalMaxMP -= maxMPToReduce; totalStaked -= amount; + } - account.accountRewardIndex = rewardIndex; + // @notice Allows an account to leave the system. This can happen when a + // user doesn't agree with an upgrade of the stake manager. + // @dev This function is protected by whitelisting the codehash of the caller. + // This ensures `StakeVault`s will call this function only if they don't + // trust the `StakeManager` (e.g. in case of an upgrade). + function leave() external onlyTrustedCodehash nonReentrant { + Account storage account = accounts[msg.sender]; + + if (account.stakedBalance > 0) { + // calling `_unstake` to update accounting accordingly + _unstake(account.stakedBalance, account, msg.sender); + + // further cleanup that isn't done in `_unstake` + account.accountRewardIndex = 0; + account.lockUntil = 0; + } } function _updateGlobalState() internal { diff --git a/src/StakeManagerProxy.sol b/src/StakeManagerProxy.sol new file mode 100644 index 0000000..6536c3a --- /dev/null +++ b/src/StakeManagerProxy.sol @@ -0,0 +1,12 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.26; + +import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; + +contract StakeManagerProxy is ERC1967Proxy { + constructor(address _implementation, bytes memory _data) ERC1967Proxy(_implementation, _data) { } + + function implementation() external view returns (address) { + return _implementation(); + } +} diff --git a/src/StakeVault.sol b/src/StakeVault.sol index 67d462b..7ab2581 100644 --- a/src/StakeVault.sol +++ b/src/StakeVault.sol @@ -4,7 +4,7 @@ pragma solidity ^0.8.26; import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol"; import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; -import { IStakeManager } from "./interfaces/IStakeManager.sol"; +import { IStakeManagerProxy } from "./interfaces/IStakeManagerProxy.sol"; /** * @title StakeVault @@ -19,12 +19,15 @@ contract StakeVault is Ownable { error StakeVault__StakingFailed(); error StakeVault__UnstakingFailed(); error StakeVault__NotAllowedToExit(); + error StakeVault__NotAllowedToLeave(); + error StakeVault__StakeManagerImplementationNotTrusted(); //STAKING_TOKEN must be kept as an immutable, otherwise, StakeManager would accept StakeVaults with any token //if is needed that STAKING_TOKEN to be a variable, StakeManager should be changed to check codehash and //StakeVault(msg.sender).STAKING_TOKEN() IERC20 public immutable STAKING_TOKEN; - IStakeManager private stakeManager; + IStakeManagerProxy private stakeManager; + address public stakeManagerImplementationAddress; /** * @dev Emitted when tokens are staked. @@ -42,14 +45,30 @@ contract StakeVault is Ownable { _; } + modifier onlyTrustedStakeManager() { + if (!_stakeManagerImplementationTrusted()) { + revert StakeVault__StakeManagerImplementationNotTrusted(); + } + _; + } + /** * @notice Initializes the contract with the owner, staked token, and stake manager. * @param _owner The address of the owner. * @param _stakeManager The address of the StakeManager contract. */ - constructor(address _owner, IStakeManager _stakeManager) Ownable(_owner) { + constructor(address _owner, IStakeManagerProxy _stakeManager) Ownable(_owner) { STAKING_TOKEN = _stakeManager.STAKING_TOKEN(); stakeManager = _stakeManager; + stakeManagerImplementationAddress = _stakeManager.implementation(); + } + + /** + * @notice Allows the owner to trust a new stake manager implementation. + * @param stakeManagerAddress The address of the new stake manager implementation. + */ + function trustStakeManager(address stakeManagerAddress) external onlyOwner { + stakeManagerImplementationAddress = stakeManagerAddress; } /** @@ -57,7 +76,7 @@ contract StakeVault is Ownable { * @param _amount The amount of tokens to stake. * @param _seconds The time period to stake for. */ - function stake(uint256 _amount, uint256 _seconds) external onlyOwner { + function stake(uint256 _amount, uint256 _seconds) external onlyOwner onlyTrustedStakeManager { _stake(_amount, _seconds, msg.sender); } @@ -67,7 +86,7 @@ contract StakeVault is Ownable { * @param _seconds The time period to stake for. * @param _from The address from which tokens will be transferred. */ - function stake(uint256 _amount, uint256 _seconds, address _from) external onlyOwner { + function stake(uint256 _amount, uint256 _seconds, address _from) external onlyOwner onlyTrustedStakeManager { _stake(_amount, _seconds, _from); } @@ -75,7 +94,7 @@ contract StakeVault is Ownable { * @notice Lock the staked amount for a specified time. * @param _seconds The time period to lock the staked amount for. */ - function lock(uint256 _seconds) external onlyOwner { + function lock(uint256 _seconds) external onlyOwner onlyTrustedStakeManager { stakeManager.lock(_seconds); } @@ -83,7 +102,7 @@ contract StakeVault is Ownable { * @notice Unstake a specified amount of tokens and send to the owner. * @param _amount The amount of tokens to unstake. */ - function unstake(uint256 _amount) external onlyOwner { + function unstake(uint256 _amount) external onlyOwner onlyTrustedStakeManager { _unstake(_amount, msg.sender); } @@ -92,10 +111,44 @@ contract StakeVault is Ownable { * @param _amount The amount of tokens to unstake. * @param _destination The address to receive the unstaked tokens. */ - function unstake(uint256 _amount, address _destination) external onlyOwner validDestination(_destination) { + function unstake( + uint256 _amount, + address _destination + ) + external + onlyOwner + validDestination(_destination) + onlyTrustedStakeManager + { _unstake(_amount, _destination); } + /** + * @notice Withdraw all tokens from the contract to the owner. + */ + function leave(address _destination) external onlyOwner validDestination(_destination) { + if (_stakeManagerImplementationTrusted()) { + // If the stakeManager is trusted, the vault cannot leave the system + // and has to properly unstake instead (which might not be possible if + // funds are locked). + revert StakeVault__NotAllowedToLeave(); + } + + // If the stakeManager is not trusted, we know there was an upgrade. + // In this case, vaults are free to leave the system and move their funds back + // to the owner. + // + // We have to `try/catch` here in case the upgrade was malicious and `leave()` + // either doesn't exist on the new stake manager or reverts for some reason. + // If it was a benign upgrade, it will cause the stake manager to properly update + // its internal accounting before we move the funds out. + try stakeManager.leave() { + STAKING_TOKEN.transfer(_destination, STAKING_TOKEN.balanceOf(address(this))); + } catch { + STAKING_TOKEN.transfer(_destination, STAKING_TOKEN.balanceOf(address(this))); + } + } + /** * @notice Withdraw tokens from the contract. * @param _token The IERC20 token to withdraw. @@ -136,13 +189,11 @@ contract StakeVault is Ownable { } function _stake(uint256 _amount, uint256 _seconds, address _source) internal { + stakeManager.stake(_amount, _seconds); bool success = STAKING_TOKEN.transferFrom(_source, address(this), _amount); if (!success) { revert StakeVault__StakingFailed(); } - - stakeManager.stake(_amount, _seconds); - emit Staked(_source, address(this), _amount, _seconds); } @@ -184,4 +235,8 @@ contract StakeVault is Ownable { STAKING_TOKEN.transfer(_destination, STAKING_TOKEN.balanceOf(address(this))); } } + + function _stakeManagerImplementationTrusted() internal view virtual returns (bool) { + return stakeManagerImplementationAddress == stakeManager.implementation(); + } } diff --git a/src/interfaces/IStakeManager.sol b/src/interfaces/IStakeManager.sol index 6bd84b2..0cdd534 100644 --- a/src/interfaces/IStakeManager.sol +++ b/src/interfaces/IStakeManager.sol @@ -5,14 +5,16 @@ import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; import { ITrustedCodehashAccess } from "./ITrustedCodehashAccess.sol"; interface IStakeManager is ITrustedCodehashAccess { - error StakeManager__FundsLocked(); - error StakeManager__InvalidLockTime(); - error StakeManager__InsufficientFunds(); - error StakeManager__StakeIsTooLow(); + error StakingManager__FundsLocked(); + error StakingManager__InvalidLockTime(); + error StakingManager__InsufficientFunds(); + error StakingManager__StakeIsTooLow(); + error StakingManager__NotAllowedToLeave(); function stake(uint256 _amount, uint256 _seconds) external; function lock(uint256 _seconds) external; function unstake(uint256 _amount) external; + function leave() external; function emergencyModeEnabled() external view returns (bool); function totalStaked() external view returns (uint256); diff --git a/src/interfaces/IStakeManagerProxy.sol b/src/interfaces/IStakeManagerProxy.sol new file mode 100644 index 0000000..c9fd318 --- /dev/null +++ b/src/interfaces/IStakeManagerProxy.sol @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.26; + +import { IStakeManager } from "./IStakeManager.sol"; + +interface IStakeManagerProxy is IStakeManager { + function implementation() external view returns (address); +} diff --git a/test/RewardsStreamerMP.t.sol b/test/RewardsStreamerMP.t.sol index 584eafd..01c7c35 100644 --- a/test/RewardsStreamerMP.t.sol +++ b/test/RewardsStreamerMP.t.sol @@ -7,6 +7,8 @@ import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; import { RewardsStreamerMP } from "../src/RewardsStreamerMP.sol"; import { StakeVault } from "../src/StakeVault.sol"; +import { IStakeManagerProxy } from "../src/interfaces/IStakeManagerProxy.sol"; +import { StakeManagerProxy } from "../src/StakeManagerProxy.sol"; import { MockToken } from "./mocks/MockToken.sol"; contract RewardsStreamerMPTest is Test { @@ -29,7 +31,7 @@ contract RewardsStreamerMPTest is Test { bytes memory initializeData = abi.encodeCall(RewardsStreamerMP.initialize, (address(this), address(stakingToken), address(rewardToken))); address impl = address(new RewardsStreamerMP()); - address proxy = address(new ERC1967Proxy(impl, initializeData)); + address proxy = address(new StakeManagerProxy(impl, initializeData)); streamer = RewardsStreamerMP(proxy); address[4] memory accounts = [alice, bob, charlie, dave]; @@ -93,7 +95,7 @@ contract RewardsStreamerMPTest is Test { function _createTestVault(address owner) internal returns (StakeVault vault) { vm.prank(owner); - vault = new StakeVault(owner, streamer); + vault = new StakeVault(owner, IStakeManagerProxy(address(streamer))); if (!streamer.isTrustedCodehash(address(vault).codehash)) { streamer.setTrustedCodehash(address(vault).codehash, true); @@ -118,6 +120,12 @@ contract RewardsStreamerMPTest is Test { vault.emergencyExit(account); } + function _leave(address account) public { + StakeVault vault = StakeVault(vaults[account]); + vm.prank(account); + vault.leave(account); + } + function _addReward(uint256 amount) public { vm.prank(admin); rewardToken.transfer(address(streamer), amount); @@ -1840,3 +1848,103 @@ contract UpgradeTest is RewardsStreamerMPTest { ); } } + +contract LeaveTest is RewardsStreamerMPTest { + function _upgradeStakeManager() internal { + address newImpl = address(new RewardsStreamerMP()); + bytes memory initializeData; + UUPSUpgradeable(streamer).upgradeToAndCall(newImpl, initializeData); + } + + function setUp() public override { + super.setUp(); + } + + function test_RevertWhenStakeManagerIsTrusted() public { + _stake(alice, 10e18, 0); + vm.expectRevert(StakeVault.StakeVault__NotAllowedToLeave.selector); + _leave(alice); + } + + function test_LeaveShouldProperlyUpdateAccounting() public { + uint256 aliceInitialBalance = stakingToken.balanceOf(alice); + + _stake(alice, 100e18, 0); + + assertEq(stakingToken.balanceOf(alice), aliceInitialBalance - 100e18, "Alice should have staked tokens"); + + checkStreamer( + CheckStreamerParams({ + totalStaked: 100e18, + totalMP: 100e18, + totalMaxMP: 500e18, + stakingBalance: 100e18, + rewardBalance: 0, + rewardIndex: 0, + accountedRewards: 0 + }) + ); + + _upgradeStakeManager(); + _leave(alice); + + // stake manager properly updates accounting + checkStreamer( + CheckStreamerParams({ + totalStaked: 0, + totalMP: 0, + totalMaxMP: 0, + stakingBalance: 0, + rewardBalance: 0, + rewardIndex: 0, + accountedRewards: 0 + }) + ); + + // vault should be empty as funds have been moved out + checkAccount( + CheckAccountParams({ + account: vaults[alice], + rewardBalance: 0, + stakedBalance: 0, + vaultBalance: 0, + rewardIndex: 0, + accountMP: 0, + maxMP: 0 + }) + ); + + assertEq(stakingToken.balanceOf(alice), aliceInitialBalance, "Alice has all her funds back"); + } + + function test_TrustNewStakeManager() public { + // first, upgrade to new stake manager, marking it as not trusted + _upgradeStakeManager(); + + // ensure vault functions revert if stake manager is not trusted + vm.expectRevert(StakeVault.StakeVault__StakeManagerImplementationNotTrusted.selector); + _stake(alice, 100e18, 0); + + // ensure vault functions revert if stake manager is not trusted + StakeVault vault = StakeVault(vaults[alice]); + vm.prank(alice); + vm.expectRevert(StakeVault.StakeVault__StakeManagerImplementationNotTrusted.selector); + vault.lock(365 days); + + // ensure vault functions revert if stake manager is not trusted + vm.expectRevert(StakeVault.StakeVault__StakeManagerImplementationNotTrusted.selector); + _unstake(alice, 100e18); + + // now, trust the new stake manager + address newStakeManagerImpl = IStakeManagerProxy(address(streamer)).implementation(); + vm.prank(alice); + vault.trustStakeManager(newStakeManagerImpl); + + // stake manager is now trusted, so functions are enabeled again + _stake(alice, 100e18, 0); + + // however, a trusted manager cannot be left + vm.expectRevert(StakeVault.StakeVault__NotAllowedToLeave.selector); + _leave(alice); + } +} diff --git a/test/StakeVault.test.sol b/test/StakeVault.test.sol index d5d44a6..8316e22 100644 --- a/test/StakeVault.test.sol +++ b/test/StakeVault.test.sol @@ -4,6 +4,8 @@ pragma solidity ^0.8.26; import { Test } from "forge-std/Test.sol"; import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; +import { IStakeManagerProxy } from "../src/interfaces/IStakeManagerProxy.sol"; +import { StakeManagerProxy } from "../src/StakeManagerProxy.sol"; import { RewardsStreamerMP } from "../src/RewardsStreamerMP.sol"; import { StakeVault } from "../src/StakeVault.sol"; import { MockToken } from "./mocks/MockToken.sol"; @@ -21,7 +23,7 @@ contract StakeVaultTest is Test { function _createTestVault(address owner) internal returns (StakeVault vault) { vm.prank(owner); - vault = new StakeVault(owner, streamer); + vault = new StakeVault(owner, IStakeManagerProxy(address(streamer))); if (!streamer.isTrustedCodehash(address(vault).codehash)) { streamer.setTrustedCodehash(address(vault).codehash, true); @@ -35,7 +37,7 @@ contract StakeVaultTest is Test { bytes memory initializeData = abi.encodeWithSelector( RewardsStreamerMP.initialize.selector, address(this), address(stakingToken), address(rewardToken) ); - address proxy = address(new ERC1967Proxy(impl, initializeData)); + address proxy = address(new StakeManagerProxy(impl, initializeData)); streamer = RewardsStreamerMP(proxy); stakingToken.mint(alice, 10_000e18); @@ -48,7 +50,7 @@ contract StakeVaultTest is Test { contract StakingTokenTest is StakeVaultTest { function setUp() public override { - StakeVaultTest.setUp(); + super.setUp(); } function testStakeToken() public view { @@ -58,7 +60,7 @@ contract StakingTokenTest is StakeVaultTest { contract WithdrawTest is StakeVaultTest { function setUp() public override { - StakeVaultTest.setUp(); + super.setUp(); } function test_CannotWithdrawStakedFunds() public {