Add Delegator rewards to different address

ava-labs · Nov 18, 2024 · 79a1984 · 79a1984
1 parent d75569e
commit 79a1984
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/contracts/validator-manager/PoSValidatorManager.sol b/contracts/validator-manager/PoSValidatorManager.sol
@@ -709,7 +709,9 @@ abstract contract PoSValidatorManager is
 
             uint256 reward = _calculateDelegationReward(delegator);
             $._redeemableDelegatorRewards[delegationID] = reward;
-            $._delegatorRewardRecipients[delegationID] = rewardRecipient;
+            if (rewardRecipient != address(0)) {
+                $._delegatorRewardRecipients[delegationID] = rewardRecipient;
+            }
 
             emit DelegatorRemovalInitialized({
                 delegationID: delegationID,
@@ -718,6 +720,9 @@ abstract contract PoSValidatorManager is
             return (reward > 0);
         } else if (validator.status == ValidatorStatus.Completed) {
             $._redeemableDelegatorRewards[delegationID] = _calculateDelegationReward(delegator);
+            if (rewardRecipient != address(0)) {
+                $._delegatorRewardRecipients[delegationID] = rewardRecipient;
+            }
 
             _completeEndDelegation(delegationID);
             // If the validator has completed, then no further uptimes may be submitted, so we always
@@ -846,7 +851,13 @@ abstract contract PoSValidatorManager is
         // Once this function completes, the delegation is completed so we can clear it from state now.
         delete $._delegatorStakes[delegationID];
 
-        address rewardRecipient = delegator.owner;
+        address rewardRecipient = $._delegatorRewardRecipients[delegationID];
+
+        if (rewardRecipient == address(0)) {
+            rewardRecipient = delegator.owner;
+        } else {
+            delete $._delegatorRewardRecipients[delegationID];
+        }
 
         (uint256 delegationRewards, uint256 validatorFees) =
             _withdrawDelegationRewards(rewardRecipient, delegationID, validationID);

diff --git a/contracts/validator-manager/tests/PoSValidatorManagerTests.t.sol b/contracts/validator-manager/tests/PoSValidatorManagerTests.t.sol
@@ -633,6 +633,7 @@ abstract contract PoSValidatorManagerTest is ValidatorManagerTest {
     function testCompleteEndDelegationWithNonDelegatorRewardRecipient() public {
         bytes32 validationID = _registerDefaultValidator();
         bytes32 delegationID = _registerDefaultDelegator(validationID);
+        address rewardRecipient = address(42);
 
         _initializeEndDelegationValidatorActiveWithChecks({
             validationID: validationID,
@@ -644,7 +645,7 @@ abstract contract PoSValidatorManagerTest is ValidatorManagerTest {
             expectedNonce: 2,
             includeUptime: true,
             force: false,
-            rewardRecipient: address(0)
+            rewardRecipient: rewardRecipient
         });
 
         uint256 expectedTotalReward = rewardCalculator.calculateReward({
@@ -669,7 +670,7 @@ abstract contract PoSValidatorManagerTest is ValidatorManagerTest {
             validatorWeight: DEFAULT_WEIGHT,
             expectedValidatorWeight: DEFAULT_WEIGHT,
             expectedNonce: 2,
-            rewardRecipient: address(0)
+            rewardRecipient: rewardRecipient
         });
     }
 
@@ -1603,7 +1604,7 @@ abstract contract PoSValidatorManagerTest is ValidatorManagerTest {
             force: force,
             setValidatorWeightPayload: setValidatorWeightPayload,
             uptimePayload: uptimeMsg,
-            rewardRecipient: address(0)
+            rewardRecipient: rewardRecipient
         });
     }