From 7ad3428d4961ebbc69428a5432f14f93fc0ed278 Mon Sep 17 00:00:00 2001 From: r4bbit <445106+0x-r4bbit@users.noreply.github.com> Date: Wed, 27 Nov 2024 12:59:29 +0100 Subject: [PATCH] feat(RewardsStreamerMP): introduce `leave()` function This function allows users to `leave()` the system if they can't or don't want to trust the stake manager. This is the case when the owner of the stake manager performs an upgrade (upgradeability lands in future changes). In case of such an upgrade, the stake manager will be marked as not trusted which prevents the user from staking, unstaking, locking etc. The user can then either explicitly trust stake manager (will happen in future changes) to enable the vault's functionality again, or, `leave()` the system at which point it will try to perform a benign `leave()` operation and then move the funds out of the vault. Closes #66 --- .gas-report | 46 ++++++++------ .gas-snapshot | 88 +++++++++++++------------- certora/specs/EmergencyMode.spec | 3 +- certora/specs/RewardsStreamerMP.spec | 1 + src/RewardsStreamerMP.sol | 28 ++++++-- src/StakeVault.sol | 61 +++++++++++++++--- src/interfaces/IStakeManager.sol | 10 +-- test/RewardsStreamerMP.t.sol | 83 ++++++++++++++++++++++++ test/harness/StakeVaultNonTrusting.sol | 17 +++++ 9 files changed, 259 insertions(+), 78 deletions(-) create mode 100644 test/harness/StakeVaultNonTrusting.sol diff --git a/.gas-report b/.gas-report index 7179b6a..81cd69a 100644 --- a/.gas-report +++ b/.gas-report @@ -15,24 +15,24 @@ | src/RewardsStreamerMP.sol:RewardsStreamerMP contract | | | | | | |------------------------------------------------------|-----------------|-------|--------|-------|---------| | Deployment Cost | Deployment Size | | | | | -| 1425736 | 6528 | | | | | +| 1464062 | 6707 | | | | | | Function Name | min | avg | median | max | # calls | | MAX_LOCKUP_PERIOD | 228 | 228 | 228 | 228 | 23 | | MAX_MULTIPLIER | 274 | 274 | 274 | 274 | 30 | | MIN_LOCKUP_PERIOD | 275 | 275 | 275 | 275 | 11 | | MP_RATE_PER_YEAR | 231 | 231 | 231 | 231 | 3 | -| SCALE_FACTOR | 295 | 295 | 295 | 295 | 41 | -| STAKING_TOKEN | 273 | 273 | 273 | 273 | 172 | -| accountedRewards | 351 | 906 | 351 | 2351 | 72 | +| SCALE_FACTOR | 273 | 273 | 273 | 273 | 41 | +| STAKING_TOKEN | 273 | 273 | 273 | 273 | 182 | +| accountedRewards | 329 | 896 | 329 | 2329 | 74 | | emergencyModeEnabled | 2377 | 2377 | 2377 | 2377 | 7 | -| enableEmergencyMode | 23504 | 40411 | 45696 | 45696 | 8 | -| getAccount | 1596 | 1596 | 1596 | 1596 | 71 | -| isTrustedCodehash | 496 | 996 | 496 | 2496 | 172 | -| rewardIndex | 373 | 400 | 373 | 2373 | 72 | -| setTrustedCodehash | 47926 | 47926 | 47926 | 47926 | 43 | -| totalMP | 330 | 330 | 330 | 330 | 75 | -| totalMaxMP | 351 | 351 | 351 | 351 | 75 | -| totalStaked | 330 | 330 | 330 | 330 | 75 | +| enableEmergencyMode | 23482 | 40389 | 45674 | 45674 | 8 | +| getAccount | 1596 | 1596 | 1596 | 1596 | 72 | +| isTrustedCodehash | 496 | 996 | 496 | 2496 | 180 | +| rewardIndex | 373 | 400 | 373 | 2373 | 74 | +| setTrustedCodehash | 47926 | 47926 | 47926 | 47926 | 47 | +| totalMP | 330 | 330 | 330 | 330 | 77 | +| totalMaxMP | 351 | 351 | 351 | 351 | 77 | +| totalStaked | 330 | 330 | 330 | 330 | 77 | | updateAccountMP | 36758 | 38996 | 39260 | 39260 | 19 | | updateGlobalState | 32134 | 60366 | 49513 | 82460 | 28 | @@ -40,12 +40,13 @@ | src/StakeVault.sol:StakeVault contract | | | | | | |----------------------------------------|-----------------|--------|--------|--------|---------| | Deployment Cost | Deployment Size | | | | | -| 1095864 | 5202 | | | | | +| 1130636 | 5366 | | | | | | Function Name | min | avg | median | max | # calls | | emergencyExit | 31410 | 43924 | 43160 | 60260 | 7 | +| leave | 24019 | 24019 | 24019 | 24019 | 1 | | lock | 38362 | 60516 | 42741 | 100445 | 3 | -| stake | 199213 | 236948 | 242906 | 263435 | 55 | -| unstake | 84988 | 113999 | 103434 | 146128 | 13 | +| stake | 199207 | 237049 | 242900 | 263429 | 56 | +| unstake | 85027 | 114046 | 103483 | 146178 | 13 | | src/XPNFTToken.sol:XPNFTToken contract | | | | | | @@ -108,6 +109,15 @@ | urlSuffix | 1250 | 1250 | 1250 | 1250 | 1 | +| test/harness/StakeVaultNonTrusting.sol:StakeVaultNonTrusting contract | | | | | | +|-----------------------------------------------------------------------|-----------------|--------|--------|--------|---------| +| Deployment Cost | Deployment Size | | | | | +| 1090189 | 5189 | | | | | +| Function Name | min | avg | median | max | # calls | +| leave | 84495 | 84495 | 84495 | 84495 | 1 | +| stakeNonTrusted | 240739 | 240739 | 240739 | 240739 | 1 | + + | test/mocks/MockMetadataGenerator.sol:MockMetadataGenerator contract | | | | | | |---------------------------------------------------------------------|-----------------|-------|--------|-------|---------| | Deployment Cost | Deployment Size | | | | | @@ -121,9 +131,9 @@ | Deployment Cost | Deployment Size | | | | | | 639406 | 3369 | | | | | | Function Name | min | avg | median | max | # calls | -| approve | 46334 | 46343 | 46346 | 46346 | 220 | -| balanceOf | 561 | 1381 | 561 | 2561 | 334 | -| mint | 51284 | 58817 | 51284 | 68384 | 236 | +| approve | 46334 | 46343 | 46346 | 46346 | 232 | +| balanceOf | 561 | 1388 | 561 | 2561 | 343 | +| mint | 51284 | 58789 | 51284 | 68384 | 246 | | transfer | 34390 | 48859 | 51490 | 51490 | 13 | diff --git a/.gas-snapshot b/.gas-snapshot index 94d93e7..a56787d 100644 --- a/.gas-snapshot +++ b/.gas-snapshot @@ -1,15 +1,17 @@ -EmergencyExitTest:test_CannotEnableEmergencyModeTwice() (gas: 79829) -EmergencyExitTest:test_CannotLeaveBeforeEmergencyMode() (gas: 283234) -EmergencyExitTest:test_EmergencyExitBasic() (gas: 379874) -EmergencyExitTest:test_EmergencyExitMultipleUsers() (gas: 788714) -EmergencyExitTest:test_EmergencyExitToAlternateAddress() (gas: 517455) -EmergencyExitTest:test_EmergencyExitWithLock() (gas: 370304) -EmergencyExitTest:test_EmergencyExitWithRewards() (gas: 507361) -EmergencyExitTest:test_OnlyOwnerCanEnableEmergencyMode() (gas: 34607) -IntegrationTest:testStakeFoo() (gas: 1490300) -LockTest:test_LockFailsWithInvalidPeriod() (gas: 290178) -LockTest:test_LockFailsWithNoStake() (gas: 53357) -LockTest:test_LockWithoutPriorLock() (gas: 383272) +EmergencyExitTest:test_CannotEnableEmergencyModeTwice() (gas: 79785) +EmergencyExitTest:test_CannotLeaveBeforeEmergencyMode() (gas: 283229) +EmergencyExitTest:test_EmergencyExitBasic() (gas: 379825) +EmergencyExitTest:test_EmergencyExitMultipleUsers() (gas: 788638) +EmergencyExitTest:test_EmergencyExitToAlternateAddress() (gas: 517427) +EmergencyExitTest:test_EmergencyExitWithLock() (gas: 370277) +EmergencyExitTest:test_EmergencyExitWithRewards() (gas: 507312) +EmergencyExitTest:test_OnlyOwnerCanEnableEmergencyMode() (gas: 34585) +IntegrationTest:testStakeFoo() (gas: 1490184) +LeaveTest:test_LeaveShouldProperlyUpdateAccounting() (gas: 365536) +LeaveTest:test_RevertWhenStakeManagerIsTrusted() (gas: 275827) +LockTest:test_LockFailsWithInvalidPeriod() (gas: 290237) +LockTest:test_LockFailsWithNoStake() (gas: 53423) +LockTest:test_LockWithoutPriorLock() (gas: 383200) NFTMetadataGeneratorSVGTest:testGenerateMetadata() (gas: 92874) NFTMetadataGeneratorSVGTest:testSetImageStrings() (gas: 60081) NFTMetadataGeneratorSVGTest:testSetImageStringsRevert() (gas: 35818) @@ -17,37 +19,37 @@ NFTMetadataGeneratorURLTest:testGenerateMetadata() (gas: 109345) NFTMetadataGeneratorURLTest:testSetBaseURL() (gas: 50653) NFTMetadataGeneratorURLTest:testSetBaseURLRevert() (gas: 35993) RewardsStreamerTest:testStake() (gas: 869874) -StakeTest:test_StakeMultipleAccounts() (gas: 493279) -StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 640763) -StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 818252) -StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 499381) -StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 520783) -StakeTest:test_StakeOneAccount() (gas: 284277) -StakeTest:test_StakeOneAccountAndRewards() (gas: 431756) -StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 498901) -StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 494078) -StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 298175) -StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 298187) -StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 298298) -UnstakeTest:test_StakeMultipleAccounts() (gas: 493323) -UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 640807) -UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 818251) -UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 499380) -UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 520827) -UnstakeTest:test_StakeOneAccount() (gas: 284300) -UnstakeTest:test_StakeOneAccountAndRewards() (gas: 431800) -UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 498945) -UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 494080) -UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 298132) -UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 298187) -UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 298298) -UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 508511) -UnstakeTest:test_UnstakeMultipleAccounts() (gas: 688755) -UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 1014239) -UnstakeTest:test_UnstakeOneAccount() (gas: 480152) -UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 496638) -UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 585964) -UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 518574) +StakeTest:test_StakeMultipleAccounts() (gas: 493245) +StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 640707) +StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 818174) +StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 499259) +StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 520661) +StakeTest:test_StakeOneAccount() (gas: 284249) +StakeTest:test_StakeOneAccountAndRewards() (gas: 431706) +StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 498829) +StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 493984) +StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 298103) +StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 298115) +StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 298226) +UnstakeTest:test_StakeMultipleAccounts() (gas: 493267) +UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 640729) +UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 818151) +UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 499236) +UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 520683) +UnstakeTest:test_StakeOneAccount() (gas: 284272) +UnstakeTest:test_StakeOneAccountAndRewards() (gas: 431728) +UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 498851) +UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 493964) +UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 298103) +UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 298115) +UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 298204) +UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 508466) +UnstakeTest:test_UnstakeMultipleAccounts() (gas: 688777) +UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 1014246) +UnstakeTest:test_UnstakeOneAccount() (gas: 480136) +UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 496593) +UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 585942) +UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 518419) XPNFTTokenTest:testApproveNotAllowed() (gas: 10507) XPNFTTokenTest:testGetApproved() (gas: 10531) XPNFTTokenTest:testIsApprovedForAll() (gas: 10705) diff --git a/certora/specs/EmergencyMode.spec b/certora/specs/EmergencyMode.spec index 1de2d50..ac62c18 100644 --- a/certora/specs/EmergencyMode.spec +++ b/certora/specs/EmergencyMode.spec @@ -47,7 +47,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); } diff --git a/certora/specs/RewardsStreamerMP.spec b/certora/specs/RewardsStreamerMP.spec index 90d105c..34def34 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 c5acb81..1388b27 100644 --- a/src/RewardsStreamerMP.sol +++ b/src/RewardsStreamerMP.sol @@ -150,13 +150,16 @@ contract RewardsStreamerMP is IStakeManager, TrustedCodehashAccess, ReentrancyGu 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; @@ -167,11 +170,28 @@ contract RewardsStreamerMP is IStakeManager, TrustedCodehashAccess, ReentrancyGu 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/StakeVault.sol b/src/StakeVault.sol index efb0322..18ab3ab 100644 --- a/src/StakeVault.sol +++ b/src/StakeVault.sol @@ -19,6 +19,8 @@ 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 @@ -42,6 +44,13 @@ 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. @@ -57,7 +66,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 +76,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 +84,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 +92,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 +101,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 +179,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 +225,8 @@ contract StakeVault is Ownable { STAKING_TOKEN.transfer(_destination, STAKING_TOKEN.balanceOf(address(this))); } } + + function _stakeManagerImplementationTrusted() internal pure virtual returns (bool) { + return true; + } } 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/test/RewardsStreamerMP.t.sol b/test/RewardsStreamerMP.t.sol index 5800f27..d12dd39 100644 --- a/test/RewardsStreamerMP.t.sol +++ b/test/RewardsStreamerMP.t.sol @@ -5,6 +5,7 @@ import { Test } from "forge-std/Test.sol"; import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol"; import { RewardsStreamerMP } from "../src/RewardsStreamerMP.sol"; import { StakeVault } from "../src/StakeVault.sol"; +import { StakeVaultNonTrusting } from "./harness/StakeVaultNonTrusting.sol"; import { MockToken } from "./mocks/MockToken.sol"; contract RewardsStreamerMPTest is Test { @@ -111,6 +112,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); @@ -1779,3 +1786,79 @@ contract EmergencyExitTest is RewardsStreamerMPTest { ); } } + +contract LeaveTest is RewardsStreamerMPTest { + StakeVaultNonTrusting public nonTrustingVault; + + function setUp() public override { + super.setUp(); + + // set up vault that doesn't trust `streamer` + vm.startPrank(alice); + nonTrustingVault = new StakeVaultNonTrusting(alice, streamer); + stakingToken.approve(address(nonTrustingVault), 10_000e18); + vm.stopPrank(); + + // ensure non trusting vault is trusted by stake manager + streamer.setTrustedCodehash(address(nonTrustingVault).codehash, true); + } + + function test_RevertWhenStakeManagerIsTrusted() public { + // this test uses the standard `StakeVault` which trusts the `streamer` + _stake(alice, 10e18, 0); + vm.expectRevert(StakeVault.StakeVault__NotAllowedToLeave.selector); + _leave(alice); + } + + function test_LeaveShouldProperlyUpdateAccounting() public { + uint256 aliceInitialBalance = stakingToken.balanceOf(alice); + + vm.prank(alice); + nonTrustingVault.stakeNonTrusted(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 + }) + ); + + vm.prank(alice); + nonTrustingVault.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: address(nonTrustingVault), + rewardBalance: 0, + stakedBalance: 0, + vaultBalance: 0, + rewardIndex: 0, + accountMP: 0, + maxMP: 0 + }) + ); + + assertEq(stakingToken.balanceOf(alice), aliceInitialBalance, "Alice has all her funds back"); + } +} diff --git a/test/harness/StakeVaultNonTrusting.sol b/test/harness/StakeVaultNonTrusting.sol new file mode 100644 index 0000000..cab0d7d --- /dev/null +++ b/test/harness/StakeVaultNonTrusting.sol @@ -0,0 +1,17 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.26; + +import { IStakeManager } from "../../src/interfaces/IStakeManager.sol"; +import { StakeVault } from "../../src/StakeVault.sol"; + +contract StakeVaultNonTrusting is StakeVault { + constructor(address _owner, IStakeManager _stakeManager) StakeVault(_owner, _stakeManager) { } + + function _stakeManagerImplementationTrusted() internal pure override returns (bool) { + return false; + } + + function stakeNonTrusted(uint256 _amount, uint256 _time) public { + _stake(_amount, _time, msg.sender); + } +}