refactor: move vault registrations into dedicated registry contract

.gas-snapshot

@@ -1,69 +1,69 @@
-EmergencyExitTest:test_CannotEnableEmergencyModeTwice() (gas: 92690)
-EmergencyExitTest:test_CannotLeaveBeforeEmergencyMode() (gas: 297868)
-EmergencyExitTest:test_EmergencyExitBasic() (gas: 384605)
-EmergencyExitTest:test_EmergencyExitMultipleUsers() (gas: 659510)
-EmergencyExitTest:test_EmergencyExitToAlternateAddress() (gas: 392525)
-EmergencyExitTest:test_EmergencyExitWithLock() (gas: 392081)
-EmergencyExitTest:test_EmergencyExitWithRewards() (gas: 377466)
-EmergencyExitTest:test_OnlyOwnerCanEnableEmergencyMode() (gas: 39430)
-IntegrationTest:testStakeFoo() (gas: 1179859)
-LeaveTest:test_LeaveShouldProperlyUpdateAccounting() (gas: 2941805)
-LeaveTest:test_RevertWhenStakeManagerIsTrusted() (gas: 294977)
-LeaveTest:test_TrustNewStakeManager() (gas: 3019253)
-LockTest:test_LockFailsWithInvalidPeriod() (gas: 310062)
-LockTest:test_LockFailsWithNoStake() (gas: 63620)
-LockTest:test_LockWithoutPriorLock() (gas: 391268)
-MaliciousUpgradeTest:test_UpgradeStackOverflowStakeManager() (gas: 1745438)
-MultipleVaultsStakeTest:test_StakeMultipleVaults() (gas: 717206)
+EmergencyExitTest:test_CannotEnableEmergencyModeTwice() (gas: 92598)
+EmergencyExitTest:test_CannotLeaveBeforeEmergencyMode() (gas: 303185)
+EmergencyExitTest:test_EmergencyExitBasic() (gas: 389831)
+EmergencyExitTest:test_EmergencyExitMultipleUsers() (gas: 670073)
+EmergencyExitTest:test_EmergencyExitToAlternateAddress() (gas: 397772)
+EmergencyExitTest:test_EmergencyExitWithLock() (gas: 397352)
+EmergencyExitTest:test_EmergencyExitWithRewards() (gas: 382738)
+EmergencyExitTest:test_OnlyOwnerCanEnableEmergencyMode() (gas: 39384)
+IntegrationTest:testStakeFoo() (gas: 1205695)
+LeaveTest:test_LeaveShouldProperlyUpdateAccounting() (gas: 3044358)
+LeaveTest:test_RevertWhenStakeManagerIsTrusted() (gas: 300338)
+LeaveTest:test_TrustNewStakeManager() (gas: 3121863)
+LockTest:test_LockFailsWithInvalidPeriod() (gas: 320790)
+LockTest:test_LockFailsWithNoStake() (gas: 69009)
+LockTest:test_LockWithoutPriorLock() (gas: 401855)
+MaliciousUpgradeTest:test_UpgradeStackOverflowStakeManager() (gas: 1750751)
+MultipleVaultsStakeTest:test_StakeMultipleVaults() (gas: 739653)
 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: 670943)
-RewardsStreamerMP_RewardsTest:testSetRewards() (gas: 160368)
-RewardsStreamerMP_RewardsTest:testSetRewards_RevertsBadAmount() (gas: 39345)
-RewardsStreamerMP_RewardsTest:testSetRewards_RevertsBadDuration() (gas: 39368)
-RewardsStreamerMP_RewardsTest:testSetRewards_RevertsNotAuthorized() (gas: 39381)
-RewardsStreamerMP_RewardsTest:testTotalRewardsSupply() (gas: 610886)
+RewardsStreamerMP_RewardsTest:testRewardsBalanceOf() (gas: 676244)
+RewardsStreamerMP_RewardsTest:testSetRewards() (gas: 160279)
+RewardsStreamerMP_RewardsTest:testSetRewards_RevertsBadAmount() (gas: 39409)
+RewardsStreamerMP_RewardsTest:testSetRewards_RevertsBadDuration() (gas: 39345)
+RewardsStreamerMP_RewardsTest:testSetRewards_RevertsNotAuthorized() (gas: 39358)
+RewardsStreamerMP_RewardsTest:testTotalRewardsSupply() (gas: 616724)
 RewardsStreamerTest:testStake() (gas: 869181)
-StakeTest:test_StakeMultipleAccounts() (gas: 494743)
-StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 500680)
-StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 831250)
-StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 517691)
-StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 539574)
-StakeTest:test_StakeOneAccount() (gas: 277083)
-StakeTest:test_StakeOneAccountAndRewards() (gas: 283051)
-StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 499980)
-StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 496361)
-StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 301892)
-StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 301881)
-StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 301948)
+StakeTest:test_StakeMultipleAccounts() (gas: 505375)
+StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 511291)
+StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 841637)
+StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 528373)
+StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 550256)
+StakeTest:test_StakeOneAccount() (gas: 282377)
+StakeTest:test_StakeOneAccountAndRewards() (gas: 288302)
+StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 505138)
+StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 501475)
+StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 307256)
+StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 307200)
+StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 307311)
 StakingTokenTest:testStakeToken() (gas: 10422)
-UnstakeTest:test_StakeMultipleAccounts() (gas: 494765)
-UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 500680)
-UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 831227)
-UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 517713)
-UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 539596)
-UnstakeTest:test_StakeOneAccount() (gas: 277106)
-UnstakeTest:test_StakeOneAccountAndRewards() (gas: 283073)
-UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 500002)
-UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 496341)
-UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 301892)
-UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 301881)
-UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 301992)
-UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 543052)
-UnstakeTest:test_UnstakeMultipleAccounts() (gas: 693591)
-UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 787431)
-UnstakeTest:test_UnstakeOneAccount() (gas: 473557)
-UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 495171)
-UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 404616)
-UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 531627)
-UpgradeTest:test_RevertWhenNotOwner() (gas: 2582744)
-UpgradeTest:test_UpgradeStakeManager() (gas: 2856354)
-VaultRegistrationTest:test_VaultRegistration() (gas: 62035)
-WithdrawTest:test_CannotWithdrawStakedFunds() (gas: 310613)
+UnstakeTest:test_StakeMultipleAccounts() (gas: 505374)
+UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 511268)
+UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 841614)
+UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 528372)
+UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 550300)
+UnstakeTest:test_StakeOneAccount() (gas: 282400)
+UnstakeTest:test_StakeOneAccountAndRewards() (gas: 288346)
+UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 505182)
+UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 501477)
+UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 307278)
+UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 307200)
+UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 307311)
+UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 553550)
+UnstakeTest:test_UnstakeMultipleAccounts() (gas: 714853)
+UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 811813)
+UnstakeTest:test_UnstakeOneAccount() (gas: 487377)
+UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 505804)
+UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 415205)
+UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 542243)
+UpgradeTest:test_RevertWhenNotOwner() (gas: 2680084)
+UpgradeTest:test_UpgradeStakeManager() (gas: 2958999)
+VaultRegistrationTest:test_VaultRegistration() (gas: 55423)
+WithdrawTest:test_CannotWithdrawStakedFunds() (gas: 315973)
 XPNFTTokenTest:testApproveNotAllowed() (gas: 10500)
 XPNFTTokenTest:testGetApproved() (gas: 10523)
 XPNFTTokenTest:testIsApprovedForAll() (gas: 10698)

certora/confs/RewardsStreamerMP.conf

@@ -1,9 +1,11 @@
 {
   "files": [
+    "src/StakeVaultRegistry.sol",
     "src/RewardsStreamerMP.sol",
     "certora/helpers/ERC20A.sol"
   ],
   "link" : [
+    "RewardsStreamerMP:stakeVaultRegistry=StakeVaultRegistry",
     "RewardsStreamerMP:STAKING_TOKEN=ERC20A"
   ],
   "msg": "Verifying RewardsStreamerMP.sol",

certora/specs/EmergencyMode.spec

@@ -20,6 +20,7 @@ definition isViewFunction(method f) returns bool = (
   f.selector == sig:streamer.totalMP().selector ||
   f.selector == sig:streamer.accounts(address).selector ||
   f.selector == sig:streamer.emergencyModeEnabled().selector ||
+  f.selector == sig:streamer.stakeVaultRegistry().selector ||
   f.selector == sig:streamer.getStakedBalance(address).selector ||
   f.selector == sig:streamer.getAccount(address).selector ||
   f.selector == sig:streamer.rewardsBalanceOf(address).selector ||
@@ -33,11 +34,7 @@ definition isViewFunction(method f) returns bool = (
   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.getUserVaults(address).selector ||
-  f.selector == sig:streamer.isVaultRegistered(address).selector ||
-  f.selector == sig:streamer.registerVault().selector ||
-  f.selector == sig:streamer.vaultOwners(address).selector ||
-  f.selector == sig:streamer.vaults(address,uint256).selector
+  f.selector == sig:streamer.registerVault().selector
 );
 
 definition isOwnableFunction(method f) returns bool = (
@@ -52,7 +49,7 @@ definition isTrustedCodehashAccessFunction(method f) returns bool = (
 );
 
 definition isInitializerFunction(method f) returns bool = (
-  f.selector == sig:streamer.initialize(address,address).selector
+  f.selector == sig:streamer.initialize(address,address,address).selector
 );
 
 definition isUUPSUpgradeableFunction(method f) returns bool = (

src/RewardsStreamerMP.sol

@@ -7,6 +7,7 @@ import { Initializable } from "@openzeppelin/contracts-upgradeable/proxy/utils/I
 import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
 import { IStakeManager } from "./interfaces/IStakeManager.sol";
 import { IStakeVault } from "./interfaces/IStakeVault.sol";
+import { IStakeVaultRegistry } from "./interfaces/IStakeVaultRegistry.sol";
 import { TrustedCodehashAccess } from "./TrustedCodehashAccess.sol";
 
 // Rewards Streamer with Multiplier Points
@@ -30,6 +31,8 @@ contract RewardsStreamerMP is
     error StakingManager__EmergencyModeEnabled();
     error StakingManager__DurationCannotBeZero();
 
+    IStakeVaultRegistry public stakeVaultRegistry;
+
     IERC20 public STAKING_TOKEN;
 
     uint256 public constant SCALE_FACTOR = 1e18;
@@ -61,12 +64,10 @@ contract RewardsStreamerMP is
         uint256 lockUntil;
     }
 
-    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)) {
+        if (stakeVaultRegistry.vaultOwners(msg.sender) == address(0)) {
             revert StakingManager__VaultNotRegistered();
         }
         _;
@@ -83,11 +84,12 @@ contract RewardsStreamerMP is
         _disableInitializers();
     }
 
-    function initialize(address _owner, address _stakingToken) public initializer {
+    function initialize(address _owner, address _stakeVaultRegistry, address _stakingToken) public initializer {
         __TrustedCodehashAccess_init(_owner);
         __UUPSUpgradeable_init();
         __ReentrancyGuard_init();
 
+        stakeVaultRegistry = IStakeVaultRegistry(_stakeVaultRegistry);
         STAKING_TOKEN = IERC20(_stakingToken);
         lastMPUpdatedTime = block.timestamp;
     }
@@ -96,41 +98,14 @@ contract RewardsStreamerMP is
         _checkOwner();
     }
 
-    /**
-     * @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);
-    }
-
-    /**
-     * @notice Get all vaults owned by an address
-     */
-    function getUserVaults(address owner) external view returns (address[] memory) {
-        return vaults[owner];
+        stakeVaultRegistry.register(owner, vault);
     }
 
     /**
@@ -140,7 +115,7 @@ contract RewardsStreamerMP is
      * @return The total multiplier points for the user
      */
     function getUserTotalMP(address user) external view returns (uint256) {
-        address[] memory userVaults = vaults[user];
+        address[] memory userVaults = stakeVaultRegistry.vaultsOf(user);
         uint256 userTotalMP = 0;
 
         for (uint256 i = 0; i < userVaults.length; i++) {
@@ -157,7 +132,7 @@ contract RewardsStreamerMP is
      * @return The total maximum multiplier points for the user
      */
     function getUserTotalMaxMP(address user) external view returns (uint256) {
-        address[] memory userVaults = vaults[user];
+        address[] memory userVaults = stakeVaultRegistry.vaultsOf(user);
         uint256 userTotalMaxMP = 0;
 
         for (uint256 i = 0; i < userVaults.length; i++) {
@@ -173,7 +148,7 @@ contract RewardsStreamerMP is
      * @return The total staked balance for the user
      */
     function getUserTotalStakedBalance(address user) external view returns (uint256) {
-        address[] memory userVaults = vaults[user];
+        address[] memory userVaults = stakeVaultRegistry.vaultsOf(user);
         uint256 userTotalStake = 0;
 
         for (uint256 i = 0; i < userVaults.length; i++) {

src/StakeVaultRegistry.sol

@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.26;
+
+import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
+import { IStakeVaultRegistry } from "./interfaces/IStakeVaultRegistry.sol";
+import { IStakeVault } from "./interfaces/IStakeVault.sol";
+
+contract StakeVaultRegistry is IStakeVaultRegistry, Ownable {
+    error StakeManagerVaultRegistry__NotAuthorized();
+    error StakeManagerVaultRegistry__VaultAlreadyRegistered();
+    error StakeManagerVaultRegistry__InvalidVault();
+
+    address public stakeManager;
+
+    mapping(address owner => address[] vault) public vaults;
+    mapping(address vault => address owner) public vaultOwners;
+
+    modifier onlyStakeManager() {
+        if (msg.sender != stakeManager) {
+            revert StakeManagerVaultRegistry__NotAuthorized();
+        }
+        _;
+    }
+
+    constructor(address _owner) Ownable(_owner) { }
+
+    function setStakeManager(address _stakeManager) external onlyOwner {
+        stakeManager = _stakeManager;
+    }
+
+    function register(address owner, address vault) external onlyStakeManager {
+        if (vaultOwners[vault] != address(0)) {
+            revert StakeManagerVaultRegistry__VaultAlreadyRegistered();
+        }
+
+        // Verify this is a legitimate vault by checking it points to stakeManager
+        if (address(IStakeVault(vault).stakeManager()) != stakeManager) {
+            revert StakeManagerVaultRegistry__InvalidVault();
+        }
+
+        vaultOwners[vault] = owner;
+        vaults[owner].push(vault);
+    }
+
+    function vaultsOf(address owner) external view returns (address[] memory) {
+        return vaults[owner];
+    }
+}

src/interfaces/IStakeVaultRegistry.sol

@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.26;
+
+interface IStakeVaultRegistry {
+    function register(address owner, address vault) external;
+    function vaultsOf(address owner) external view returns (address[] memory);
+    function vaultOwners(address vault) external view returns (address);
+}

test/RewardsStreamerMP.t.sol

@@ -7,6 +7,7 @@ import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils
 import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
 import { RewardsStreamerMP } from "../src/RewardsStreamerMP.sol";
 import { StakeVault } from "../src/StakeVault.sol";
+import { StakeVaultRegistry } from "../src/StakeVaultRegistry.sol";
 import { IStakeManagerProxy } from "../src/interfaces/IStakeManagerProxy.sol";
 import { StakeManagerProxy } from "../src/StakeManagerProxy.sol";
 import { MockToken } from "./mocks/MockToken.sol";
@@ -15,6 +16,7 @@ import { StackOverflowStakeManager } from "./mocks/StackOverflowStakeManager.sol
 contract RewardsStreamerMPTest is Test {
     MockToken stakingToken;
     RewardsStreamerMP public streamer;
+    StakeVaultRegistry public vaultRegistry;
 
     address admin = makeAddr("admin");
     address alice = makeAddr("alice");
@@ -26,12 +28,16 @@ contract RewardsStreamerMPTest is Test {
 
     function setUp() public virtual {
         stakingToken = new MockToken("Staking Token", "ST");
+        vaultRegistry = new StakeVaultRegistry(address(this));
 
-        bytes memory initializeData = abi.encodeCall(RewardsStreamerMP.initialize, (admin, address(stakingToken)));
+        bytes memory initializeData =
+            abi.encodeCall(RewardsStreamerMP.initialize, (admin, address(vaultRegistry), address(stakingToken)));
         address impl = address(new RewardsStreamerMP());
         address proxy = address(new StakeManagerProxy(impl, initializeData));
         streamer = RewardsStreamerMP(proxy);
 
+        vaultRegistry.setStakeManager(address(streamer));
+
         // Create a temporary vault just to get the codehash
         StakeVault tempVault = new StakeVault(address(this), IStakeManagerProxy(address(streamer)));
         bytes32 vaultCodeHash = address(tempVault).codehash;
@@ -167,7 +173,7 @@ contract VaultRegistrationTest is RewardsStreamerMPTest {
     function test_VaultRegistration() 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]);
+            address[] memory userVaults = vaultRegistry.vaultsOf(accounts[i]);
             assertEq(userVaults.length, 1, "wrong number of vaults");
             assertEq(userVaults[0], vaults[accounts[i]], "wrong vault address");
         }