Check available balance in previewRewards (#413)

* Check available balance in `previewRewards` * Chore: forge fmt --------- Co-authored-by: Daniel Von Fange <[email protected]>
OriginProtocol · May 8, 2024 · 2a51241 · 2a51241
1 parent 51035a9
commit 2a51241
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/contracts/FixedRateRewardsSource.sol b/contracts/FixedRateRewardsSource.sol
@@ -87,12 +87,18 @@ contract FixedRateRewardsSource is Governable, Initializable {
 
     /// @dev Compute pending rewards since last collect
     /// @return rewardAmount Amount of reward that'll be distributed if collected now
-    function previewRewards() public view returns (uint256) {
+    function previewRewards() public view returns (uint256 rewardAmount) {
         RewardConfig memory _config = rewardConfig;
+
         if (_config.lastCollect == 0) {
             return 0;
         }
-        return (block.timestamp - _config.lastCollect) * _config.rewardsPerSecond;
+
+        rewardAmount = (block.timestamp - _config.lastCollect) * _config.rewardsPerSecond;
+        uint256 balance = IERC20(rewardToken).balanceOf(address(this));
+        if (rewardAmount > balance) {
+            rewardAmount = balance;
+        }
     }
 
     /// @dev Set address of the strategist

diff --git a/tests/staking/FixedRateRewardsSource.t.sol b/tests/staking/FixedRateRewardsSource.t.sol
@@ -100,6 +100,18 @@ contract FixedRateRewardsSourceTest is Test {
         assertEq(rewards.previewRewards(), 0 ether, "Pending reward mismatch");
     }
 
+    function testLowBalanceCollection() public {
+        // Should also allow disabling rewards
+        vm.prank(strategist);
+        rewards.setRewardsPerSecond(2000000 ether);
+
+        // Should never show more than balance
+        vm.warp(block.number + 10);
+        assertEq(rewards.previewRewards(), 1000000 ether, "Pending reward mismatch");
+        vm.warp(block.number + 123);
+        assertEq(rewards.previewRewards(), 1000000 ether, "Pending reward mismatch");
+    }
+
     function testRewardRatePermission() public {
         // Should allow Strategist to change
         vm.prank(strategist);