From 9d45b6414349b858b62fe5d7a1c15d7108dd70a3 Mon Sep 17 00:00:00 2001 From: r4bbit <445106+0x-r4bbit@users.noreply.github.com> Date: Fri, 25 Oct 2024 10:21:51 +0200 Subject: [PATCH] feat(StakeManager): add capabilities to register vaults This commit introduces changes related to vault registrations in the stake manager. The stake manager needs to keep track of the vaults a users creates so it can aggregate accumulated MP across vaults for any given user. The `StakeVault` now comes with a `register()` function which needs to be called to register itself with the stake manager. `StakeManager` has a new `onlyRegisteredVault` modifier that ensures only registered vaults can actually `stake` and `unstake`. Closes #70 --- .gas-report | 41 ++++++++------- .gas-snapshot | 88 +++++++++++++++++--------------- src/RewardsStreamerMP.sol | 78 ++++++++++++++++++++++++++-- src/StakeVault.sol | 19 ++++++- src/interfaces/IStakeManager.sol | 3 ++ src/interfaces/IStakeVault.sol | 10 ++++ test/RewardsStreamerMP.t.sol | 21 ++++++-- 7 files changed, 189 insertions(+), 71 deletions(-) create mode 100644 src/interfaces/IStakeVault.sol diff --git a/.gas-report b/.gas-report index 7179b6a..0dfee06 100644 --- a/.gas-report +++ b/.gas-report @@ -15,37 +15,40 @@ | src/RewardsStreamerMP.sol:RewardsStreamerMP contract | | | | | | |------------------------------------------------------|-----------------|-------|--------|-------|---------| | Deployment Cost | Deployment Size | | | | | -| 1425736 | 6528 | | | | | +| 1760847 | 8090 | | | | | | 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 | +| MAX_LOCKUP_PERIOD | 273 | 273 | 273 | 273 | 23 | +| MAX_MULTIPLIER | 296 | 296 | 296 | 296 | 30 | +| MIN_LOCKUP_PERIOD | 297 | 297 | 297 | 297 | 11 | | MP_RATE_PER_YEAR | 231 | 231 | 231 | 231 | 3 | | SCALE_FACTOR | 295 | 295 | 295 | 295 | 41 | -| STAKING_TOKEN | 273 | 273 | 273 | 273 | 172 | +| STAKING_TOKEN | 273 | 273 | 273 | 273 | 245 | | accountedRewards | 351 | 906 | 351 | 2351 | 72 | | 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 | +| getAccount | 1625 | 1625 | 1625 | 1625 | 71 | +| getUserVaults | 5181 | 5181 | 5181 | 5181 | 24 | +| rewardIndex | 395 | 422 | 395 | 2395 | 72 | +| setTrustedCodehash | 47948 | 47948 | 47948 | 47948 | 49 | +| totalMP | 352 | 352 | 352 | 352 | 75 | | totalMaxMP | 351 | 351 | 351 | 351 | 75 | -| totalStaked | 330 | 330 | 330 | 330 | 75 | -| updateAccountMP | 36758 | 38996 | 39260 | 39260 | 19 | +| totalStaked | 352 | 352 | 352 | 352 | 75 | +| updateAccountMP | 36809 | 39047 | 39311 | 39311 | 19 | | updateGlobalState | 32134 | 60366 | 49513 | 82460 | 28 | | src/StakeVault.sol:StakeVault contract | | | | | | |----------------------------------------|-----------------|--------|--------|--------|---------| | Deployment Cost | Deployment Size | | | | | -| 1095864 | 5202 | | | | | +| 1136448 | 5391 | | | | | | Function Name | min | avg | median | max | # calls | -| emergencyExit | 31410 | 43924 | 43160 | 60260 | 7 | -| lock | 38362 | 60516 | 42741 | 100445 | 3 | -| stake | 199213 | 236948 | 242906 | 263435 | 55 | -| unstake | 84988 | 113999 | 103434 | 146128 | 13 | +| emergencyExit | 31432 | 43946 | 43182 | 60282 | 7 | +| lock | 40553 | 62707 | 44932 | 102636 | 3 | +| owner | 2363 | 2363 | 2363 | 2363 | 196 | +| register | 99106 | 99106 | 99106 | 99106 | 196 | +| stake | 201338 | 239073 | 245031 | 265560 | 55 | +| stakeManager | 369 | 369 | 369 | 369 | 196 | +| unstake | 86776 | 116130 | 105669 | 148363 | 13 | | src/XPNFTToken.sol:XPNFTToken contract | | | | | | @@ -121,9 +124,9 @@ | Deployment Cost | Deployment Size | | | | | | 639406 | 3369 | | | | | | Function Name | min | avg | median | max | # calls | -| approve | 46334 | 46343 | 46346 | 46346 | 220 | +| approve | 46334 | 46343 | 46346 | 46346 | 250 | | balanceOf | 561 | 1381 | 561 | 2561 | 334 | -| mint | 51284 | 58817 | 51284 | 68384 | 236 | +| mint | 51284 | 58739 | 51284 | 68384 | 266 | | transfer | 34390 | 48859 | 51490 | 51490 | 13 | diff --git a/.gas-snapshot b/.gas-snapshot index 94d93e7..06a1c56 100644 --- a/.gas-snapshot +++ b/.gas-snapshot @@ -1,53 +1,59 @@ 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_CannotLeaveBeforeEmergencyMode() (gas: 285403) +EmergencyExitTest:test_EmergencyExitBasic() (gas: 382116) +EmergencyExitTest:test_EmergencyExitMultipleUsers() (gas: 793133) +EmergencyExitTest:test_EmergencyExitToAlternateAddress() (gas: 519631) +EmergencyExitTest:test_EmergencyExitWithLock() (gas: 372473) +EmergencyExitTest:test_EmergencyExitWithRewards() (gas: 509508) 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_VaultsRegistered() (gas: 56670) +IntegrationTest:testStakeFoo() (gas: 1502203) +IntegrationTest:test_VaultsRegistered() (gas: 56494) +LockTest:test_LockFailsWithInvalidPeriod() (gas: 294494) +LockTest:test_LockFailsWithNoStake() (gas: 55548) +LockTest:test_LockWithoutPriorLock() (gas: 387713) +LockTest:test_VaultsRegistered() (gas: 56472) NFTMetadataGeneratorSVGTest:testGenerateMetadata() (gas: 92874) NFTMetadataGeneratorSVGTest:testSetImageStrings() (gas: 60081) NFTMetadataGeneratorSVGTest:testSetImageStringsRevert() (gas: 35818) NFTMetadataGeneratorURLTest:testGenerateMetadata() (gas: 109345) NFTMetadataGeneratorURLTest:testSetBaseURL() (gas: 50653) NFTMetadataGeneratorURLTest:testSetBaseURLRevert() (gas: 35993) +RewardsStreamerMPTest:test_VaultsRegistered() (gas: 56472) 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: 497653) +StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 645225) +StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 823122) +StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 503853) +StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 525300) +StakeTest:test_StakeOneAccount() (gas: 286497) +StakeTest:test_StakeOneAccountAndRewards() (gas: 434064) +StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 501406) +StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 496605) +StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 300435) +StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 300467) +StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 300578) +StakeTest:test_VaultsRegistered() (gas: 56451) +UnstakeTest:test_StakeMultipleAccounts() (gas: 497675) +UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 645247) +UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 823099) +UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 503830) +UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 525322) +UnstakeTest:test_StakeOneAccount() (gas: 286520) +UnstakeTest:test_StakeOneAccountAndRewards() (gas: 434086) +UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 501428) +UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 496585) +UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 300435) +UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 300467) +UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 300556) +UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 513208) +UnstakeTest:test_UnstakeMultipleAccounts() (gas: 696829) +UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 1024775) +UnstakeTest:test_UnstakeOneAccount() (gas: 486087) +UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 501254) +UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 590580) +UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 523428) +UnstakeTest:test_VaultsRegistered() (gas: 56473) XPNFTTokenTest:testApproveNotAllowed() (gas: 10507) XPNFTTokenTest:testGetApproved() (gas: 10531) XPNFTTokenTest:testIsApprovedForAll() (gas: 10705) diff --git a/src/RewardsStreamerMP.sol b/src/RewardsStreamerMP.sol index c5acb81..9f8b3c7 100644 --- a/src/RewardsStreamerMP.sol +++ b/src/RewardsStreamerMP.sol @@ -4,10 +4,14 @@ pragma solidity ^0.8.26; import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; import { ReentrancyGuard } from "@openzeppelin/contracts/utils/ReentrancyGuard.sol"; import { IStakeManager } from "./interfaces/IStakeManager.sol"; +import { IStakeVault } from "./interfaces/IStakeVault.sol"; import { TrustedCodehashAccess } from "./TrustedCodehashAccess.sol"; // Rewards Streamer with Multiplier Points contract RewardsStreamerMP is IStakeManager, TrustedCodehashAccess, ReentrancyGuard { + error StakingManager__InvalidVault(); + error StakingManager__VaultNotRegistered(); + error StakingManager__VaultAlreadyRegistered(); error StakingManager__AmountCannotBeZero(); error StakingManager__TransferFailed(); error StakingManager__InsufficientBalance(); @@ -44,7 +48,16 @@ contract RewardsStreamerMP is IStakeManager, TrustedCodehashAccess, ReentrancyGu 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) { @@ -59,7 +72,50 @@ contract RewardsStreamerMP is IStakeManager, TrustedCodehashAccess, ReentrancyGu lastMPUpdatedTime = block.timestamp; } - 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(); } @@ -108,7 +164,13 @@ contract RewardsStreamerMP is IStakeManager, TrustedCodehashAccess, ReentrancyGu 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(); } @@ -141,7 +203,13 @@ contract RewardsStreamerMP is IStakeManager, TrustedCodehashAccess, ReentrancyGu 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(); @@ -227,7 +295,7 @@ contract RewardsStreamerMP is IStakeManager, TrustedCodehashAccess, ReentrancyGu } } - 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 efb0322..cede6bd 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 { IStakeManager } from "./interfaces/IStakeManager.sol"; +import { IStakeVault } from "./interfaces/IStakeVault.sol"; /** * @title StakeVault @@ -12,7 +13,7 @@ import { IStakeManager } from "./interfaces/IStakeManager.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__NoEnoughAvailableBalance(); error StakeVault__InvalidDestinationAddress(); error StakeVault__UpdateNotAvailable(); @@ -24,7 +25,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; - IStakeManager private stakeManager; + IStakeManager public stakeManager; /** * @dev Emitted when tokens are staked. @@ -52,6 +53,20 @@ contract StakeVault is Ownable { stakeManager = _stakeManager; } + /** + * @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 6bd84b2..e796513 100644 --- a/src/interfaces/IStakeManager.sol +++ b/src/interfaces/IStakeManager.sol @@ -5,11 +5,14 @@ import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; import { ITrustedCodehashAccess } from "./ITrustedCodehashAccess.sol"; interface IStakeManager is ITrustedCodehashAccess { + error StakeManager__InvalidRegistered(); + error StakeManager__VaultAlreadyRegistered(); error StakeManager__FundsLocked(); error StakeManager__InvalidLockTime(); error StakeManager__InsufficientFunds(); error StakeManager__StakeIsTooLow(); + 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..8f159a2 --- /dev/null +++ b/src/interfaces/IStakeVault.sol @@ -0,0 +1,10 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.26; + +import { IStakeManager } from "./IStakeManager.sol"; + +interface IStakeVault { + function owner() external view returns (address); + function stakeManager() external view returns (IStakeManager); + function register() external; +} diff --git a/test/RewardsStreamerMP.t.sol b/test/RewardsStreamerMP.t.sol index 5800f27..614e3c5 100644 --- a/test/RewardsStreamerMP.t.sol +++ b/test/RewardsStreamerMP.t.sol @@ -25,6 +25,13 @@ contract RewardsStreamerMPTest is Test { stakingToken = new MockToken("Staking Token", "ST"); streamer = new RewardsStreamerMP(address(this), address(stakingToken), address(rewardToken)); + // Create a temporary vault just to get the codehash + StakeVault tempVault = new StakeVault(address(this), streamer); + bytes32 vaultCodeHash = address(tempVault).codehash; + + // Register the codehash before creating any user vaults + streamer.setTrustedCodehash(vaultCodeHash, true); + address[4] memory accounts = [alice, bob, charlie, dave]; for (uint256 i = 0; i < accounts.length; i++) { // ensure user has tokens @@ -87,10 +94,7 @@ contract RewardsStreamerMPTest is Test { function _createTestVault(address owner) internal returns (StakeVault vault) { vm.prank(owner); vault = new StakeVault(owner, streamer); - - if (!streamer.isTrustedCodehash(address(vault).codehash)) { - streamer.setTrustedCodehash(address(vault).codehash, true); - } + vault.register(); } function _stake(address account, uint256 amount, uint256 lockupTime) public { @@ -133,6 +137,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 {