diff --git a/.gas-report b/.gas-report index 7553415..c589bf0 100644 --- a/.gas-report +++ b/.gas-report @@ -15,87 +15,93 @@ | src/RewardsStreamerMP.sol:RewardsStreamerMP contract | | | | | | |------------------------------------------------------|-----------------|--------|--------|--------|---------| | Deployment Cost | Deployment Size | | | | | -| 2116504 | 9735 | | | | | +| 2445388 | 11266 | | | | | | Function Name | min | avg | median | max | # calls | | MAX_LOCKUP_PERIOD | 294 | 294 | 294 | 294 | 23 | | MAX_MULTIPLIER | 251 | 251 | 251 | 251 | 30 | -| MIN_LOCKUP_PERIOD | 252 | 252 | 252 | 252 | 11 | +| MIN_LOCKUP_PERIOD | 274 | 274 | 274 | 274 | 11 | | MP_RATE_PER_YEAR | 253 | 253 | 253 | 253 | 3 | | SCALE_FACTOR | 273 | 273 | 273 | 273 | 41 | -| STAKING_TOKEN | 2404 | 2404 | 2404 | 2404 | 222 | -| emergencyModeEnabled | 2420 | 2420 | 2420 | 2420 | 7 | +| STAKING_TOKEN | 2404 | 2404 | 2404 | 2404 | 329 | +| emergencyModeEnabled | 2398 | 2398 | 2398 | 2398 | 7 | | enableEmergencyMode | 2439 | 19346 | 24631 | 24631 | 8 | -| getAccount | 1621 | 1621 | 1621 | 1621 | 72 | -| getStakedBalance | 2611 | 2611 | 2611 | 2611 | 1 | -| initialize | 115536 | 115536 | 115536 | 115536 | 57 | -| isTrustedCodehash | 541 | 1054 | 541 | 2541 | 222 | -| lastRewardTime | 351 | 1351 | 1351 | 2351 | 2 | -| leave | 56077 | 56077 | 56077 | 56077 | 1 | -| lock | 9861 | 31868 | 14190 | 71554 | 3 | +| getAccount | 1615 | 1615 | 1615 | 1615 | 72 | +| getStakedBalance | 2562 | 2562 | 2562 | 2562 | 1 | +| getUserVaults | 5197 | 5197 | 5197 | 5197 | 40 | +| initialize | 115551 | 115551 | 115551 | 115551 | 67 | +| isTrustedCodehash | 563 | 563 | 563 | 563 | 260 | +| lastRewardTime | 373 | 1373 | 1373 | 2373 | 2 | +| leave | 56122 | 56122 | 56122 | 56122 | 1 | +| lock | 12029 | 34036 | 16358 | 73722 | 3 | | proxiableUUID | 319 | 319 | 319 | 319 | 3 | -| rewardEndTime | 373 | 1373 | 1373 | 2373 | 2 | -| rewardStartTime | 396 | 1396 | 1396 | 2396 | 2 | -| rewardsBalanceOf | 1276 | 1276 | 1276 | 1276 | 4 | +| registerVault | 72848 | 72848 | 72848 | 72848 | 262 | +| rewardEndTime | 395 | 1395 | 1395 | 2395 | 2 | +| rewardStartTime | 374 | 1374 | 1374 | 2374 | 2 | +| rewardsBalanceOf | 1292 | 1292 | 1292 | 1292 | 4 | | setReward | 2537 | 50851 | 60232 | 102549 | 7 | -| setTrustedCodehash | 26197 | 26197 | 26197 | 26197 | 57 | -| stake | 128652 | 168926 | 175469 | 195948 | 63 | +| setTrustedCodehash | 26174 | 26174 | 26174 | 26174 | 67 | +| stake | 130820 | 171094 | 177637 | 198116 | 63 | | totalMP | 373 | 373 | 373 | 373 | 80 | -| totalMaxMP | 372 | 372 | 372 | 372 | 80 | +| totalMaxMP | 350 | 350 | 350 | 350 | 80 | | totalRewardsAccrued | 373 | 373 | 373 | 373 | 3 | | totalRewardsSupply | 1003 | 1964 | 1767 | 6743 | 30 | | totalStaked | 351 | 351 | 351 | 351 | 81 | -| unstake | 57991 | 58605 | 57991 | 61984 | 13 | -| updateAccountMP | 15397 | 18475 | 17899 | 34999 | 21 | +| unstake | 60225 | 60839 | 60225 | 64218 | 13 | +| updateAccountMP | 15391 | 18469 | 17893 | 34993 | 21 | | updateGlobalState | 11066 | 28094 | 25315 | 110295 | 21 | -| upgradeToAndCall | 3125 | 9237 | 10766 | 10766 | 5 | +| upgradeToAndCall | 3141 | 9258 | 10782 | 10804 | 5 | | src/StakeManagerProxy.sol:StakeManagerProxy contract | | | | | | |------------------------------------------------------|-----------------|-------|--------|--------|---------| | Deployment Cost | Deployment Size | | | | | -| 256004 | 1231 | | | | | +| 256019 | 1231 | | | | | | Function Name | min | avg | median | max | # calls | | MAX_LOCKUP_PERIOD | 721 | 1503 | 721 | 5221 | 23 | | MAX_MULTIPLIER | 678 | 1578 | 678 | 5178 | 30 | -| MIN_LOCKUP_PERIOD | 679 | 3951 | 5179 | 5179 | 11 | +| MIN_LOCKUP_PERIOD | 701 | 3973 | 5201 | 5201 | 11 | | MP_RATE_PER_YEAR | 680 | 680 | 680 | 680 | 3 | | SCALE_FACTOR | 700 | 700 | 700 | 700 | 41 | -| STAKING_TOKEN | 7331 | 7331 | 7331 | 7331 | 222 | -| emergencyModeEnabled | 7347 | 7347 | 7347 | 7347 | 7 | +| STAKING_TOKEN | 7331 | 7331 | 7331 | 7331 | 329 | +| emergencyModeEnabled | 7325 | 7325 | 7325 | 7325 | 7 | | enableEmergencyMode | 28434 | 45335 | 50619 | 50619 | 8 | -| getAccount | 2075 | 2075 | 2075 | 2075 | 72 | -| getStakedBalance | 7541 | 7541 | 7541 | 7541 | 1 | -| implementation | 343 | 899 | 343 | 2343 | 309 | -| isTrustedCodehash | 971 | 1484 | 971 | 2971 | 222 | -| lastRewardTime | 778 | 1778 | 1778 | 2778 | 2 | -| rewardEndTime | 800 | 1800 | 1800 | 2800 | 2 | -| rewardStartTime | 823 | 4073 | 4073 | 7323 | 2 | -| rewardsBalanceOf | 1706 | 1706 | 1706 | 1706 | 4 | +| getAccount | 2069 | 2069 | 2069 | 2069 | 72 | +| getStakedBalance | 7492 | 7492 | 7492 | 7492 | 1 | +| getUserVaults | 5633 | 6758 | 5633 | 10133 | 40 | +| implementation | 343 | 756 | 343 | 2343 | 416 | +| isTrustedCodehash | 993 | 993 | 993 | 993 | 260 | +| lastRewardTime | 800 | 1800 | 1800 | 2800 | 2 | +| rewardEndTime | 822 | 1822 | 1822 | 2822 | 2 | +| rewardStartTime | 801 | 4051 | 4051 | 7301 | 2 | +| rewardsBalanceOf | 1722 | 1722 | 1722 | 1722 | 4 | | setReward | 28817 | 77165 | 86590 | 128835 | 7 | -| setTrustedCodehash | 52843 | 52843 | 52843 | 52843 | 57 | +| setTrustedCodehash | 52820 | 52820 | 52820 | 52820 | 67 | | totalMP | 800 | 800 | 800 | 800 | 80 | -| totalMaxMP | 799 | 799 | 799 | 799 | 80 | +| totalMaxMP | 777 | 777 | 777 | 777 | 80 | | totalRewardsAccrued | 800 | 800 | 800 | 800 | 3 | | totalRewardsSupply | 1430 | 2541 | 2194 | 11670 | 30 | | totalStaked | 778 | 778 | 778 | 778 | 81 | -| updateAccountMP | 41756 | 44834 | 44258 | 61358 | 21 | +| updateAccountMP | 41750 | 44828 | 44252 | 61352 | 21 | | updateGlobalState | 37054 | 54082 | 51303 | 136283 | 21 | -| upgradeToAndCall | 29768 | 35875 | 37402 | 37402 | 5 | +| upgradeToAndCall | 29784 | 35895 | 37418 | 37440 | 5 | | src/StakeVault.sol:StakeVault contract | | | | | | |----------------------------------------|-----------------|--------|--------|--------|---------| | Deployment Cost | Deployment Size | | | | | -| 1374566 | 6483 | | | | | +| 1420401 | 6695 | | | | | | Function Name | min | avg | median | max | # calls | | STAKING_TOKEN | 216 | 216 | 216 | 216 | 1 | -| emergencyExit | 36375 | 48879 | 48113 | 65213 | 7 | -| leave | 33507 | 131405 | 60569 | 370978 | 4 | -| lock | 33245 | 58948 | 48599 | 105351 | 4 | -| stake | 33454 | 240127 | 250016 | 270543 | 64 | -| trustStakeManager | 28997 | 28997 | 28997 | 28997 | 1 | -| unstake | 33260 | 94528 | 99898 | 107496 | 14 | -| withdraw | 42271 | 42271 | 42271 | 42271 | 1 | +| emergencyExit | 36353 | 48857 | 48091 | 65191 | 7 | +| leave | 33507 | 131414 | 60587 | 370978 | 4 | +| lock | 33245 | 60574 | 50767 | 107519 | 4 | +| owner | 2339 | 2339 | 2339 | 2339 | 262 | +| register | 103997 | 103997 | 103997 | 103997 | 262 | +| stake | 33411 | 242219 | 252141 | 272668 | 64 | +| stakeManager | 368 | 368 | 368 | 368 | 262 | +| trustStakeManager | 28953 | 28953 | 28953 | 28953 | 1 | +| unstake | 33282 | 96560 | 102154 | 109752 | 14 | +| withdraw | 42222 | 42222 | 42222 | 42222 | 1 | | src/XPNFTToken.sol:XPNFTToken contract | | | | | | @@ -171,19 +177,19 @@ | Deployment Cost | Deployment Size | | | | | | 625454 | 3260 | | | | | | Function Name | min | avg | median | max | # calls | -| approve | 46330 | 46339 | 46342 | 46342 | 227 | +| approve | 46330 | 46339 | 46342 | 46342 | 267 | | balanceOf | 558 | 989 | 558 | 2558 | 139 | -| mint | 51279 | 56555 | 51279 | 68379 | 243 | +| mint | 51279 | 56413 | 51279 | 68379 | 283 | | transfer | 34384 | 42934 | 42934 | 51484 | 2 | | test/mocks/StackOverflowStakeManager.sol:StackOverflowStakeManager contract | | | | | | |-----------------------------------------------------------------------------|-----------------|--------|--------|--------|---------| | Deployment Cost | Deployment Size | | | | | -| 1033443 | 4615 | | | | | +| 1026739 | 4584 | | | | | | Function Name | min | avg | median | max | # calls | | leave | 391 | 161316 | 161316 | 322322 | 334 | -| proxiableUUID | 319 | 319 | 319 | 319 | 1 | +| proxiableUUID | 341 | 341 | 341 | 341 | 1 | | test/mocks/XPProviderMock.sol:XPProviderMock contract | | | | | | diff --git a/.gas-snapshot b/.gas-snapshot index 115d2e1..36ab500 100644 --- a/.gas-snapshot +++ b/.gas-snapshot @@ -1,67 +1,77 @@ EmergencyExitTest:test_CannotEnableEmergencyModeTwice() (gas: 92598) -EmergencyExitTest:test_CannotLeaveBeforeEmergencyMode() (gas: 295308) -EmergencyExitTest:test_EmergencyExitBasic() (gas: 381937) -EmergencyExitTest:test_EmergencyExitMultipleUsers() (gas: 654329) -EmergencyExitTest:test_EmergencyExitToAlternateAddress() (gas: 389879) -EmergencyExitTest:test_EmergencyExitWithLock() (gas: 389475) -EmergencyExitTest:test_EmergencyExitWithRewards() (gas: 374838) +EmergencyExitTest:test_CannotLeaveBeforeEmergencyMode() (gas: 297433) +EmergencyExitTest:test_EmergencyExitBasic() (gas: 384012) +EmergencyExitTest:test_EmergencyExitMultipleUsers() (gas: 658414) +EmergencyExitTest:test_EmergencyExitToAlternateAddress() (gas: 391976) +EmergencyExitTest:test_EmergencyExitWithLock() (gas: 391600) +EmergencyExitTest:test_EmergencyExitWithRewards() (gas: 376941) EmergencyExitTest:test_OnlyOwnerCanEnableEmergencyMode() (gas: 39362) -IntegrationTest:testStakeFoo() (gas: 1166661) -LeaveTest:test_LeaveShouldProperlyUpdateAccounting() (gas: 2548215) -LeaveTest:test_RevertWhenStakeManagerIsTrusted() (gas: 292439) -LeaveTest:test_TrustNewStakeManager() (gas: 2623611) -LockTest:test_LockFailsWithInvalidPeriod() (gas: 305322) -LockTest:test_LockFailsWithNoStake() (gas: 61418) -LockTest:test_LockWithoutPriorLock() (gas: 385914) -MaliciousUpgradeTest:test_UpgradeStackOverflowStakeManager() (gas: 1749466) +EmergencyExitTest:test_VaultsRegistered() (gas: 62064) +IntegrationTest:testStakeFoo() (gas: 1177254) +IntegrationTest:test_VaultsRegistered() (gas: 62064) +LeaveTest:test_LeaveShouldProperlyUpdateAccounting() (gas: 2879645) +LeaveTest:test_RevertWhenStakeManagerIsTrusted() (gas: 294542) +LeaveTest:test_TrustNewStakeManager() (gas: 2957495) +LeaveTest:test_VaultsRegistered() (gas: 61999) +LockTest:test_LockFailsWithInvalidPeriod() (gas: 309615) +LockTest:test_LockFailsWithNoStake() (gas: 63586) +LockTest:test_LockWithoutPriorLock() (gas: 390173) +LockTest:test_VaultsRegistered() (gas: 62042) +MaliciousUpgradeTest:test_UpgradeStackOverflowStakeManager() (gas: 1744894) +MaliciousUpgradeTest:test_VaultsRegistered() (gas: 62042) NFTMetadataGeneratorSVGTest:testGenerateMetadata() (gas: 85934) NFTMetadataGeneratorSVGTest:testSetImageStrings() (gas: 58332) NFTMetadataGeneratorSVGTest:testSetImageStringsRevert() (gas: 35804) NFTMetadataGeneratorURLTest:testGenerateMetadata() (gas: 102512) NFTMetadataGeneratorURLTest:testSetBaseURL() (gas: 49555) NFTMetadataGeneratorURLTest:testSetBaseURLRevert() (gas: 35979) -RewardsStreamerMP_RewardsTest:testRewardsBalanceOf() (gas: 668216) -RewardsStreamerMP_RewardsTest:testSetRewards() (gas: 160234) +RewardsStreamerMPTest:test_VaultsRegistered() (gas: 62042) +RewardsStreamerMP_RewardsTest:testRewardsBalanceOf() (gas: 670393) +RewardsStreamerMP_RewardsTest:testSetRewards() (gas: 160278) RewardsStreamerMP_RewardsTest:testSetRewards_RevertsBadAmount() (gas: 39364) RewardsStreamerMP_RewardsTest:testSetRewards_RevertsBadDuration() (gas: 39300) RewardsStreamerMP_RewardsTest:testSetRewards_RevertsNotAuthorized() (gas: 39335) -RewardsStreamerMP_RewardsTest:testTotalRewardsSupply() (gas: 609395) +RewardsStreamerMP_RewardsTest:testTotalRewardsSupply() (gas: 611520) +RewardsStreamerMP_RewardsTest:test_VaultsRegistered() (gas: 62021) RewardsStreamerTest:testStake() (gas: 869181) -StakeTest:test_StakeMultipleAccounts() (gas: 489543) -StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 495503) -StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 825688) -StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 512574) -StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 534457) -StakeTest:test_StakeOneAccount() (gas: 274483) -StakeTest:test_StakeOneAccountAndRewards() (gas: 280407) -StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 497162) -StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 493499) -StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 299311) -StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 299300) -StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 299411) +StakeTest:test_StakeMultipleAccounts() (gas: 493759) +StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 499674) +StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 829812) +StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 516824) +StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 538707) +StakeTest:test_StakeOneAccount() (gas: 276580) +StakeTest:test_StakeOneAccountAndRewards() (gas: 282504) +StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 499197) +StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 495534) +StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 301459) +StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 301425) +StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 301536) +StakeTest:test_VaultsRegistered() (gas: 62021) StakingTokenTest:testStakeToken() (gas: 10422) -UnstakeTest:test_StakeMultipleAccounts() (gas: 489587) -UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 495480) -UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 825687) -UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 512573) -UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 534501) -UnstakeTest:test_StakeOneAccount() (gas: 274506) -UnstakeTest:test_StakeOneAccountAndRewards() (gas: 280451) -UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 497206) -UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 493501) -UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 299356) -UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 299300) -UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 299411) -UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 537481) -UnstakeTest:test_UnstakeMultipleAccounts() (gas: 682793) -UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 774872) -UnstakeTest:test_UnstakeOneAccount() (gas: 466459) -UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 489723) -UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 399173) -UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 526156) -UpgradeTest:test_RevertWhenNotOwner() (gas: 2192157) -UpgradeTest:test_UpgradeStakeManager() (gas: 2463165) -WithdrawTest:test_CannotWithdrawStakedFunds() (gas: 308100) +UnstakeTest:test_StakeMultipleAccounts() (gas: 493781) +UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 499674) +UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 829789) +UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 516801) +UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 538729) +UnstakeTest:test_StakeOneAccount() (gas: 276603) +UnstakeTest:test_StakeOneAccountAndRewards() (gas: 282526) +UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 499219) +UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 495514) +UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 301459) +UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 301425) +UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 301514) +UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 541772) +UnstakeTest:test_UnstakeMultipleAccounts() (gas: 691487) +UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 784848) +UnstakeTest:test_UnstakeOneAccount() (gas: 472094) +UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 494004) +UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 403476) +UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 530509) +UnstakeTest:test_VaultsRegistered() (gas: 62043) +UpgradeTest:test_RevertWhenNotOwner() (gas: 2521503) +UpgradeTest:test_UpgradeStakeManager() (gas: 2794606) +UpgradeTest:test_VaultsRegistered() (gas: 61999) +WithdrawTest:test_CannotWithdrawStakedFunds() (gas: 310176) XPNFTTokenTest:testApproveNotAllowed() (gas: 10500) XPNFTTokenTest:testGetApproved() (gas: 10523) XPNFTTokenTest:testIsApprovedForAll() (gas: 10698) diff --git a/src/RewardsStreamerMP.sol b/src/RewardsStreamerMP.sol index f0af4ee..435477d 100644 --- a/src/RewardsStreamerMP.sol +++ b/src/RewardsStreamerMP.sol @@ -6,6 +6,7 @@ import { ReentrancyGuardUpgradeable } from "@openzeppelin/contracts-upgradeable/ import { Initializable } from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol"; import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol"; import { IStakeManager } from "./interfaces/IStakeManager.sol"; +import { IStakeVault } from "./interfaces/IStakeVault.sol"; import { TrustedCodehashAccess } from "./TrustedCodehashAccess.sol"; // Rewards Streamer with Multiplier Points @@ -16,6 +17,9 @@ contract RewardsStreamerMP is TrustedCodehashAccess, ReentrancyGuardUpgradeable { + error StakingManager__InvalidVault(); + error StakingManager__VaultNotRegistered(); + error StakingManager__VaultAlreadyRegistered(); error StakingManager__AmountCannotBeZero(); error StakingManager__TransferFailed(); error StakingManager__InsufficientBalance(); @@ -57,7 +61,16 @@ contract RewardsStreamerMP is uint256 lockUntil; } - mapping(address account => Account data) public accounts; + mapping(address owner => address[] vault) public vaults; + mapping(address vault => address owner) public vaultOwners; + mapping(address vault => Account data) public accounts; + + modifier onlyRegisteredVault() { + if (!isVaultRegistered(msg.sender)) { + revert StakingManager__VaultNotRegistered(); + } + _; + } modifier onlyNotEmergencyMode() { if (emergencyModeEnabled) { @@ -83,7 +96,50 @@ contract RewardsStreamerMP is _checkOwner(); } - function stake(uint256 amount, uint256 lockPeriod) external onlyTrustedCodehash onlyNotEmergencyMode nonReentrant { + /** + * @notice Check if a vault is registered + * @param vault The address of the vault to check + * @return true if the vault is registered, false otherwise + */ + function isVaultRegistered(address vault) public view returns (bool) { + return vaultOwners[vault] != address(0); + } + + /** + * @notice Registers a vault with its owner. Called by the vault itself during initialization. + * @dev Only callable by contracts with trusted codehash + */ + function registerVault() external onlyTrustedCodehash { + address vault = msg.sender; + address owner = IStakeVault(vault).owner(); + + if (vaultOwners[vault] != address(0)) { + revert StakingManager__VaultAlreadyRegistered(); + } + + // Verify this is a legitimate vault by checking it points to us + if (address(IStakeVault(vault).stakeManager()) != address(this)) { + revert StakingManager__InvalidVault(); + } + + vaultOwners[vault] = owner; + vaults[owner].push(vault); + } + + function getUserVaults(address owner) external view returns (address[] memory) { + return vaults[owner]; + } + + function stake( + uint256 amount, + uint256 lockPeriod + ) + external + onlyTrustedCodehash + onlyNotEmergencyMode + onlyRegisteredVault + nonReentrant + { if (amount == 0) { revert StakingManager__AmountCannotBeZero(); } @@ -127,7 +183,13 @@ contract RewardsStreamerMP is account.lastMPUpdateTime = block.timestamp; } - function lock(uint256 lockPeriod) external onlyTrustedCodehash onlyNotEmergencyMode nonReentrant { + function lock(uint256 lockPeriod) + external + onlyTrustedCodehash + onlyNotEmergencyMode + onlyRegisteredVault + nonReentrant + { if (lockPeriod < MIN_LOCKUP_PERIOD || lockPeriod > MAX_LOCKUP_PERIOD) { revert StakingManager__InvalidLockingPeriod(); } @@ -160,7 +222,13 @@ contract RewardsStreamerMP is account.lastMPUpdateTime = block.timestamp; } - function unstake(uint256 amount) external onlyTrustedCodehash onlyNotEmergencyMode nonReentrant { + function unstake(uint256 amount) + external + onlyTrustedCodehash + onlyNotEmergencyMode + onlyRegisteredVault + nonReentrant + { Account storage account = accounts[msg.sender]; if (amount > account.stakedBalance) { revert StakingManager__InsufficientBalance(); @@ -316,7 +384,7 @@ contract RewardsStreamerMP is lastRewardTime = block.timestamp < rewardEndTime ? block.timestamp : rewardEndTime; } - function _calculateBonusMP(uint256 amount, uint256 lockPeriod) internal view returns (uint256) { + function _calculateBonusMP(uint256 amount, uint256 lockPeriod) internal pure returns (uint256) { uint256 lockMultiplier = (lockPeriod * MAX_MULTIPLIER * SCALE_FACTOR) / MAX_LOCKUP_PERIOD; return amount * lockMultiplier / SCALE_FACTOR; } diff --git a/src/StakeVault.sol b/src/StakeVault.sol index 7ab2581..3d09caa 100644 --- a/src/StakeVault.sol +++ b/src/StakeVault.sol @@ -5,6 +5,7 @@ pragma solidity ^0.8.26; import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol"; import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; import { IStakeManagerProxy } from "./interfaces/IStakeManagerProxy.sol"; +import { IStakeVault } from "./interfaces/IStakeVault.sol"; /** * @title StakeVault @@ -12,7 +13,7 @@ import { IStakeManagerProxy } from "./interfaces/IStakeManagerProxy.sol"; * @notice A contract to secure user stakes and manage staking with IStakeManager. * @dev This contract is owned by the user and allows staking, unstaking, and withdrawing tokens. */ -contract StakeVault is Ownable { +contract StakeVault is IStakeVault, Ownable { error StakeVault__NotEnoughAvailableBalance(); error StakeVault__InvalidDestinationAddress(); error StakeVault__UpdateNotAvailable(); @@ -26,7 +27,7 @@ contract StakeVault is Ownable { //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; - IStakeManagerProxy private stakeManager; + IStakeManagerProxy public stakeManager; address public stakeManagerImplementationAddress; /** @@ -71,6 +72,20 @@ contract StakeVault is Ownable { stakeManagerImplementationAddress = stakeManagerAddress; } + /** + * @notice Registers the vault with the stake manager. + */ + function register() public { + stakeManager.registerVault(); + } + + /** + * @notice Returns the address of the current owner. + */ + function owner() public view override(Ownable, IStakeVault) returns (address) { + return super.owner(); + } + /** * @notice Stake tokens for a specified time. * @param _amount The amount of tokens to stake. diff --git a/src/interfaces/IStakeManager.sol b/src/interfaces/IStakeManager.sol index 6b44e22..0cc62ea 100644 --- a/src/interfaces/IStakeManager.sol +++ b/src/interfaces/IStakeManager.sol @@ -11,6 +11,7 @@ interface IStakeManager is ITrustedCodehashAccess { error StakingManager__StakeIsTooLow(); error StakingManager__NotAllowedToLeave(); + function registerVault() external; function stake(uint256 _amount, uint256 _seconds) external; function lock(uint256 _seconds) external; function unstake(uint256 _amount) external; diff --git a/src/interfaces/IStakeVault.sol b/src/interfaces/IStakeVault.sol new file mode 100644 index 0000000..dcc96c4 --- /dev/null +++ b/src/interfaces/IStakeVault.sol @@ -0,0 +1,10 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.26; + +import { IStakeManagerProxy } from "./IStakeManagerProxy.sol"; + +interface IStakeVault { + function owner() external view returns (address); + function stakeManager() external view returns (IStakeManagerProxy); + function register() external; +} diff --git a/test/RewardsStreamerMP.t.sol b/test/RewardsStreamerMP.t.sol index c4297a1..70147ee 100644 --- a/test/RewardsStreamerMP.t.sol +++ b/test/RewardsStreamerMP.t.sol @@ -32,6 +32,14 @@ contract RewardsStreamerMPTest is Test { address proxy = address(new StakeManagerProxy(impl, initializeData)); streamer = RewardsStreamerMP(proxy); + // Create a temporary vault just to get the codehash + StakeVault tempVault = new StakeVault(address(this), IStakeManagerProxy(address(streamer))); + bytes32 vaultCodeHash = address(tempVault).codehash; + + // Register the codehash before creating any user vaults + vm.prank(admin); + streamer.setTrustedCodehash(vaultCodeHash, true); + address[4] memory accounts = [alice, bob, charlie, dave]; for (uint256 i = 0; i < accounts.length; i++) { // ensure user has tokens @@ -88,6 +96,7 @@ contract RewardsStreamerMPTest is Test { function _createTestVault(address owner) internal returns (StakeVault vault) { vm.prank(owner); vault = new StakeVault(owner, IStakeManagerProxy(address(streamer))); + vault.register(); if (!streamer.isTrustedCodehash(address(vault).codehash)) { vm.prank(admin); @@ -135,6 +144,15 @@ contract RewardsStreamerMPTest is Test { uint256 timeInSeconds = (maxMP * 365 days) / mpPerYear; return timeInSeconds; } + + function test_VaultsRegistered() public view { + address[4] memory accounts = [alice, bob, charlie, dave]; + for (uint256 i = 0; i < accounts.length; i++) { + address[] memory userVaults = streamer.getUserVaults(accounts[i]); + assertEq(userVaults.length, 1, "wrong number of vaults"); + assertEq(userVaults[0], vaults[accounts[i]], "wrong vault address"); + } + } } contract IntegrationTest is RewardsStreamerMPTest { diff --git a/test/StakeVault.test.sol b/test/StakeVault.test.sol index 8316e22..ca32530 100644 --- a/test/StakeVault.test.sol +++ b/test/StakeVault.test.sol @@ -24,10 +24,7 @@ contract StakeVaultTest is Test { function _createTestVault(address owner) internal returns (StakeVault vault) { vm.prank(owner); vault = new StakeVault(owner, IStakeManagerProxy(address(streamer))); - - if (!streamer.isTrustedCodehash(address(vault).codehash)) { - streamer.setTrustedCodehash(address(vault).codehash, true); - } + vault.register(); } function setUp() public virtual { @@ -41,6 +38,14 @@ contract StakeVaultTest is Test { streamer = RewardsStreamerMP(proxy); stakingToken.mint(alice, 10_000e18); + + // Create a temporary vault just to get the codehash + StakeVault tempVault = new StakeVault(address(this), IStakeManagerProxy(address(streamer))); + bytes32 vaultCodeHash = address(tempVault).codehash; + + // Register the codehash before creating any user vaults + streamer.setTrustedCodehash(vaultCodeHash, true); + stakeVault = _createTestVault(alice); vm.prank(alice); diff --git a/test/mocks/StackOverflowStakeManager.sol b/test/mocks/StackOverflowStakeManager.sol index a57e0c9..e76a427 100644 --- a/test/mocks/StackOverflowStakeManager.sol +++ b/test/mocks/StackOverflowStakeManager.sol @@ -65,4 +65,8 @@ contract StackOverflowStakeManager is function getAccount(address _account) external view returns (Account memory) { return accounts[_account]; } + + function registerVault() external override { + // implementation + } }