diff --git a/.gas-report b/.gas-report index d25e428..fb0a0ff 100644 --- a/.gas-report +++ b/.gas-report @@ -15,33 +15,36 @@ | src/RewardsStreamerMP.sol:RewardsStreamerMP contract | | | | | | |------------------------------------------------------|-----------------|-------|--------|-------|---------| | Deployment Cost | Deployment Size | | | | | -| 1216069 | 5542 | | | | | +| 1521308 | 6964 | | | | | | Function Name | min | avg | median | max | # calls | | MAX_LOCKUP_PERIOD | 228 | 228 | 228 | 228 | 22 | -| MAX_MULTIPLIER | 229 | 229 | 229 | 229 | 28 | -| MIN_LOCKUP_PERIOD | 275 | 275 | 275 | 275 | 11 | +| MAX_MULTIPLIER | 296 | 296 | 296 | 296 | 28 | +| MIN_LOCKUP_PERIOD | 297 | 297 | 297 | 297 | 11 | | MP_RATE_PER_YEAR | 231 | 231 | 231 | 231 | 3 | | SCALE_FACTOR | 229 | 229 | 229 | 229 | 39 | -| STAKING_TOKEN | 273 | 273 | 273 | 273 | 128 | +| STAKING_TOKEN | 273 | 273 | 273 | 273 | 180 | | accountedRewards | 395 | 953 | 395 | 2395 | 68 | -| getAccount | 1596 | 1596 | 1596 | 1596 | 65 | -| isTrustedCodehash | 496 | 996 | 496 | 2496 | 128 | -| rewardIndex | 351 | 380 | 351 | 2351 | 68 | -| setTrustedCodehash | 47926 | 47926 | 47926 | 47926 | 32 | +| getAccount | 1625 | 1625 | 1625 | 1625 | 65 | +| getUserVaults | 5159 | 5159 | 5159 | 5159 | 16 | +| rewardIndex | 373 | 402 | 373 | 2373 | 68 | +| setTrustedCodehash | 47948 | 47948 | 47948 | 47948 | 36 | | totalMP | 352 | 352 | 352 | 352 | 71 | -| totalMaxMP | 395 | 395 | 395 | 395 | 71 | +| totalMaxMP | 373 | 373 | 373 | 373 | 71 | | totalStaked | 374 | 374 | 374 | 374 | 71 | -| updateAccountMP | 34654 | 36892 | 37156 | 37156 | 19 | +| updateAccountMP | 34683 | 36921 | 37185 | 37185 | 19 | | updateGlobalState | 30008 | 55588 | 47387 | 80334 | 25 | | src/StakeVault.sol:StakeVault contract | | | | | | |----------------------------------------|-----------------|--------|--------|--------|---------| | Deployment Cost | Deployment Size | | | | | -| 857122 | 4070 | | | | | +| 897733 | 4258 | | | | | | Function Name | min | avg | median | max | # calls | -| stake | 196978 | 234038 | 240671 | 261155 | 46 | -| unstake | 83287 | 111971 | 101308 | 144002 | 13 | +| owner | 2318 | 2318 | 2318 | 2318 | 144 | +| register | 99083 | 99083 | 99083 | 99083 | 144 | +| stake | 199169 | 236229 | 242862 | 263346 | 46 | +| stakeManager | 391 | 391 | 391 | 391 | 144 | +| unstake | 85075 | 114102 | 103543 | 146237 | 13 | | src/XPNFTToken.sol:XPNFTToken contract | | | | | | @@ -117,9 +120,9 @@ | Deployment Cost | Deployment Size | | | | | | 639406 | 3369 | | | | | | Function Name | min | avg | median | max | # calls | -| approve | 46334 | 46343 | 46346 | 46346 | 165 | +| approve | 46334 | 46343 | 46346 | 46346 | 185 | | balanceOf | 561 | 1351 | 561 | 2561 | 286 | -| mint | 51284 | 59028 | 51284 | 68384 | 181 | +| mint | 51284 | 58938 | 51284 | 68384 | 201 | | transfer | 34390 | 48070 | 51490 | 51490 | 10 | diff --git a/.gas-snapshot b/.gas-snapshot index 9467ae2..0c27519 100644 --- a/.gas-snapshot +++ b/.gas-snapshot @@ -1,42 +1,46 @@ -IntegrationTest:testStakeFoo() (gas: 1471979) +IntegrationTest:testStakeFoo() (gas: 1483486) +IntegrationTest:test_VaultsRegistered() (gas: 56341) 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: 56384) RewardsStreamerTest:testStake() (gas: 869874) -StakeTest:test_StakeMultipleAccounts() (gas: 488941) -StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 634452) -StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 801369) -StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 494644) -StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 516001) -StakeTest:test_StakeOneAccount() (gas: 282173) -StakeTest:test_StakeOneAccountAndRewards() (gas: 427680) -StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 488578) -StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 483621) -StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 295784) -StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 295862) -StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 295973) -UnstakeTest:test_StakeMultipleAccounts() (gas: 488963) -UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 634429) -UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 801391) -UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 494621) -UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 516023) -UnstakeTest:test_StakeOneAccount() (gas: 282196) -UnstakeTest:test_StakeOneAccountAndRewards() (gas: 427702) -UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 488600) -UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 483601) -UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 295829) -UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 295862) -UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 295951) -UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 500003) -UnstakeTest:test_UnstakeMultipleAccounts() (gas: 681122) -UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 1002814) -UnstakeTest:test_UnstakeOneAccount() (gas: 474888) -UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 488421) -UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 579871) -UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 509781) +StakeTest:test_StakeMultipleAccounts() (gas: 493447) +StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 638892) +StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 806198) +StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 499160) +StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 520517) +StakeTest:test_StakeOneAccount() (gas: 284393) +StakeTest:test_StakeOneAccountAndRewards() (gas: 429900) +StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 490930) +StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 486040) +StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 298042) +StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 298142) +StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 298231) +StakeTest:test_VaultsRegistered() (gas: 56363) +UnstakeTest:test_StakeMultipleAccounts() (gas: 493403) +UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 638869) +UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 806197) +UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 499159) +UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 520561) +UnstakeTest:test_StakeOneAccount() (gas: 284416) +UnstakeTest:test_StakeOneAccountAndRewards() (gas: 429878) +UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 490974) +UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 486042) +UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 298087) +UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 298142) +UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 298231) +UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 504568) +UnstakeTest:test_UnstakeMultipleAccounts() (gas: 689196) +UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 1013152) +UnstakeTest:test_UnstakeOneAccount() (gas: 480713) +UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 492905) +UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 584355) +UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 514481) +UnstakeTest:test_VaultsRegistered() (gas: 56407) XPNFTTokenTest:testApproveNotAllowed() (gas: 10507) XPNFTTokenTest:testGetApproved() (gas: 10531) XPNFTTokenTest:testIsApprovedForAll() (gas: 10705) diff --git a/src/RewardsStreamerMP.sol b/src/RewardsStreamerMP.sol index b173c5d..37a9edc 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(); @@ -41,7 +45,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(); + } + _; + } constructor(address _owner, address _stakingToken, address _rewardToken) TrustedCodehashAccess(_owner) { STAKING_TOKEN = IERC20(_stakingToken); @@ -49,7 +62,41 @@ contract RewardsStreamerMP is IStakeManager, TrustedCodehashAccess, ReentrancyGu lastMPUpdatedTime = block.timestamp; } - function stake(uint256 amount, uint256 lockPeriod) external onlyTrustedCodehash 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 onlyRegisteredVault nonReentrant { if (amount == 0) { revert StakingManager__AmountCannotBeZero(); } @@ -99,7 +146,7 @@ contract RewardsStreamerMP is IStakeManager, TrustedCodehashAccess, ReentrancyGu account.lastMPUpdateTime = block.timestamp; } - function unstake(uint256 amount) external onlyTrustedCodehash nonReentrant { + function unstake(uint256 amount) external onlyTrustedCodehash onlyRegisteredVault nonReentrant { Account storage account = accounts[msg.sender]; if (amount > account.stakedBalance) { revert StakingManager__InsufficientBalance(); diff --git a/src/StakeVault.sol b/src/StakeVault.sol index 2c63ee3..58b6e14 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(); @@ -23,7 +24,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. @@ -51,6 +52,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 97d1574..df5972a 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 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 2fe79df..9de9648 100644 --- a/test/RewardsStreamerMP.t.sol +++ b/test/RewardsStreamerMP.t.sol @@ -24,6 +24,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 @@ -85,10 +92,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 { @@ -125,6 +129,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 {