Merge branch 'staking-contract' into churn-tracker

ava-labs · Sep 12, 2024 · e367998 · e367998
2 parents 4281078 + 006de7e
commit e367998
Show file tree

Hide file tree

Showing 7 changed files with 98 additions and 161 deletions.
diff --git a/abi-bindings/go/staking/ERC20TokenStakingManager/ERC20TokenStakingManager.go b/abi-bindings/go/staking/ERC20TokenStakingManager/ERC20TokenStakingManager.go
diff --git a/abi-bindings/go/staking/NativeTokenStakingManager/NativeTokenStakingManager.go b/abi-bindings/go/staking/NativeTokenStakingManager/NativeTokenStakingManager.go
diff --git a/abi-bindings/go/staking/PoAValidatorManager/PoAValidatorManager.go b/abi-bindings/go/staking/PoAValidatorManager/PoAValidatorManager.go
diff --git a/contracts/staking/PoSValidatorManager.sol b/contracts/staking/PoSValidatorManager.sol
@@ -31,10 +31,6 @@ abstract contract PoSValidatorManager is IPoSValidatorManager, ValidatorManager
         IRewardCalculator _rewardCalculator;
         /// @notice Maps the delegationID to the delegator information.
         mapping(bytes32 delegationID => Delegator) _delegatorStakes;
-        /// @notice Maps the delegationID to pending register delegator messages.
-        mapping(bytes32 delegationID => bytes) _pendingRegisterDelegatorMessages;
-        /// @notice Maps the delegationID to pending end delegator messages.
-        mapping(bytes32 delegationID => bytes) _pendingEndDelegatorMessages;
     }
     // solhint-enable private-vars-leading-underscore
 
@@ -173,14 +169,15 @@ abstract contract PoSValidatorManager is IPoSValidatorManager, ValidatorManager
 
         // Construct the delegation ID. This is guaranteed to be unique since it is
         // constructed using a new nonce.
-        uint64 nonce = _getAndIncrementNonce(validationID);
+        uint64 nonce = _incrementAndGetNonce(validationID);
         bytes32 delegationID = keccak256(abi.encodePacked(validationID, delegatorAddress, nonce));
 
         // Submit the message to the Warp precompile.
-        bytes memory setValidatorWeightPayload = ValidatorMessages
-            .packSetSubnetValidatorWeightMessage(validationID, nonce, newValidatorWeight);
-        $._pendingRegisterDelegatorMessages[delegationID] = setValidatorWeightPayload;
-        bytes32 messageID = WARP_MESSENGER.sendWarpMessage(setValidatorWeightPayload);
+        bytes32 messageID = WARP_MESSENGER.sendWarpMessage(
+            ValidatorMessages.packSetSubnetValidatorWeightMessage(
+                validationID, nonce, newValidatorWeight
+            )
+        );
 
         // Store the delegation information. Set the delegator status to pending added,
         // so that it can be properly started in the complete step, even if the delivered
@@ -208,22 +205,13 @@ abstract contract PoSValidatorManager is IPoSValidatorManager, ValidatorManager
         return delegationID;
     }
 
-    function resendDelegatorRegistration(bytes32 delegationID) external {
-        _checkPendingRegisterDelegatorMessages(delegationID);
-        PoSValidatorManagerStorage storage $ = _getPoSValidatorManagerStorage();
-        // Submit the message to the Warp precompile.
-        WARP_MESSENGER.sendWarpMessage($._pendingRegisterDelegatorMessages[delegationID]);
-    }
-
     function completeDelegatorRegistration(uint32 messageIndex, bytes32 delegationID) external {
         PoSValidatorManagerStorage storage $ = _getPoSValidatorManagerStorage();
 
         // Unpack the Warp message
         WarpMessage memory warpMessage = _getPChainWarpMessage(messageIndex);
         (bytes32 validationID, uint64 nonce,) =
             ValidatorMessages.unpackSubnetValidatorWeightUpdateMessage(warpMessage.payload);
-        _checkPendingRegisterDelegatorMessages(delegationID);
-        delete $._pendingRegisterDelegatorMessages[delegationID];
 
         Validator memory validator = _getValidator(validationID);
 
@@ -283,7 +271,7 @@ abstract contract PoSValidatorManager is IPoSValidatorManager, ValidatorManager
         require(
             delegator.owner == _msgSender(), "PoSValidatorManager: delegation not owned by sender"
         );
-        uint64 nonce = _getAndIncrementNonce(validationID);
+        uint64 nonce = _incrementAndGetNonce(validationID);
 
         // Check that removing this delegator would not exceed the maximum churn rate.
         // TODO this check won't be necessary for a delegator whose validator has already initialized ending their validation period.
@@ -304,10 +292,11 @@ abstract contract PoSValidatorManager is IPoSValidatorManager, ValidatorManager
         _setValidatorWeight(validationID, newValidatorWeight);
 
         // Submit the message to the Warp precompile.
-        bytes memory setValidatorWeightPayload = ValidatorMessages
-            .packSetSubnetValidatorWeightMessage(validationID, nonce, newValidatorWeight);
-        $._pendingEndDelegatorMessages[delegationID] = setValidatorWeightPayload;
-        bytes32 messageID = WARP_MESSENGER.sendWarpMessage(setValidatorWeightPayload);
+        bytes32 messageID = WARP_MESSENGER.sendWarpMessage(
+            ValidatorMessages.packSetSubnetValidatorWeightMessage(
+                validationID, nonce, newValidatorWeight
+            )
+        );
 
         emit DelegatorRemovalInitialized({
             delegationID: delegationID,
@@ -319,10 +308,31 @@ abstract contract PoSValidatorManager is IPoSValidatorManager, ValidatorManager
         });
     }
 
-    function resendEndDelegation(bytes32 delegationID) external {
-        _checkPendingEndDelegatorMessage(delegationID);
+    /**
+     * @dev Resending the latest validator weight with the latest nonce is safe because all weight changes are
+     * cumulative, so the latest weight change will always include the weight change for any added delegators.
+     */
+    function resendUpdateDelegation(bytes32 delegationID) external {
         PoSValidatorManagerStorage storage $ = _getPoSValidatorManagerStorage();
-        WARP_MESSENGER.sendWarpMessage($._pendingEndDelegatorMessages[delegationID]);
+        Delegator memory delegator = $._delegatorStakes[delegationID];
+        require(
+            delegator.status == DelegatorStatus.PendingAdded
+                || delegator.status == DelegatorStatus.PendingRemoved,
+            "PoSValidatorManager: delegation status not pending"
+        );
+
+        Validator memory validator = _getValidator(delegator.validationID);
+        require(
+            validator.messageNonce != 0,
+            "PoSValidatorManager: could not find validator for delegation ID"
+        );
+
+        // Submit the message to the Warp precompile.
+        WARP_MESSENGER.sendWarpMessage(
+            ValidatorMessages.packSetSubnetValidatorWeightMessage(
+                delegator.validationID, validator.messageNonce, validator.weight
+            )
+        );
     }
 
     function completeEndDelegation(uint32 messageIndex, bytes32 delegationID) external {
@@ -332,8 +342,6 @@ abstract contract PoSValidatorManager is IPoSValidatorManager, ValidatorManager
         WarpMessage memory warpMessage = _getPChainWarpMessage(messageIndex);
         (bytes32 validationID, uint64 nonce,) =
             ValidatorMessages.unpackSubnetValidatorWeightUpdateMessage(warpMessage.payload);
-        _checkPendingEndDelegatorMessage(delegationID);
-        delete $._pendingEndDelegatorMessages[delegationID];
 
         Validator memory validator = _getValidator(validationID);
         // The received nonce should be no greater than the highest sent nonce. This should never
@@ -367,22 +375,4 @@ abstract contract PoSValidatorManager is IPoSValidatorManager, ValidatorManager
 
         emit DelegationEnded(delegationID, validationID, nonce);
     }
-
-    function _checkPendingEndDelegatorMessage(bytes32 delegationID) private view {
-        PoSValidatorManagerStorage storage $ = _getPoSValidatorManagerStorage();
-        require(
-            $._pendingEndDelegatorMessages[delegationID].length > 0
-                && $._delegatorStakes[delegationID].status == DelegatorStatus.PendingRemoved,
-            "PoSValidatorManager: delegation removal not pending"
-        );
-    }
-
-    function _checkPendingRegisterDelegatorMessages(bytes32 delegationID) private view {
-        PoSValidatorManagerStorage storage $ = _getPoSValidatorManagerStorage();
-        require(
-            $._pendingRegisterDelegatorMessages[delegationID].length > 0
-                && $._delegatorStakes[delegationID].status == DelegatorStatus.PendingAdded,
-            "PoSValidatorManager: delegation registration not pending"
-        );
-    }
 }
diff --git a/contracts/staking/ValidatorManager.sol b/contracts/staking/ValidatorManager.sol
@@ -45,8 +45,6 @@ abstract contract ValidatorManager is
         ValidatorChurnPeriod _churnTracker;
         /// @notice Maps the validationID to the registration message such that the message can be re-sent if needed.
         mapping(bytes32 => bytes) _pendingRegisterValidationMessages;
-        /// @notice Maps the validationID to the end validation message such that the message can be re-sent if needed.
-        mapping(bytes32 => bytes) _pendingEndValidationMessages;
         /// @notice Maps the validationID to the validator information.
         mapping(bytes32 => Validator) _validationPeriods;
         /// @notice Maps the nodeID to the validationID for active validation periods.
@@ -166,9 +164,6 @@ abstract contract ValidatorManager is
             endedAt: 0
         });
 
-        // Increment the nonce for the next usage.
-        _getAndIncrementNonce(validationID);
-
         emit ValidationPeriodCreated(validationID, nodeID, messageID, weight, registrationExpiry);
 
         return validationID;
@@ -255,8 +250,7 @@ abstract contract ValidatorManager is
 
         // Submit the message to the Warp precompile.
         bytes memory setValidatorWeightPayload = ValidatorMessages
-            .packSetSubnetValidatorWeightMessage(validationID, _getAndIncrementNonce(validationID), 0);
-        $._pendingEndValidationMessages[validationID] = setValidatorWeightPayload;
+            .packSetSubnetValidatorWeightMessage(validationID, _incrementAndGetNonce(validationID), 0);
 
         bytes32 messageID = WARP_MESSENGER.sendWarpMessage(setValidatorWeightPayload);
 
@@ -271,14 +265,18 @@ abstract contract ValidatorManager is
     // solhint-disable-next-line
     function resendEndValidatorMessage(bytes32 validationID) external {
         ValidatorManagerStorage storage $ = _getValidatorManagerStorage();
+        Validator memory validator = $._validationPeriods[validationID];
 
         require(
-            $._pendingEndValidationMessages[validationID].length > 0
-                && $._validationPeriods[validationID].status == ValidatorStatus.PendingRemoved,
+            validator.status == ValidatorStatus.PendingRemoved,
             "ValidatorManager: Validator not pending removal"
         );
 
-        WARP_MESSENGER.sendWarpMessage($._pendingEndValidationMessages[validationID]);
+        WARP_MESSENGER.sendWarpMessage(
+            ValidatorMessages.packSetSubnetValidatorWeightMessage(
+                validationID, validator.messageNonce, 0
+            )
+        );
     }
 
     /**
@@ -312,7 +310,6 @@ abstract contract ValidatorManager is
 
         if (validator.status == ValidatorStatus.PendingRemoved) {
             endStatus = ValidatorStatus.Completed;
-            delete $._pendingEndValidationMessages[validationID];
         } else {
             endStatus = ValidatorStatus.Invalidated;
         }
@@ -339,11 +336,10 @@ abstract contract ValidatorManager is
         return _getValidator(validationID).weight;
     }
 
-    function _getAndIncrementNonce(bytes32 validationID) internal returns (uint64) {
+    function _incrementAndGetNonce(bytes32 validationID) internal returns (uint64) {
         ValidatorManagerStorage storage $ = _getValidatorManagerStorage();
-        uint64 currentNonce = $._validationPeriods[validationID].messageNonce;
         $._validationPeriods[validationID].messageNonce++;
-        return currentNonce;
+        return $._validationPeriods[validationID].messageNonce;
     }
 
     function _getPChainWarpMessage(uint32 messageIndex)

diff --git a/contracts/staking/interfaces/IPoSValidatorManager.sol b/contracts/staking/interfaces/IPoSValidatorManager.sol
@@ -111,13 +111,6 @@ interface IPoSValidatorManager is IValidatorManager {
         uint32 messageIndex
     ) external;
 
-    /**
-     * @notice Resubmits a delegator registration message to be sent to the P-Chain.
-     * Only necessary if the original message can't be delivered due to validator churn.
-     * @param delegationID The ID of the delegation being registered.
-     */
-    function resendDelegatorRegistration(bytes32 delegationID) external;
-
     /**
      * @notice Completes the delegator registration process by returning an acknowledgement of the registration of a
      * validationID from the P-Chain. After this function is called, the validator's weight is updated in the contract state.
@@ -149,11 +142,11 @@ interface IPoSValidatorManager is IValidatorManager {
     ) external;
 
     /**
-     * @notice Resubmits a delegator end message to be sent to the P-Chain.
+     * @notice Resubmits a delegator registration or delegator end message to be sent to the P-Chain.
      * Only necessary if the original message can't be delivered due to validator churn.
-     * @param delegationID The ID of the delegation being removed.
+     * @param delegationID The ID of the delegation.
      */
-    function resendEndDelegation(bytes32 delegationID) external;
+    function resendUpdateDelegation(bytes32 delegationID) external;
 
     /**
      * @notice Completes the process of ending a delegation by receiving an acknowledgement from the P-Chain.

diff --git a/contracts/staking/tests/PoSValidatorManagerTests.t.sol b/contracts/staking/tests/PoSValidatorManagerTests.t.sol
@@ -201,7 +201,7 @@ abstract contract PoSValidatorManagerTest is ValidatorManagerTest {
             validationID, 1, DEFAULT_WEIGHT + DEFAULT_DELEGATOR_WEIGHT
         );
         _mockSendWarpMessage(setValidatorWeightPayload, bytes32(0));
-        posValidatorManager.resendDelegatorRegistration(delegationID);
+        posValidatorManager.resendUpdateDelegation(delegationID);
     }
 
     function testCompleteDelegatorRegistration() public {
@@ -381,7 +381,7 @@ abstract contract PoSValidatorManagerTest is ValidatorManagerTest {
         bytes memory setValidatorWeightPayload =
             ValidatorMessages.packSetSubnetValidatorWeightMessage(validationID, 2, DEFAULT_WEIGHT);
         _mockSendWarpMessage(setValidatorWeightPayload, bytes32(0));
-        posValidatorManager.resendEndDelegation(delegationID);
+        posValidatorManager.resendUpdateDelegation(delegationID);
     }
 
     function testCompleteEndDelegation() public {