Streamed rewards

.gas-snapshot

@@ -1,61 +1,67 @@
-EmergencyExitTest:test_CannotEnableEmergencyModeTwice() (gas: 89712)
-EmergencyExitTest:test_CannotLeaveBeforeEmergencyMode() (gas: 296084)
-EmergencyExitTest:test_EmergencyExitBasic() (gas: 398741)
-EmergencyExitTest:test_EmergencyExitMultipleUsers() (gas: 830337)
-EmergencyExitTest:test_EmergencyExitToAlternateAddress() (gas: 541602)
-EmergencyExitTest:test_EmergencyExitWithLock() (gas: 387785)
-EmergencyExitTest:test_EmergencyExitWithRewards() (gas: 533248)
-EmergencyExitTest:test_OnlyOwnerCanEnableEmergencyMode() (gas: 39377)
-IntegrationTest:testStakeFoo() (gas: 1578393)
-LeaveTest:test_LeaveShouldProperlyUpdateAccounting() (gas: 2592444)
-LeaveTest:test_RevertWhenStakeManagerIsTrusted() (gas: 293227)
-LeaveTest:test_TrustNewStakeManager() (gas: 2642695)
-LockTest:test_LockFailsWithInvalidPeriod() (gas: 306175)
-LockTest:test_LockFailsWithNoStake() (gas: 61440)
-LockTest:test_LockWithoutPriorLock() (gas: 403273)
-MaliciousUpgradeTest:test_UpgradeStackOverflowStakeManager() (gas: 1761895)
+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)
 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)
 RewardsStreamerTest:testStake() (gas: 869181)
-StakeTest:test_StakeMultipleAccounts() (gas: 512711)
-StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 668852)
-StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 879853)
-StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 527355)
-StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 549222)
-StakeTest:test_StakeOneAccount() (gas: 293776)
-StakeTest:test_StakeOneAccountAndRewards() (gas: 449912)
-StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 540462)
-StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 536907)
-StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 314452)
-StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 314419)
-StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 314530)
+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)
 StakingTokenTest:testStakeToken() (gas: 10422)
-UnstakeTest:test_StakeMultipleAccounts() (gas: 512733)
-UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 668874)
-UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 879830)
-UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 527332)
-UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 549244)
-UnstakeTest:test_StakeOneAccount() (gas: 293799)
-UnstakeTest:test_StakeOneAccountAndRewards() (gas: 449934)
-UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 540484)
-UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 536887)
-UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 314452)
-UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 314419)
-UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 314508)
-UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 549390)
-UnstakeTest:test_UnstakeMultipleAccounts() (gas: 728319)
-UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 1073610)
-UnstakeTest:test_UnstakeOneAccount() (gas: 508967)
-UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 531442)
-UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 615950)
-UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 564034)
-UpgradeTest:test_RevertWhenNotOwner() (gas: 2210367)
-UpgradeTest:test_UpgradeStakeManager() (gas: 2499517)
-WithdrawTest:test_CannotWithdrawStakedFunds() (gas: 308844)
+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)
 XPNFTTokenTest:testApproveNotAllowed() (gas: 10500)
 XPNFTTokenTest:testGetApproved() (gas: 10523)
 XPNFTTokenTest:testIsApprovedForAll() (gas: 10698)

certora/confs/EmergencyMode.conf

@@ -4,8 +4,7 @@
     "certora/helpers/ERC20A.sol"
   ],
   "link" : [
-    "RewardsStreamerMP:STAKING_TOKEN=ERC20A",
-    "RewardsStreamerMP:REWARD_TOKEN=ERC20A"
+    "RewardsStreamerMP:STAKING_TOKEN=ERC20A"
   ],
   "msg": "Verifying RewardsStreamerMP.sol",
   "rule_sanity": "basic",

certora/confs/RewardsStreamerMP.conf

@@ -4,8 +4,7 @@
     "certora/helpers/ERC20A.sol"
   ],
   "link" : [
-    "RewardsStreamerMP:STAKING_TOKEN=ERC20A",
-    "RewardsStreamerMP:REWARD_TOKEN=ERC20A"
+    "RewardsStreamerMP:STAKING_TOKEN=ERC20A"
   ],
   "msg": "Verifying RewardsStreamerMP.sol",
   "rule_sanity": "basic",

certora/confs/StakeVault.conf

@@ -7,7 +7,6 @@
   "link" : [
     "StakeVault:STAKING_TOKEN=ERC20A",
     "RewardsStreamerMP:STAKING_TOKEN=ERC20A",
-    "RewardsStreamerMP:REWARD_TOKEN=ERC20A",
     "StakeVault:stakeManager=RewardsStreamerMP"
   ],
   "msg": "Verifying StakeVault.sol",

certora/specs/EmergencyMode.spec

@@ -7,7 +7,6 @@ methods {
 
 definition isViewFunction(method f) returns bool = (
   f.selector == sig:streamer.STAKING_TOKEN().selector ||
-  f.selector == sig:streamer.REWARD_TOKEN().selector ||
   f.selector == sig:streamer.SCALE_FACTOR().selector ||
   f.selector == sig:streamer.MP_RATE_PER_YEAR().selector ||
   f.selector == sig:streamer.MIN_LOCKUP_PERIOD().selector ||

src/RewardsStreamerMP.sol

@@ -3,7 +3,7 @@
 
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import { ReentrancyGuardUpgradeable } from "@openzeppelin/contracts-upgradeable/utils/ReentrancyGuardUpgradeable.sol";
 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 { TrustedCodehashAccess } from "./TrustedCodehashAccess.sol";
@@ -18,9 +18,9 @@
     error StakingManager__TokensAreLocked();
     error StakingManager__AlreadyLocked();
     error StakingManager__EmergencyModeEnabled();
+    error StakingManager__DurationCannotBeZero();
 
     IERC20 public STAKING_TOKEN;
-    IERC20 public REWARD_TOKEN;
 
     uint256 public constant SCALE_FACTOR = 1e18;
     uint256 public constant MP_RATE_PER_YEAR = 1e18;
@@ -33,10 +33,15 @@
     uint256 public totalMP;
     uint256 public totalMaxMP;
     uint256 public rewardIndex;
-    uint256 public accountedRewards;
     uint256 public lastMPUpdatedTime;
     bool public emergencyModeEnabled;
 
+    uint256 public totalRewardsAccrued;
+    uint256 public rewardAmount;
+    uint256 public lastRewardTime;
+    uint256 public rewardStartTime;
+    uint256 public rewardEndTime;
+
     struct Account {
         uint256 stakedBalance;
         uint256 accountRewardIndex;
@@ -59,13 +64,12 @@
         _disableInitializers();
     }
 
-    function initialize(address _owner, address _stakingToken, address _rewardToken) public initializer {
+    function initialize(address _owner, address _stakingToken) public initializer {
         __TrustedCodehashAccess_init(_owner);
         __UUPSUpgradeable_init();
         __ReentrancyGuard_init();
 
         STAKING_TOKEN = IERC20(_stakingToken);
-        REWARD_TOKEN = IERC20(_rewardToken);
         lastMPUpdatedTime = block.timestamp;
     }
 
@@ -90,11 +94,6 @@
             revert StakingManager__CannotRestakeWithLockedFunds();
         }
 
-        uint256 accountRewards = calculateAccountRewards(msg.sender);
-        if (accountRewards > 0) {
-            distributeRewards(msg.sender, accountRewards);
-        }
-
         account.stakedBalance += amount;
         totalStaked += amount;
 
@@ -171,11 +170,6 @@
         _updateGlobalState();
         _updateAccountMP(accountAddress);
 
-        uint256 accountRewards = calculateAccountRewards(accountAddress);
-        if (accountRewards > 0) {
-            distributeRewards(accountAddress, accountRewards);
-        }
-
         uint256 previousStakedBalance = account.stakedBalance;
 
         uint256 mpToReduce = (account.accountMP * amount * SCALE_FACTOR) / (previousStakedBalance * SCALE_FACTOR);
@@ -237,6 +231,7 @@
         // Adjust rewardIndex before updating totalMP
         uint256 previousTotalWeight = totalStaked + totalMP;
         totalMP += accruedMP;
+
         uint256 newTotalWeight = totalStaked + totalMP;
 
         if (previousTotalWeight != 0 && newTotalWeight != previousTotalWeight) {
@@ -246,19 +241,73 @@
         lastMPUpdatedTime = currentTime;
     }
 
+    function setReward(uint256 amount, uint256 duration) external onlyOwner {
+        if (duration == 0) {
+            revert StakingManager__DurationCannotBeZero();
+        }
+
+        if (amount == 0) {
+            revert StakingManager__AmountCannotBeZero();
+        }
+
+        // this will call _updateRewardIndex and update the totalRewardsAccrued
+        _updateGlobalState();
+
+        // in case _updateRewardIndex returns earlier,
+        // we still update the lastRewardTime
+        lastRewardTime = block.timestamp;
+        rewardAmount = amount;
+        rewardStartTime = block.timestamp;
+        rewardEndTime = block.timestamp + duration;
+    }
+
+    function _calculateAccruedRewards() internal view returns (uint256) {
+        if (rewardEndTime <= rewardStartTime) {
+            // No active reward period
+            return 0;
+        }
+
+        uint256 currentTime = block.timestamp < rewardEndTime ? block.timestamp : rewardEndTime;
+
+        if (currentTime <= lastRewardTime) {
+            // No new rewards have accrued since lastRewardTime
+            return 0;
+        }
+
+        uint256 timeElapsed = currentTime - lastRewardTime;
+        uint256 duration = rewardEndTime - rewardStartTime;
+
+        if (duration == 0) {
+            // Prevent division by zero
+            return 0;
+        }
+
+        uint256 accruedRewards = (timeElapsed * rewardAmount) / duration;
+        return accruedRewards;
+    }
+
     function updateRewardIndex() internal {
         uint256 totalWeight = totalStaked + totalMP;
         if (totalWeight == 0) {
             return;
         }
 
-        uint256 rewardBalance = REWARD_TOKEN.balanceOf(address(this));
-        uint256 newRewards = rewardBalance > accountedRewards ? rewardBalance - accountedRewards : 0;
+        uint256 currentTime = block.timestamp;
+        uint256 applicableTime = rewardEndTime > currentTime ? currentTime : rewardEndTime;
+        uint256 elapsedTime = applicableTime - lastRewardTime;
+
+        if (elapsedTime == 0) {
+            return;
+        }
 
-        if (newRewards > 0) {
-            rewardIndex += (newRewards * SCALE_FACTOR) / totalWeight;
-            accountedRewards += newRewards;
+        uint256 newRewards = _calculateAccruedRewards();
+        if (newRewards == 0) {
+            return;
         }
+
+        totalRewardsAccrued += newRewards;
+        rewardIndex += (newRewards * SCALE_FACTOR) / totalWeight;
+        lastRewardTime = block.timestamp < rewardEndTime ? block.timestamp : rewardEndTime;
     }
 
     function _calculateBonusMP(uint256 amount, uint256 lockPeriod) internal view returns (uint256) {
@@ -295,24 +344,11 @@
 
     function calculateAccountRewards(address accountAddress) public view returns (uint256) {
         Account storage account = accounts[accountAddress];
+
         uint256 accountWeight = account.stakedBalance + account.accountMP;
         uint256 deltaRewardIndex = rewardIndex - account.accountRewardIndex;
-        return (accountWeight * deltaRewardIndex) / SCALE_FACTOR;
-    }
-
-    function distributeRewards(address to, uint256 amount) internal {
-        uint256 rewardBalance = REWARD_TOKEN.balanceOf(address(this));
-        // If amount is higher than the contract's balance (for rounding error), transfer the balance.
-        if (amount > rewardBalance) {
-            amount = rewardBalance;
-        }
-
-        accountedRewards -= amount;
 
-        bool success = REWARD_TOKEN.transfer(to, amount);
-        if (!success) {
-            revert StakingManager__TransferFailed();
-        }
+        return (accountWeight * deltaRewardIndex) / SCALE_FACTOR;
     }
 
     function enableEmergencyMode() external onlyOwner {
@@ -326,11 +362,15 @@
         return accounts[accountAddress].stakedBalance;
     }
 
-    function getPendingRewards(address accountAddress) external view returns (uint256) {
-        return calculateAccountRewards(accountAddress);
-    }
-
     function getAccount(address accountAddress) external view returns (Account memory) {
         return accounts[accountAddress];
     }
+
+    function totalRewardsSupply() public view returns (uint256) {
+        return totalRewardsAccrued + _calculateAccruedRewards();
+    }
+
+    function rewardsBalanceOf(address accountAddress) external view returns (uint256) {
+        return calculateAccountRewards(accountAddress);
+    }
 }

src/interfaces/IStakeManager.sol

@@ -23,7 +23,6 @@ interface IStakeManager is ITrustedCodehashAccess {
     function getStakedBalance(address _vault) external view returns (uint256 _balance);
 
     function STAKING_TOKEN() external view returns (IERC20);
-    function REWARD_TOKEN() external view returns (IERC20);
     function MIN_LOCKUP_PERIOD() external view returns (uint256);
     function MAX_LOCKUP_PERIOD() external view returns (uint256);
     function MP_RATE_PER_YEAR() external view returns (uint256);