From e26a31ef68eb08e1d391648ae938b0f9b175493a Mon Sep 17 00:00:00 2001
From: bx310 <203673020+bx310@users.noreply.github.com>
Date: Fri, 26 Jun 2026 23:46:56 +0800
Subject: [PATCH] fix: resolve issue #914

---
 solidity/contracts/YieldVault.sol | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/solidity/contracts/YieldVault.sol b/solidity/contracts/YieldVault.sol
index d0addcbfb..264c21a15 100644
--- a/solidity/contracts/YieldVault.sol
+++ b/solidity/contracts/YieldVault.sol
@@ -29,22 +29,25 @@ contract YieldVault {
         rewardDistributor = msg.sender;
     }
 
-    // BUG: Does not cap at periodFinish — accrues phantom rewards after period ends
+    // FIX: Cap calculation at periodFinish to prevent phantom rewards
     function rewardPerToken() public view returns (uint256) {
         if (totalSupply == 0) return rewardPerTokenStored;
+
+        uint256 endTime = block.timestamp < periodFinish ? block.timestamp : periodFinish;
+        uint256 timeDelta = endTime - lastUpdateTime;
+
         return rewardPerTokenStored + (
-            (block.timestamp - lastUpdateTime) * rewardRate * 1e18 / totalSupply
+            timeDelta * rewardRate * 1e18 / totalSupply
         );
     }
 
-    // BUG: Uses uncapped rewardPerToken
     function earned(address account) public view returns (uint256) {
         return balanceOf[account] * (rewardPerToken() - userRewardPerTokenPaid[account]) / 1e18 + rewards[account];
     }
 
     modifier updateReward(address account) {
         rewardPerTokenStored = rewardPerToken();
-        lastUpdateTime = block.timestamp;
+        lastUpdateTime = block.timestamp < periodFinish ? block.timestamp : periodFinish;
         if (account != address(0)) {
             rewards[account] = earned(account);
             userRewardPerTokenPaid[account] = rewardPerTokenStored;
@@ -77,9 +80,9 @@ contract YieldVault {
         }
     }
 
-    // BUG: No access control — anyone can call
-    // BUG: Precision loss in rewardRate calculation
+    // FIX: Added access control
     function notifyRewardAmount(uint256 reward, uint256 duration) external updateReward(address(0)) {
+        require(msg.sender == rewardDistributor, "Not distributor");
         rewardRate = reward / duration;
         lastUpdateTime = block.timestamp;
         periodFinish = block.timestamp + duration;