From 2a512418c7027e6d0056d74f4b3f1f5fe7881c43 Mon Sep 17 00:00:00 2001
From: Shahul Hameed <10547529+shahthepro@users.noreply.github.com>
Date: Wed, 8 May 2024 20:57:34 +0530
Subject: [PATCH] Check available balance in `previewRewards` (#413)

* Check available balance in `previewRewards`

* Chore: forge fmt

---------

Co-authored-by: Daniel Von Fange <daniel@leancoder.com>
---
 contracts/FixedRateRewardsSource.sol       | 10 ++++++++--
 tests/staking/FixedRateRewardsSource.t.sol | 12 ++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/contracts/FixedRateRewardsSource.sol b/contracts/FixedRateRewardsSource.sol
index 0727b0f1..4ef9f431 100644
--- 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
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);