From 37f70481a6de87e83b8dbf7cd357efae6c57baf1 Mon Sep 17 00:00:00 2001
From: Ana Julia <anajuliabit@gmail.com>
Date: Mon, 17 Jun 2024 09:38:51 -0300
Subject: [PATCH] updateRewards

---
 src/Staking.sol | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/Staking.sol b/src/Staking.sol
index 6ade9a0..12e6493 100644
--- a/src/Staking.sol
+++ b/src/Staking.sol
@@ -121,7 +121,26 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         // Distribute rewards
         rewardsDistributor.distributeRewards();
 
-        if (caller != address(0)) {}
+        address[] rewardTokens = rewardsDistributor.rewardTokens();
+
+        if (caller != address(0)) {
+            for (uint256 i = 0; i < rewardTokens.length; i++) {
+                address token = rewardTokens[i];
+                // ignore staking token as it will be auto compounded
+                if (token == address(stakingToken)) {
+                    continue;
+                }
+
+                uint256 rewardPerToken = rewardPerToken(token);
+
+                keyperRewards[caller][token].earned += (balanceOf(caller) *
+                    (rewardPerToken -
+                        keyperRewards[caller][token].userRewardPerTokenPaid));
+
+                keyperRewards[caller][token]
+                    .userRewardPerTokenPaid = rewardPerToken;
+            }
+        }
 
         // If the caller has no assets or is the zero address, skip compound
         if (assetsBefore == 0 || caller == address(0)) {