Fix: Proxy version storage collision and KPI reward window snipe

src/modules/logicModule/LM_PC_KPIRewarder_v2.sol

@@ -70,8 +70,8 @@ contract LM_PC_KPIRewarder_v2 is
     //   or not (see below).
     // - An external actor with the ASSERTER role can trigger the posting of an assertion to the UMA Oracle, specifying
     //   the value to be asserted and the KPI to use for the reward distrbution in case it resolves
-    // - To ensure fairness, all new staking requests are queued until the next KPI assertion is resolved.
-    //   They will be added before posting the next assertion.
+    // - To ensure fairness, staking is blocked while an assertion is pending and until the reward window
+    //   opened by its resolution has closed, so a resolved round only pays the stakers that were staked during it.
     // - Once the assertion resolves, the UMA oracle triggers the assertionResolvedCallback() function.
     //   This will calculate the final reward value and distribute it to the stakers.
 
@@ -86,9 +86,11 @@ contract LM_PC_KPIRewarder_v2 is
     mapping(bytes32 => RewardRoundConfiguration) internal assertionConfig;
     /// @dev	 For locking certain utilities when there are assertions open.
     bool internal assertionPending;
+    /// @dev	End of the reward window opened by the last resolved assertion. Staking is blocked until then.
+    uint internal rewardWindowEnd;
 
     /// @dev	Storage gap for future upgrades.
-    uint[50] private __gap;
+    uint[49] private __gap;
 
     /*
     Tranche Example:
@@ -298,6 +300,12 @@ contract LM_PC_KPIRewarder_v2 is
             );
         }
 
+        // Block staking in the reward window of a resolved assertion, otherwise the assertion could be settled and
+        // staked into within the same transaction to capture a share of a round the staker was never part of.
+        if (block.timestamp < rewardWindowEnd) {
+            revert Module__LM_PC_KPIRewarder_v2__CannotStakeDuringRewardWindow();
+        }
+
         address sender = _msgSender();
 
         _stake(sender, amount);
@@ -395,6 +403,7 @@ contract LM_PC_KPIRewarder_v2 is
             }
 
             _setRewards(rewardAmount, 1);
+            rewardWindowEnd = block.timestamp + 1;
             assertionConfig[assertionId].distributed = true;
         } else {
             // To keep in line with the upstream contract. If the assertion was false, we delete the corresponding

src/modules/logicModule/interfaces/ILM_PC_KPIRewarder_v2.sol

@@ -55,6 +55,9 @@ interface ILM_PC_KPIRewarder_v2 {
     /// @notice The user cannot stake while an assertion is unresolved.
     error Module__LM_PC_KPIRewarder_v2__CannotStakeWhenAssertionPending();
 
+    /// @notice The user cannot stake while the reward window of a resolved assertion is open.
+    error Module__LM_PC_KPIRewarder_v2__CannotStakeDuringRewardWindow();
+
     /// @notice Callback received references non existent assertionId.
     error Module__LM_PC_KPIRewarder_v2__NonExistentAssertionId(
         bytes32 assertionId

src/proxies/InverterTransparentUpgradeableProxy_v1.sol

@@ -66,14 +66,19 @@ contract InverterTransparentUpgradeableProxy_v1 is ERC1967Proxy {
     /// @dev	The address of the beacon that is used to fetch the implementation address.
     IInverterBeacon_v1 internal immutable _beacon;
 
-    /// @dev	The major version of the implementation.
-    uint internal majorVersion;
-
-    /// @dev	The minor version of the implementation.
-    uint internal minorVersion;
+    /// @dev	The version of the linked implementation, stored in an ERC-7201 namespace so it cannot
+    ///         collide with the storage of the implementation the proxy delegates to.
+    /// @custom:storage-location erc7201:inverter.storage.InverterTransparentUpgradeableProxy_v1
+    struct VersionStorage {
+        uint majorVersion;
+        uint minorVersion;
+        uint patchVersion;
+    }
 
-    /// @dev	The patch version of the implementation.
-    uint internal patchVersion;
+    /// @dev	keccak256(abi.encode(uint(keccak256("inverter.storage.InverterTransparentUpgradeableProxy_v1")) - 1))
+    ///         & ~bytes32(uint(0xff))
+    bytes32 private constant VERSION_STORAGE_LOCATION =
+        0x7f64a4379a384334ad31a069fd0e3183465498d81cb6b14f5f286747a479a800;
 
     //--------------------------------------------------------------------------
     // Constructor
@@ -93,7 +98,7 @@ contract InverterTransparentUpgradeableProxy_v1 is ERC1967Proxy {
 
         _beacon = beacon;
         _admin = initialOwner;
-        (majorVersion, minorVersion, patchVersion) = _beacon.version();
+        _setVersion();
         // Set the storage value and emit an event for ERC-1967 compatibility
         ERC1967Utils.changeAdmin(initialOwner);
     }
@@ -108,7 +113,8 @@ contract InverterTransparentUpgradeableProxy_v1 is ERC1967Proxy {
     /// @return The minor version.
     /// @return The patch version.
     function version() external view returns (uint, uint, uint) {
-        return (majorVersion, minorVersion, patchVersion);
+        VersionStorage storage $ = _getVersionStorage();
+        return ($.majorVersion, $.minorVersion, $.patchVersion);
     }
 
     /// @dev    Fallback function to delegate calls to the implementation contract
@@ -125,7 +131,24 @@ contract InverterTransparentUpgradeableProxy_v1 is ERC1967Proxy {
         // Override implementation
         ERC1967Utils.upgradeToAndCall(_beacon.implementation(), bytes(""));
         // Override version
-        (majorVersion, minorVersion, patchVersion) = _beacon.version();
+        _setVersion();
+    }
+
+    /// @dev	Stores the current version of the linked beacon.
+    function _setVersion() internal {
+        VersionStorage storage $ = _getVersionStorage();
+        ($.majorVersion, $.minorVersion, $.patchVersion) = _beacon.version();
+    }
+
+    /// @dev	Returns the namespaced storage struct holding the version.
+    function _getVersionStorage()
+        private
+        pure
+        returns (VersionStorage storage $)
+    {
+        assembly {
+            $.slot := VERSION_STORAGE_LOCATION
+        }
     }
 
     /// @dev	If caller is the admin process the call internally, otherwise transparently fallback to the proxy behavior.

test/unit/modules/logicModule/LM_PC_KPIRewarder_v2.t.sol

@@ -769,6 +769,8 @@ stakeTest
 │   └── it should revert
 ├── when there is an unresolved assertion live
 │   └── it should revert
+├── when the reward window of a resolved assertion is open
+│   └── it should revert
 ├── when the caller does not have sufficient funds
 │   └── it should revert
 └── when the caller has sufficient funds
@@ -857,6 +859,49 @@ contract LM_PC_KPIRewarder_v2_stakeTest is LM_PC_KPIRewarder_v2Test {
         vm.stopPrank();
     }
 
+    function test_RevertWhen_TheRewardWindowIsOpen() external {
+        // it should revert
+
+        address[] memory users = new address[](2);
+        users[0] = makeAddr("staker1");
+        users[1] = makeAddr("staker2");
+        uint[] memory amounts = new uint[](2);
+        amounts[0] = 1000e18;
+        amounts[1] = 1000e18;
+
+        (bytes32 createdID,,,) =
+            setUpStateForAssertionResolution(users, amounts, 250e18, true);
+
+        // Settling the assertion and staking in the same transaction must not be possible
+        vm.prank(address(ooV3));
+        kpiManager.assertionResolvedCallback(createdID, true);
+
+        address sniper = makeAddr("sniper");
+        uint sniperAmount = 18_000e18;
+        stakingToken.mint(sniper, sniperAmount);
+        vm.startPrank(sniper);
+        stakingToken.approve(address(kpiManager), sniperAmount);
+        vm.expectRevert(
+            ILM_PC_KPIRewarder_v2
+                .Module__LM_PC_KPIRewarder_v2__CannotStakeDuringRewardWindow
+                .selector
+        );
+        kpiManager.stake(sniperAmount);
+
+        // Once the window has closed, staking works again but earns nothing from the resolved round
+        vm.warp(block.timestamp + 1);
+        kpiManager.stake(sniperAmount);
+        vm.stopPrank();
+
+        vm.warp(block.timestamp + 10);
+        assertEq(kpiManager.getEarned(sniper), 0);
+        assertApproxEqAbs(
+            kpiManager.getEarned(users[0]) + kpiManager.getEarned(users[1]),
+            kpiManager.getRewardRate(),
+            1e8
+        );
+    }
+
     function test_WhenTheCallerHasSufficientFunds(uint amount) external {
         // it should stake the funds normally
 

test/unit/proxies/InverterTransparentUpgradeableProxy_v1.t.sol

@@ -157,4 +157,37 @@ contract InverterTransparentUpgradeableProxyV1Test is Test {
         assertEq(returnedMinorVersion, initialMinorVersion + 1);
         assertEq(returnedPatchVersion, initialPatchVersion + 1);
     }
+
+    function test_upgradeToNewestVersionKeepsImplementationStorage() public {
+        // Simulate implementation state in the first storage slots, e.g. the orchestrator of a module
+        for (uint i; i < 3; i++) {
+            vm.store(address(proxy), bytes32(i), bytes32(uint(0xdead + i)));
+        }
+
+        beacon.overrideImplementation(address(implementation2));
+        beacon.overrideVersion(
+            initialMajorVersion + 1,
+            initialMinorVersion + 1,
+            initialPatchVersion + 1
+        );
+
+        vm.prank(admin);
+        IInverterTransparentUpgradeableProxy_v1(address(proxy))
+            .upgradeToNewestVersion();
+
+        for (uint i; i < 3; i++) {
+            assertEq(
+                vm.load(address(proxy), bytes32(i)), bytes32(uint(0xdead + i))
+            );
+        }
+
+        (
+            uint returnedMajorVersion,
+            uint returnedMinorVersion,
+            uint returnedPatchVersion
+        ) = proxy.version();
+        assertEq(returnedMajorVersion, initialMajorVersion + 1);
+        assertEq(returnedMinorVersion, initialMinorVersion + 1);
+        assertEq(returnedPatchVersion, initialPatchVersion + 1);
+    }
 }