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,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)

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;
     }

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;

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;
+}

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 {

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);

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
+    }
 }