Feat/on chain accounting

foundry/src/FoxStakingV1.sol

@@ -7,14 +7,16 @@ import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/U
 import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import {PausableUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.sol";
+import {ReentrancyGuardUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/ReentrancyGuardUpgradeable.sol";
 import {StakingInfo} from "./StakingInfo.sol";
 import {UnstakingRequest} from "./UnstakingRequest.sol";
 
 contract FoxStakingV1 is
     Initializable,
     PausableUpgradeable,
     UUPSUpgradeable,
-    OwnableUpgradeable
+    OwnableUpgradeable,
+    ReentrancyGuardUpgradeable
 {
     using SafeERC20 for IERC20;
     IERC20 public foxToken;
@@ -24,6 +26,14 @@ contract FoxStakingV1 is
     bool public unstakingPaused;
     uint256 public cooldownPeriod;
 
+    uint256 public totalStaked;
+    uint256 public totalCoolingDown;
+
+    uint256 public constant REWARD_RATE = 1_000_000_000;
+    uint256 public constant WAD = 1e18;
+    uint256 public lastUpdateTime;
+    uint256 public rewardPerTokenStored;
+
     event UpdateCooldownPeriod(uint256 newCooldownPeriod);
     event Stake(
         address indexed account,
@@ -56,6 +66,8 @@ contract FoxStakingV1 is
         withdrawalsPaused = false;
         unstakingPaused = false;
         cooldownPeriod = 28 days;
+        rewardPerTokenStored = rewardPerToken();
+        lastUpdateTime = block.timestamp;
     }
 
     function _authorizeUpgrade(
@@ -121,19 +133,60 @@ contract FoxStakingV1 is
         _;
     }
 
+    /// @notice Updates all variables when changes to staking amounts are made.
+    modifier updateReward(address account) {
+        rewardPerTokenStored = rewardPerToken();
+        lastUpdateTime = block.timestamp;
+        StakingInfo storage info = stakingInfo[account];
+        info.earnedRewards = earned(account);
+        info.rewardPerTokenStored = rewardPerTokenStored;
+        _;
+    }
+
+    /// @notice Sets the cooldown period for unstaking requests.
+    /// @param newCooldownPeriod The new cooldown period to be set.
     function setCooldownPeriod(uint256 newCooldownPeriod) external onlyOwner {
         cooldownPeriod = newCooldownPeriod;
         emit UpdateCooldownPeriod(newCooldownPeriod);
     }
 
+    /// @notice Returns the current amount of reward allocated per staked token.
+    function rewardPerToken() public view returns (uint256) {
+        if (totalStaked == 0) {
+            return rewardPerTokenStored;
+        }
+        return
+            rewardPerTokenStored +
+            (((block.timestamp - lastUpdateTime) * REWARD_RATE * 1e18) /
+                totalStaked);
+    }
+
+    /// Returns the total reward earnings associated with a given address for its
+    /// entire lifetime of staking.
+    /// @param account The address we're getting the earned rewards for.
+    function earned(address account) public view returns (uint256) {
+        StakingInfo memory info = stakingInfo[account];
+        return
+            (info.stakingBalance *
+                (rewardPerToken() - info.rewardPerTokenStored)) /
+            1e18 +
+            info.earnedRewards;
+    }
+
     /// @notice Allows a user to stake a specified amount of FOX tokens and assign a RUNE address for rewards - which can be changed later on.
     /// This has to be initiated by the user itself i.e msg.sender only, cannot be called by an address for another
     /// @param amount The amount of FOX tokens to be staked.
     /// @param runeAddress The RUNE address to be associated with the user's staked FOX position.
     function stake(
         uint256 amount,
         string memory runeAddress
-    ) external whenNotPaused whenStakingNotPaused {
+    )
+        external
+        whenNotPaused
+        whenStakingNotPaused
+        updateReward(msg.sender)
+        nonReentrant
+    {
         require(
             bytes(runeAddress).length == 43,
             "Rune address must be 43 characters"
@@ -144,6 +197,7 @@ contract FoxStakingV1 is
         StakingInfo storage info = stakingInfo[msg.sender];
         info.stakingBalance += amount;
         info.runeAddress = runeAddress;
+        totalStaked += amount;
 
         emit Stake(msg.sender, amount, runeAddress);
     }
@@ -153,7 +207,13 @@ contract FoxStakingV1 is
     /// @param amount The amount of FOX tokens to be unstaked.
     function unstake(
         uint256 amount
-    ) external whenNotPaused whenUnstakingNotPaused {
+    )
+        external
+        whenNotPaused
+        whenUnstakingNotPaused
+        updateReward(msg.sender)
+        nonReentrant
+    {
         require(amount > 0, "Cannot unstake 0");
         StakingInfo storage info = stakingInfo[msg.sender];
 
@@ -166,6 +226,8 @@ contract FoxStakingV1 is
         // Set staking / unstaking amounts
         info.stakingBalance -= amount;
         info.unstakingBalance += amount;
+        totalStaked -= amount;
+        totalCoolingDown += amount;
 
         UnstakingRequest memory unstakingRequest = UnstakingRequest({
             unstakingBalance: amount,
@@ -181,7 +243,7 @@ contract FoxStakingV1 is
     /// @param index The index of the claim to withdraw
     function withdraw(
         uint256 index
-    ) public whenNotPaused whenWithdrawalsNotPaused {
+    ) public whenNotPaused whenWithdrawalsNotPaused nonReentrant {
         StakingInfo storage info = stakingInfo[msg.sender];
         require(
             info.unstakingRequests.length > 0,
@@ -211,6 +273,7 @@ contract FoxStakingV1 is
             delete info.unstakingRequests;
         }
         info.unstakingBalance -= unstakingRequest.unstakingBalance;
+        totalCoolingDown -= unstakingRequest.unstakingBalance;
         foxToken.safeTransfer(msg.sender, unstakingRequest.unstakingBalance);
         emit Withdraw(msg.sender, unstakingRequest.unstakingBalance);
     }

foundry/src/StakingInfo.sol

@@ -6,6 +6,8 @@ import {UnstakingRequest} from "./UnstakingRequest.sol";
 struct StakingInfo {
     uint256 stakingBalance;
     uint256 unstakingBalance;
+    uint256 earnedRewards; // earnedRewards for user since last update
+    uint256 rewardPerTokenStored; // user level reward per token stored
     string runeAddress;
     UnstakingRequest[] unstakingRequests;
 }