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-snapshot

@@ -1,67 +1,69 @@
-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_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_CannotEnableEmergencyModeTwice() (gas: 92646)
+EmergencyExitTest:test_CannotLeaveBeforeEmergencyMode() (gas: 297824)
+EmergencyExitTest:test_EmergencyExitBasic() (gas: 384518)
+EmergencyExitTest:test_EmergencyExitMultipleUsers() (gas: 659358)
+EmergencyExitTest:test_EmergencyExitToAlternateAddress() (gas: 392437)
+EmergencyExitTest:test_EmergencyExitWithLock() (gas: 392015)
+EmergencyExitTest:test_EmergencyExitWithRewards() (gas: 377401)
+EmergencyExitTest:test_OnlyOwnerCanEnableEmergencyMode() (gas: 39408)
+IntegrationTest:testStakeFoo() (gas: 1179318)
+LeaveTest:test_LeaveShouldProperlyUpdateAccounting() (gas: 2927052)
+LeaveTest:test_RevertWhenStakeManagerIsTrusted() (gas: 294955)
+LeaveTest:test_TrustNewStakeManager() (gas: 3004493)
+LockTest:test_LockFailsWithInvalidPeriod() (gas: 310018)
+LockTest:test_LockFailsWithNoStake() (gas: 63598)
+LockTest:test_LockWithoutPriorLock() (gas: 391137)
+MaliciousUpgradeTest:test_UpgradeStackOverflowStakeManager() (gas: 1745462)
+MultipleVaultsStakeTest:test_StakeMultipleVaults() (gas: 717207)
 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)
-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:testRewardsBalanceOf() (gas: 670984)
+RewardsStreamerMP_RewardsTest:testSetRewards() (gas: 160214)
+RewardsStreamerMP_RewardsTest:testSetRewards_RevertsBadAmount() (gas: 39323)
+RewardsStreamerMP_RewardsTest:testSetRewards_RevertsBadDuration() (gas: 39346)
+RewardsStreamerMP_RewardsTest:testSetRewards_RevertsNotAuthorized() (gas: 39359)
+RewardsStreamerMP_RewardsTest:testTotalRewardsSupply() (gas: 611915)
 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: 494656)
+StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 500594)
+StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 831165)
+StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 517651)
+StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 539579)
+StakeTest:test_StakeOneAccount() (gas: 277040)
+StakeTest:test_StakeOneAccountAndRewards() (gas: 283009)
+StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 499939)
+StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 496276)
+StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 301895)
+StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 301884)
+StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 301951)
 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: 494678)
+UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 500594)
+UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 831142)
+UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 517673)
+UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 539601)
+UnstakeTest:test_StakeOneAccount() (gas: 277063)
+UnstakeTest:test_StakeOneAccountAndRewards() (gas: 283031)
+UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 499961)
+UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 496256)
+UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 301895)
+UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 301884)
+UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 301995)
+UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 542969)
+UnstakeTest:test_UnstakeMultipleAccounts() (gas: 693417)
+UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 787224)
+UnstakeTest:test_UnstakeOneAccount() (gas: 473460)
+UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 495130)
+UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 404531)
+UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 531635)
+UpgradeTest:test_RevertWhenNotOwner() (gas: 2568028)
+UpgradeTest:test_UpgradeStakeManager() (gas: 2841596)
+VaultRegistrationTest:test_VaultRegistration() (gas: 62211)
+WithdrawTest:test_CannotWithdrawStakedFunds() (gas: 310679)
 XPNFTTokenTest:testApproveNotAllowed() (gas: 10500)
 XPNFTTokenTest:testGetApproved() (gas: 10523)
 XPNFTTokenTest:testIsApprovedForAll() (gas: 10698)

certora/specs/EmergencyMode.spec

@@ -29,7 +29,14 @@ definition isViewFunction(method f) returns bool = (
   f.selector == sig:streamer.rewardAmount().selector ||
   f.selector == sig:streamer.totalRewardsAccrued().selector ||
   f.selector == sig:streamer.rewardStartTime().selector ||
-  f.selector == sig:streamer.rewardEndTime().selector
+  f.selector == sig:streamer.rewardEndTime().selector ||
+  f.selector == sig:streamer.getUserTotalMP(address).selector ||
+  f.selector == sig:streamer.getUserTotalMaxMP(address).selector ||
+  f.selector == sig:streamer.getUserTotalStakedBalance(address).selector ||
+  f.selector == sig:streamer.vaults(address,uint256).selector ||
+  f.selector == sig:streamer.vaultOwners(address).selector ||
+  f.selector == sig:streamer.registerVault().selector ||
+  f.selector == sig:streamer.getUserVaults(address).selector
 );
 
 definition isOwnableFunction(method f) returns bool = (

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 vault => Account data) public accounts;
+    mapping(address owner => address[] vault) public vaults;
+    mapping(address vault => address owner) public vaultOwners;
+
+    modifier onlyRegisteredVault() {
+        if (vaultOwners[msg.sender] == address(0)) {
+            revert StakingManager__VaultNotRegistered();
+        }
+        _;
+    }
 
     modifier onlyNotEmergencyMode() {
         if (emergencyModeEnabled) {
@@ -83,7 +96,95 @@ contract RewardsStreamerMP is
         _checkOwner();
     }
 
-    function stake(uint256 amount, uint256 lockPeriod) external onlyTrustedCodehash onlyNotEmergencyMode nonReentrant {
+    /**
+     * @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 stakeManager
+        if (address(IStakeVault(vault).stakeManager()) != address(this)) {
+            revert StakingManager__InvalidVault();
+        }
+
+        vaultOwners[vault] = owner;
+        vaults[owner].push(vault);
+    }
+
+    /**
+     * @notice Get the vaults owned by a user
+     * @param user The address of the user
+     * @return The vaults owned by the user
+     */
+    function getUserVaults(address user) external view returns (address[] memory) {
+        return vaults[user];
+    }
+
+    /**
+     * @notice Get the total multiplier points for a user
+     * @dev Iterates over all vaults owned by the user and sums the multiplier points
+     * @param user The address of the user
+     * @return The total multiplier points for the user
+     */
+    function getUserTotalMP(address user) external view returns (uint256) {
+        address[] memory userVaults = vaults[user];
+        uint256 userTotalMP = 0;
+
+        for (uint256 i = 0; i < userVaults.length; i++) {
+            Account storage account = accounts[userVaults[i]];
+            userTotalMP += account.accountMP + _getAccountAccruedMP(account);
+        }
+        return userTotalMP;
+    }
+
+    /**
+     * @notice Get the total maximum multiplier points for a user
+     * @dev Iterates over all vaults owned by the user and sums the maximum multiplier points
+     * @param user The address of the user
+     * @return The total maximum multiplier points for the user
+     */
+    function getUserTotalMaxMP(address user) external view returns (uint256) {
+        address[] memory userVaults = vaults[user];
+        uint256 userTotalMaxMP = 0;
+
+        for (uint256 i = 0; i < userVaults.length; i++) {
+            userTotalMaxMP += accounts[userVaults[i]].maxMP;
+        }
+        return userTotalMaxMP;
+    }
+
+    /**
+     * @notice Get the total staked balance for a user
+     * @dev Iterates over all vaults owned by the user and sums the staked balances
+     * @param user The address of the user
+     * @return The total staked balance for the user
+     */
+    function getUserTotalStakedBalance(address user) external view returns (uint256) {
+        address[] memory userVaults = vaults[user];
+        uint256 userTotalStake = 0;
+
+        for (uint256 i = 0; i < userVaults.length; i++) {
+            userTotalStake += accounts[userVaults[i]].stakedBalance;
+        }
+        return userTotalStake;
+    }
+
+    function stake(
+        uint256 amount,
+        uint256 lockPeriod
+    )
+        external
+        onlyTrustedCodehash
+        onlyNotEmergencyMode
+        onlyRegisteredVault
+        nonReentrant
+    {
         if (amount == 0) {
             revert StakingManager__AmountCannotBeZero();
         }
@@ -127,7 +228,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 +267,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,29 +429,32 @@ 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;
     }
 
-    function _updateAccountMP(address accountAddress) internal {
-        Account storage account = accounts[accountAddress];
-
+    function _getAccountAccruedMP(Account storage account) internal view returns (uint256) {
         if (account.maxMP == 0 || account.stakedBalance == 0) {
-            account.lastMPUpdateTime = block.timestamp;
-            return;
+            return 0;
         }
 
         uint256 timeDiff = block.timestamp - account.lastMPUpdateTime;
         if (timeDiff == 0) {
-            return;
+            return 0;
         }
 
         uint256 accruedMP = (timeDiff * account.stakedBalance * MP_RATE_PER_YEAR) / (365 days * SCALE_FACTOR);
 
         if (account.accountMP + accruedMP > account.maxMP) {
             accruedMP = account.maxMP - account.accountMP;
         }
+        return accruedMP;
+    }
+
+    function _updateAccountMP(address accountAddress) internal {
+        Account storage account = accounts[accountAddress];
+        uint256 accruedMP = _getAccountAccruedMP(account);
 
         account.accountMP += accruedMP;
         account.lastMPUpdateTime = block.timestamp;

src/StakeVault.sol

@@ -5,14 +5,15 @@ 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
  * @author Ricardo Guilherme Schmidt <[email protected]>
  * @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.

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;