From 9b985c786c93c495052413d920c00c9bc8d9185a Mon Sep 17 00:00:00 2001
From: Shahul Hameed <10547529+shahthepro@users.noreply.github.com>
Date: Thu, 2 May 2024 15:17:25 +0530
Subject: [PATCH 1/2] Check available balance in `previewRewards`

---
 contracts/FixedRateRewardsSource.sol       | 11 ++++++++---
 tests/staking/FixedRateRewardsSource.t.sol | 12 ++++++++++++
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/contracts/FixedRateRewardsSource.sol b/contracts/FixedRateRewardsSource.sol
index 2c51f4d4..1fb97db5 100644
--- a/contracts/FixedRateRewardsSource.sol
+++ b/contracts/FixedRateRewardsSource.sol
@@ -87,10 +87,15 @@ 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;
-        return (block.timestamp - _config.lastCollect) * _config.rewardsPerSecond;
-        // return _previewRewards(rewardConfig);
+
+        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
index fce7c11b..07db745d 100644
--- 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);

From 454af8d9eb49e6428e18840c85b09b33aa107330 Mon Sep 17 00:00:00 2001
From: Daniel Von Fange <daniel@leancoder.com>
Date: Wed, 8 May 2024 11:22:12 -0400
Subject: [PATCH 2/2] Chore: forge fmt

---
 contracts/FixedRateRewardsSource.sol | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/contracts/FixedRateRewardsSource.sol b/contracts/FixedRateRewardsSource.sol
index 30f52139..4ef9f431 100644
--- a/contracts/FixedRateRewardsSource.sol
+++ b/contracts/FixedRateRewardsSource.sol
@@ -89,11 +89,11 @@ contract FixedRateRewardsSource is Governable, Initializable {
     /// @return rewardAmount Amount of reward that'll be distributed if collected now
     function previewRewards() public view returns (uint256 rewardAmount) {
         RewardConfig memory _config = rewardConfig;
-        
+
         if (_config.lastCollect == 0) {
             return 0;
         }
-        
+
         rewardAmount = (block.timestamp - _config.lastCollect) * _config.rewardsPerSecond;
         uint256 balance = IERC20(rewardToken).balanceOf(address(this));
         if (rewardAmount > balance) {