From a03e74734a23302240819b79235ab43cbf72e5e9 Mon Sep 17 00:00:00 2001
From: charles2023wood <charles2022wood@gmail.com>
Date: Mon, 15 Jul 2024 15:49:04 +0800
Subject: [PATCH] Fix delegator ref when claimReward

Signed-off-by: charles2023wood <charles2022wood@gmail.com>
---
 contracts/BitcoinLSTStake.sol | 10 +++++-----
 contracts/BitcoinStake.sol    |  7 ++++---
 contracts/HashPowerAgent.sol  |  5 +++--
 contracts/PledgeAgent.sol     | 22 +++++++++++-----------
 4 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/contracts/BitcoinLSTStake.sol b/contracts/BitcoinLSTStake.sol
index 300aa9c1..0b7b1210 100644
--- a/contracts/BitcoinLSTStake.sol
+++ b/contracts/BitcoinLSTStake.sol
@@ -256,7 +256,7 @@ contract BitcoinLSTStake is IBitcoinStake, System, IParamSubscriber, ReentrancyG
   /// Claim reward for delegator
   /// @return reward Amount claimed
   function claimReward() external override onlyBtcAgent returns (uint256 reward) {
-    return _updateUserRewards(msg.sender, true);
+    return _updateUserRewards(tx.origin, true);
   }
 
   /*********************** External implementations ***************************/
@@ -437,12 +437,12 @@ contract BitcoinLSTStake is IBitcoinStake, System, IParamSubscriber, ReentrancyG
       user.changeRound = roundTag;
     }
     if (claim) {
-      if (rewardMap[msg.sender] != 0) {
-        reward += rewardMap[msg.sender];
-        rewardMap[msg.sender] = 0;
+      if (rewardMap[userAddress] != 0) {
+        reward += rewardMap[userAddress];
+        rewardMap[userAddress] = 0;
       }
     } else {
-      rewardMap[msg.sender] += reward;
+      rewardMap[userAddress] += reward;
     }
   }
 
diff --git a/contracts/BitcoinStake.sol b/contracts/BitcoinStake.sol
index 60b5bbe7..0cb57154 100644
--- a/contracts/BitcoinStake.sol
+++ b/contracts/BitcoinStake.sol
@@ -236,11 +236,12 @@ contract BitcoinStake is IBitcoinStake, System, IParamSubscriber {
   /// Claim reward for delegator
   /// @return reward Amount claimed
   function claimReward() external override onlyBtcAgent returns (uint256 reward) {
-    reward = rewardMap[msg.sender];
+    address delegator = tx.origin;
+    reward = rewardMap[delegator];
     if (reward != 0) {
-      rewardMap[msg.sender] = 0;
+      rewardMap[delegator] = 0;
     }
-    bytes32[] storage txids = delegatorMap[msg.sender].txids;
+    bytes32[] storage txids = delegatorMap[delegator].txids;
     for (uint256 i = txids.length; i != 0; i--) {
       bytes32 txid = txids[i - 1];
       DepositReceipt storage dr = receiptMap[txid];
diff --git a/contracts/HashPowerAgent.sol b/contracts/HashPowerAgent.sol
index 053ef09a..1a598495 100644
--- a/contracts/HashPowerAgent.sol
+++ b/contracts/HashPowerAgent.sol
@@ -79,9 +79,10 @@ contract HashPowerAgent is IAgent, System, IParamSubscriber {
   /// Claim reward for delegator
   /// @return reward Amount claimed
   function claimReward() external override onlyStakeHub returns (uint256) {
-    uint256 rewardSum = rewardMap[msg.sender];
+    address delegator = tx.origin;
+    uint256 rewardSum = rewardMap[delegator];
     if (rewardSum != 0) {
-      rewardMap[msg.sender] = 0;
+      rewardMap[delegator] = 0;
     }
     return rewardSum;
   }
diff --git a/contracts/PledgeAgent.sol b/contracts/PledgeAgent.sol
index 7b07bda9..3c555531 100644
--- a/contracts/PledgeAgent.sol
+++ b/contracts/PledgeAgent.sol
@@ -283,7 +283,6 @@ contract PledgeAgent is IAgent, System, IParamSubscriber {
     uint256 deposit = undelegateCoin(candidate, msg.sender, amount, false);
     Address.sendValue(payable(msg.sender), deposit);
     emit undelegatedCoin(candidate, msg.sender, deposit);
-    // TODO removeCandidate(msg.sender, address candidate)
   }
 
   /// Transfer coin stake to a new validator
@@ -352,34 +351,35 @@ contract PledgeAgent is IAgent, System, IParamSubscriber {
   }
 
   function claimReward() external override onlyStakeHub returns (uint256) {
+    address delegator = tx.origin;
     uint256 reward;
-    uint256 rewardSum = rewardMap[msg.sender];
+    uint256 rewardSum = rewardMap[delegator];
     if (rewardSum != 0) {
-      rewardMap[msg.sender] = 0;
+      rewardMap[delegator] = 0;
     }
 
-    Delegator storage delegator = delegatorsMap[msg.sender];
-    uint256 candidateSize = delegator.candidates.length;
+    address[] storage candidates = delegatorsMap[delegator].candidates;
+    uint256 candidateSize = candidates.length;
     for (uint256 i = candidateSize; i != 0;) {
       --i;
-      Agent storage a = agentsMap[delegator.candidates[i]];
+      Agent storage a = agentsMap[candidates[i]];
       if (a.rewardSet.length == 0) {
         continue;
       }
-      CoinDelegator storage d = a.cDelegatorMap[msg.sender];
+      CoinDelegator storage d = a.cDelegatorMap[delegator];
       if (d.newDeposit == 0 && d.transferOutDeposit == 0) {
         continue;
       }
       reward = collectCoinReward(a, d, 0xFFFFFFFF);
       rewardSum += reward;
       if (d.newDeposit == 0 && d.transferOutDeposit == 0) {
-        delete a.cDelegatorMap[msg.sender];
-        removeCandidate(msg.sender, delegator.candidates[i]);
+        delete a.cDelegatorMap[delegator];
+        removeCandidate(delegator, candidates[i]);
       }
     }
 
     if (rewardSum != 0) {
-      distributeReward(payable(msg.sender), rewardSum);
+      distributeReward(payable(delegator), rewardSum);
     }
     return rewardSum;
   }
@@ -469,7 +469,7 @@ contract PledgeAgent is IAgent, System, IParamSubscriber {
   /*********************** Internal methods ***************************/
   function distributeReward(address payable delegator, uint256 reward) internal {
     Address.sendValue(delegator, reward);
-    emit claimedReward(delegator, msg.sender, reward, true);
+    emit claimedReward(delegator, tx.origin, reward, true);
   }
 
   function delegateCoin(address agent, address delegator, uint256 deposit) internal returns (uint256) {