WIP: upgradeability and trust assumptions

.gas-snapshot

@@ -1,84 +1,83 @@
-EmergencyExitTest:test_CannotEnableEmergencyModeTwice() (gas: 79785)
-EmergencyExitTest:test_CannotLeaveBeforeEmergencyMode() (gas: 283229)
-EmergencyExitTest:test_EmergencyExitBasic() (gas: 379825)
-EmergencyExitTest:test_EmergencyExitMultipleUsers() (gas: 788638)
-EmergencyExitTest:test_EmergencyExitToAlternateAddress() (gas: 517427)
-EmergencyExitTest:test_EmergencyExitWithLock() (gas: 370277)
-EmergencyExitTest:test_EmergencyExitWithRewards() (gas: 507312)
-EmergencyExitTest:test_OnlyOwnerCanEnableEmergencyMode() (gas: 34585)
-IntegrationTest:testStakeFoo() (gas: 1490184)
-LeaveTest:test_LeaveShouldProperlyUpdateAccounting() (gas: 365536)
-LeaveTest:test_RevertWhenStakeManagerIsTrusted() (gas: 275827)
-LockTest:test_LockFailsWithInvalidPeriod() (gas: 290237)
-LockTest:test_LockFailsWithNoStake() (gas: 53423)
-LockTest:test_LockWithoutPriorLock() (gas: 383200)
-NFTMetadataGeneratorSVGTest:testGenerateMetadata() (gas: 92874)
-NFTMetadataGeneratorSVGTest:testSetImageStrings() (gas: 60081)
-NFTMetadataGeneratorSVGTest:testSetImageStringsRevert() (gas: 35818)
-NFTMetadataGeneratorURLTest:testGenerateMetadata() (gas: 109345)
-NFTMetadataGeneratorURLTest:testSetBaseURL() (gas: 50653)
-NFTMetadataGeneratorURLTest:testSetBaseURLRevert() (gas: 35993)
-RewardsStreamerTest:testStake() (gas: 869874)
-StakeTest:test_StakeMultipleAccounts() (gas: 493245)
-StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 640707)
-StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 818174)
-StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 499259)
-StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 520661)
-StakeTest:test_StakeOneAccount() (gas: 284249)
-StakeTest:test_StakeOneAccountAndRewards() (gas: 431706)
-StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 498829)
-StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 493984)
-StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 298103)
-StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 298115)
-StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 298226)
-UnstakeTest:test_StakeMultipleAccounts() (gas: 493267)
-UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 640729)
-UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 818151)
-UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 499236)
-UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 520683)
-UnstakeTest:test_StakeOneAccount() (gas: 284272)
-UnstakeTest:test_StakeOneAccountAndRewards() (gas: 431728)
-UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 498851)
-UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 493964)
-UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 298103)
-UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 298115)
-UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 298204)
-UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 508466)
-UnstakeTest:test_UnstakeMultipleAccounts() (gas: 688777)
-UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 1014246)
-UnstakeTest:test_UnstakeOneAccount() (gas: 480136)
-UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 496593)
-UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 585942)
-UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 518419)
-XPNFTTokenTest:testApproveNotAllowed() (gas: 10507)
-XPNFTTokenTest:testGetApproved() (gas: 10531)
-XPNFTTokenTest:testIsApprovedForAll() (gas: 10705)
-XPNFTTokenTest:testSafeTransferNotAllowed() (gas: 10688)
-XPNFTTokenTest:testSafeTransferWithDataNotAllowed() (gas: 10906)
-XPNFTTokenTest:testSetApprovalForAllNotAllowed() (gas: 8474)
-XPNFTTokenTest:testSetMetadataGenerator() (gas: 1014075)
-XPNFTTokenTest:testSetMetadataGeneratorRevert() (gas: 1010606)
-XPNFTTokenTest:testTokenURI() (gas: 111913)
-XPNFTTokenTest:testTransferNotAllowed() (gas: 10723)
-XPTokenMintAllowanceTest:testAddXPProviderOnlyOwner() (gas: 285721)
-XPTokenMintAllowanceTest:testBalanceOf() (gas: 294561)
-XPTokenMintAllowanceTest:testBalanceOfWithNoSystemTotalXP() (gas: 43428)
-XPTokenMintAllowanceTest:testMintAllowance_Available() (gas: 205169)
-XPTokenMintAllowanceTest:testMintAllowance_NotAvailable() (gas: 205107)
-XPTokenMintAllowanceTest:testMintOnlyOwner() (gas: 241956)
-XPTokenMintAllowanceTest:testMint_Ok() (gas: 264233)
-XPTokenMintAllowanceTest:testMint_RevertWithAllowanceExceeded() (gas: 246656)
-XPTokenMintAllowanceTest:testRemoveXPProviderIndexOutOfBounds() (gas: 36295)
-XPTokenMintAllowanceTest:testRemoveXPProviderOnlyOwner() (gas: 72176)
-XPTokenMintAllowanceTest:testTotalSupply() (gas: 202454)
-XPTokenMintAllowanceTest:testTransfersNotAllowed() (gas: 20658)
-XPTokenOwnershipTest:testInitialOwner() (gas: 12649)
-XPTokenOwnershipTest:testOwnershipTransfer() (gas: 87271)
-XPTokenTest:testAddXPProviderOnlyOwner() (gas: 285753)
-XPTokenTest:testBalanceOf() (gas: 294565)
-XPTokenTest:testBalanceOfWithNoSystemTotalXP() (gas: 43405)
-XPTokenTest:testMintOnlyOwner() (gas: 241932)
-XPTokenTest:testRemoveXPProviderIndexOutOfBounds() (gas: 36277)
-XPTokenTest:testRemoveXPProviderOnlyOwner() (gas: 72141)
-XPTokenTest:testTotalSupply() (gas: 202403)
-XPTokenTest:testTransfersNotAllowed() (gas: 20702)
+EmergencyExitTest:test_CannotEnableEmergencyModeTwice() (gas: 89712)
+EmergencyExitTest:test_CannotLeaveBeforeEmergencyMode() (gas: 296040)
+EmergencyExitTest:test_EmergencyExitBasic() (gas: 398697)
+EmergencyExitTest:test_EmergencyExitMultipleUsers() (gas: 830249)
+EmergencyExitTest:test_EmergencyExitToAlternateAddress() (gas: 541558)
+EmergencyExitTest:test_EmergencyExitWithLock() (gas: 387741)
+EmergencyExitTest:test_EmergencyExitWithRewards() (gas: 533204)
+EmergencyExitTest:test_OnlyOwnerCanEnableEmergencyMode() (gas: 39377)
+IntegrationTest:testStakeFoo() (gas: 1578327)
+LeaveTest:test_LeaveShouldProperlyUpdateAccounting() (gas: 2592493)
+LockTest:test_LockFailsWithInvalidPeriod() (gas: 306108)
+LockTest:test_LockFailsWithNoStake() (gas: 61395)
+LockTest:test_LockWithoutPriorLock() (gas: 403206)
+NFTMetadataGeneratorSVGTest:testGenerateMetadata() (gas: 85934)
+NFTMetadataGeneratorSVGTest:testSetImageStrings() (gas: 58332)
+NFTMetadataGeneratorSVGTest:testSetImageStringsRevert() (gas: 35804)
+NFTMetadataGeneratorURLTest:testGenerateMetadata() (gas: 102512)
+NFTMetadataGeneratorURLTest:testSetBaseURL() (gas: 49555)
+NFTMetadataGeneratorURLTest:testSetBaseURLRevert() (gas: 35979)
+RewardsStreamerTest:testStake() (gas: 869181)
+StakeTest:test_StakeMultipleAccounts() (gas: 512667)
+StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 668808)
+StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 879809)
+StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 527311)
+StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 549178)
+StakeTest:test_StakeOneAccount() (gas: 293754)
+StakeTest:test_StakeOneAccountAndRewards() (gas: 449890)
+StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 540440)
+StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 536885)
+StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 314430)
+StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 314397)
+StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 314508)
+UnstakeTest:test_StakeMultipleAccounts() (gas: 512689)
+UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 668830)
+UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 879786)
+UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 527288)
+UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 549200)
+UnstakeTest:test_StakeOneAccount() (gas: 293777)
+UnstakeTest:test_StakeOneAccountAndRewards() (gas: 449912)
+UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 540462)
+UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 536865)
+UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 314430)
+UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 314397)
+UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 314486)
+UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 549368)
+UnstakeTest:test_UnstakeMultipleAccounts() (gas: 728275)
+UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 1073566)
+UnstakeTest:test_UnstakeOneAccount() (gas: 508945)
+UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 531420)
+UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 615928)
+UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 564012)
+XPNFTTokenTest:testApproveNotAllowed() (gas: 10500)
+XPNFTTokenTest:testGetApproved() (gas: 10523)
+XPNFTTokenTest:testIsApprovedForAll() (gas: 10698)
+XPNFTTokenTest:testSafeTransferNotAllowed() (gas: 10680)
+XPNFTTokenTest:testSafeTransferWithDataNotAllowed() (gas: 10897)
+XPNFTTokenTest:testSetApprovalForAllNotAllowed() (gas: 8467)
+XPNFTTokenTest:testSetMetadataGenerator() (gas: 969770)
+XPNFTTokenTest:testSetMetadataGeneratorRevert() (gas: 966301)
+XPNFTTokenTest:testTokenURI() (gas: 103894)
+XPNFTTokenTest:testTransferNotAllowed() (gas: 10715)
+XPTokenMintAllowanceTest:testAddXPProviderOnlyOwner() (gas: 282868)
+XPTokenMintAllowanceTest:testBalanceOf() (gas: 294494)
+XPTokenMintAllowanceTest:testBalanceOfWithNoSystemTotalXP() (gas: 43385)
+XPTokenMintAllowanceTest:testMintAllowance_Available() (gas: 205108)
+XPTokenMintAllowanceTest:testMintAllowance_NotAvailable() (gas: 205044)
+XPTokenMintAllowanceTest:testMintOnlyOwner() (gas: 241885)
+XPTokenMintAllowanceTest:testMint_Ok() (gas: 264142)
+XPTokenMintAllowanceTest:testMint_RevertWithAllowanceExceeded() (gas: 246592)
+XPTokenMintAllowanceTest:testRemoveXPProviderIndexOutOfBounds() (gas: 36286)
+XPTokenMintAllowanceTest:testRemoveXPProviderOnlyOwner() (gas: 72143)
+XPTokenMintAllowanceTest:testTotalSupply() (gas: 202403)
+XPTokenMintAllowanceTest:testTransfersNotAllowed() (gas: 20631)
+XPTokenOwnershipTest:testInitialOwner() (gas: 12645)
+XPTokenOwnershipTest:testOwnershipTransfer() (gas: 87252)
+XPTokenTest:testAddXPProviderOnlyOwner() (gas: 282900)
+XPTokenTest:testBalanceOf() (gas: 294498)
+XPTokenTest:testBalanceOfWithNoSystemTotalXP() (gas: 43362)
+XPTokenTest:testMintOnlyOwner() (gas: 241861)
+XPTokenTest:testRemoveXPProviderIndexOutOfBounds() (gas: 36268)
+XPTokenTest:testRemoveXPProviderOnlyOwner() (gas: 72108)
+XPTokenTest:testTotalSupply() (gas: 202352)
+XPTokenTest:testTransfersNotAllowed() (gas: 20675)

src/RewardsStreamerMP.sol

@@ -60,8 +60,9 @@ contract RewardsStreamerMP is UUPSUpgradeable, IStakeManager, TrustedCodehashAcc
     }
 
     function initialize(address _owner, address _stakingToken, address _rewardToken) public initializer {
-        __ReentrancyGuard_init();
         __TrustedCodehashAccess_init(_owner);
+        __UUPSUpgradeable_init();
+        __ReentrancyGuard_init();
 
         STAKING_TOKEN = IERC20(_stakingToken);
         REWARD_TOKEN = IERC20(_rewardToken);

src/StakeManagerProxy.sol

@@ -0,0 +1,12 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.26;
+
+import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
+
+contract StakeManagerProxy is ERC1967Proxy {
+    constructor(address _implementation, bytes memory _data) ERC1967Proxy(_implementation, _data) { }
+
+    function implementation() external view returns (address) {
+        return _implementation();
+    }
+}

src/StakeVault.sol

@@ -4,7 +4,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 { IStakeManagerProxy } from "./interfaces/IStakeManagerProxy.sol";
 
 /**
  * @title StakeVault
@@ -26,7 +26,8 @@ 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;
+    IStakeManagerProxy private stakeManager;
+    address public stakeManagerImplementationAddress;
 
     /**
      * @dev Emitted when tokens are staked.
@@ -56,9 +57,10 @@ contract StakeVault is Ownable {
      * @param _owner The address of the owner.
      * @param _stakeManager The address of the StakeManager contract.
      */
-    constructor(address _owner, IStakeManager _stakeManager) Ownable(_owner) {
+    constructor(address _owner, IStakeManagerProxy _stakeManager) Ownable(_owner) {
         STAKING_TOKEN = _stakeManager.STAKING_TOKEN();
         stakeManager = _stakeManager;
+        stakeManagerImplementationAddress = _stakeManager.implementation();
     }
 
     /**
@@ -226,7 +228,7 @@ contract StakeVault is Ownable {
         }
     }
 
-    function _stakeManagerImplementationTrusted() internal pure virtual returns (bool) {
-        return true;
+    function _stakeManagerImplementationTrusted() internal view virtual returns (bool) {
+        return stakeManagerImplementationAddress == stakeManager.implementation();
     }
 }

src/interfaces/IStakeManagerProxy.sol

@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.26;
+
+import { IStakeManager } from "./IStakeManager.sol";
+
+interface IStakeManagerProxy is IStakeManager {
+    function implementation() external view returns (address);
+}

test/RewardsStreamerMP.t.sol

@@ -1,18 +1,19 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.26;
 
-import { UnsafeUpgrades } from "openzeppelin-foundry-upgrades/Upgrades.sol";
 import { Test } from "forge-std/Test.sol";
 import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
+import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
 import { RewardsStreamerMP } from "../src/RewardsStreamerMP.sol";
 import { StakeVault } from "../src/StakeVault.sol";
+import { IStakeManagerProxy } from "../src/interfaces/IStakeManagerProxy.sol";
+import { StakeManagerProxy } from "../src/StakeManagerProxy.sol";
 import { StakeVaultNonTrusting } from "./harness/StakeVaultNonTrusting.sol";
 import { MockToken } from "./mocks/MockToken.sol";
 
 contract RewardsStreamerMPTest is Test {
     MockToken rewardToken;
     MockToken stakingToken;
-    RewardsStreamerMP public streamer_implementation;
     RewardsStreamerMP public streamer;
 
     address admin = makeAddr("admin");
@@ -26,15 +27,12 @@ contract RewardsStreamerMPTest is Test {
     function setUp() public virtual {
         rewardToken = new MockToken("Reward Token", "RT");
         stakingToken = new MockToken("Staking Token", "ST");
-        streamer_implementation = new RewardsStreamerMP();
-        streamer = RewardsStreamerMP(
-            UnsafeUpgrades.deployUUPSProxy(
-                address(streamer_implementation),
-                abi.encodeCall(
-                    RewardsStreamerMP.initialize, (address(this), address(stakingToken), address(rewardToken))
-                )
-            )
-        );
+
+        bytes memory initializeData =
+            abi.encodeCall(RewardsStreamerMP.initialize, (address(this), address(stakingToken), address(rewardToken)));
+        address impl = address(new RewardsStreamerMP());
+        address proxy = address(new StakeManagerProxy(impl, initializeData));
+        streamer = RewardsStreamerMP(proxy);
 
         address[4] memory accounts = [alice, bob, charlie, dave];
         for (uint256 i = 0; i < accounts.length; i++) {
@@ -97,7 +95,7 @@ contract RewardsStreamerMPTest is Test {
 
     function _createTestVault(address owner) internal returns (StakeVault vault) {
         vm.prank(owner);
-        vault = new StakeVault(owner, streamer);
+        vault = new StakeVault(owner, IStakeManagerProxy(address(streamer)));
 
         if (!streamer.isTrustedCodehash(address(vault).codehash)) {
             streamer.setTrustedCodehash(address(vault).codehash, true);
@@ -1798,33 +1796,40 @@ contract EmergencyExitTest is RewardsStreamerMPTest {
 }
 
 contract LeaveTest is RewardsStreamerMPTest {
-    StakeVaultNonTrusting public nonTrustingVault;
+    // StakeVaultNonTrusting public nonTrustingVault;
+
+    function _upgradeStakeManager() internal {
+        address newImpl = address(new RewardsStreamerMP());
+        bytes memory data;
+        UUPSUpgradeable(streamer).upgradeToAndCall(newImpl, data);
+    }
 
     function setUp() public override {
         super.setUp();
 
-        // set up vault that doesn't trust `streamer`
-        vm.startPrank(alice);
-        nonTrustingVault = new StakeVaultNonTrusting(alice, streamer);
-        stakingToken.approve(address(nonTrustingVault), 10_000e18);
-        vm.stopPrank();
-
-        // ensure non trusting vault is trusted by stake manager
-        streamer.setTrustedCodehash(address(nonTrustingVault).codehash, true);
+        // // set up vault that doesn't trust `streamer`
+        // vm.startPrank(alice);
+        // nonTrustingVault = new StakeVaultNonTrusting(alice, IStakeManagerProxy(address(streamer)));
+        // stakingToken.approve(address(nonTrustingVault), 10_000e18);
+        // vm.stopPrank();
+        //
+        // // ensure non trusting vault is trusted by stake manager
+        // streamer.setTrustedCodehash(address(nonTrustingVault).codehash, true);
     }
 
-    function test_RevertWhenStakeManagerIsTrusted() public {
-        // this test uses the standard `StakeVault` which trusts the `streamer`
-        _stake(alice, 10e18, 0);
-        vm.expectRevert(StakeVault.StakeVault__NotAllowedToLeave.selector);
-        _leave(alice);
-    }
+    // function test_RevertWhenStakeManagerIsTrusted() public {
+    //     // this test uses the standard `StakeVault` which trusts the `streamer`
+    //     _stake(alice, 10e18, 0);
+    //     vm.expectRevert(StakeVault.StakeVault__NotAllowedToLeave.selector);
+    //     _leave(alice);
+    // }
 
     function test_LeaveShouldProperlyUpdateAccounting() public {
         uint256 aliceInitialBalance = stakingToken.balanceOf(alice);
 
-        vm.prank(alice);
-        nonTrustingVault.stakeNonTrusted(100e18, 0);
+        // vm.prank(alice);
+        // nonTrustingVault.stakeNonTrusted(100e18, 0);
+        _stake(alice, 100e18, 0);
 
         assertEq(stakingToken.balanceOf(alice), aliceInitialBalance - 100e18, "Alice should have staked tokens");
 
@@ -1840,8 +1845,12 @@ contract LeaveTest is RewardsStreamerMPTest {
             })
         );
 
-        vm.prank(alice);
-        nonTrustingVault.leave(alice);
+        // upgrade stake manager, causing it to no longer be trusted
+        _upgradeStakeManager();
+
+        // vm.prank(alice);
+        // nonTrustingVault.leave(alice);
+        _leave(alice);
 
         // stake manager properly updates accounting
         checkStreamer(
@@ -1859,7 +1868,8 @@ contract LeaveTest is RewardsStreamerMPTest {
         // vault should be empty as funds have been moved out
         checkAccount(
             CheckAccountParams({
-                account: address(nonTrustingVault),
+                // account: address(nonTrustingVault),
+                account: vaults[alice],
                 rewardBalance: 0,
                 stakedBalance: 0,
                 vaultBalance: 0,