internal functions return validator state

ava-labs · Jan 15, 2025 · 2804658 · 2804658
1 parent 2465e62
commit 2804658
Show file tree

Hide file tree

Showing 7 changed files with 51 additions and 37 deletions.
diff --git a/abi-bindings/go/validator-manager/ERC20TokenStakingManager/ERC20TokenStakingManager.go b/abi-bindings/go/validator-manager/ERC20TokenStakingManager/ERC20TokenStakingManager.go
diff --git a/abi-bindings/go/validator-manager/NativeTokenStakingManager/NativeTokenStakingManager.go b/abi-bindings/go/validator-manager/NativeTokenStakingManager/NativeTokenStakingManager.go
diff --git a/abi-bindings/go/validator-manager/PoAValidatorManager/PoAValidatorManager.go b/abi-bindings/go/validator-manager/PoAValidatorManager/PoAValidatorManager.go
diff --git a/contracts/validator-manager/ACP99Manager.sol b/contracts/validator-manager/ACP99Manager.sol
@@ -143,6 +143,7 @@ abstract contract ACP99Manager {
      * @param disableOwner The disable owner of the validator.
      * @param weight The weight of the node on the L1.
      * @return validationID The ID of the registered validator.
+     * @return validator The registered validator.
      */
     function _initiateValidatorRegistration(
         bytes memory nodeID,
@@ -151,7 +152,7 @@ abstract contract ACP99Manager {
         PChainOwner memory remainingBalanceOwner,
         PChainOwner memory disableOwner,
         uint64 weight
-    ) internal virtual returns (bytes32 validationID);
+    ) internal virtual returns (bytes32 validationID, Validator memory validator);
 
     /**
      * @notice Completes the validator registration process by returning an acknowledgement of the registration of a
@@ -174,8 +175,12 @@ abstract contract ACP99Manager {
      * Emits an {InitiatedValidatorRemoval} on success.
      *
      * @param validationID The ID of the validator to remove.
+     * @return validator The validator to remove.
      */
-    function _initiateValidatorRemoval(bytes32 validationID) internal virtual;
+    function _initiateValidatorRemoval(bytes32 validationID)
+        internal
+        virtual
+        returns (Validator memory validator);
 
     /**
      * @notice Completes validator removal by consuming an RegisterL1ValidatorMessage from the P-Chain acknowledging
@@ -198,13 +203,13 @@ abstract contract ACP99Manager {
      *
      * @param validationID The ID of the validator to modify.
      * @param weight The new weight of the validator.
-     * @return nonce The validator nonce associated with the weight change.
      * @return messageID The ID of the L1ValidatorWeightMessage used to update the validator's weight.
+     * @return validator The updated validator.
      */
     function _initiateValidatorWeightUpdate(
         bytes32 validationID,
         uint64 weight
-    ) internal virtual returns (uint64 nonce, bytes32 messageID);
+    ) internal virtual returns (bytes32 messageID, Validator memory validator);
 
     /**
      * @notice Completes the validator weight update process by consuming a L1ValidatorWeightMessage from the P-Chain

diff --git a/contracts/validator-manager/PoAValidatorManager.sol b/contracts/validator-manager/PoAValidatorManager.sol
@@ -55,7 +55,7 @@ contract PoAValidatorManager is IPoAValidatorManager, ValidatorManager, OwnableU
         ValidatorRegistrationInput calldata registrationInput,
         uint64 weight
     ) external onlyOwner returns (bytes32 validationID) {
-        return _initiateValidatorRegistration({
+        (validationID,) = _initiateValidatorRegistration({
             nodeID: registrationInput.nodeID,
             blsPublicKey: registrationInput.blsPublicKey,
             registrationExpiry: registrationInput.registrationExpiry,

diff --git a/contracts/validator-manager/PoSValidatorManager.sol b/contracts/validator-manager/PoSValidatorManager.sol
@@ -323,10 +323,7 @@ abstract contract PoSValidatorManager is
     ) internal returns (bool) {
         PoSValidatorManagerStorage storage $ = _getPoSValidatorManagerStorage();
 
-        // TODONOW: Functions that modify Validator state invalid any references to the Validator object.
-        // How can we make this safer to use?
-        _initiateValidatorRemoval(validationID);
-        Validator memory validator = getValidator(validationID);
+        Validator memory validator = _initiateValidatorRemoval(validationID);
 
         // Non-PoS validators are required to boostrap the network, but are not eligible for rewards.
         if (!_isPoSValidator(validationID)) {
@@ -482,7 +479,7 @@ abstract contract PoSValidatorManager is
         uint256 lockedValue = _lock(stakeAmount);
 
         uint64 weight = valueToWeight(lockedValue);
-        bytes32 validationID = _initiateValidatorRegistration({
+        (bytes32 validationID,) = _initiateValidatorRegistration({
             nodeID: registrationInput.nodeID,
             blsPublicKey: registrationInput.blsPublicKey,
             registrationExpiry: registrationInput.registrationExpiry,
@@ -559,10 +556,10 @@ abstract contract PoSValidatorManager is
             revert MaxWeightExceeded(newValidatorWeight);
         }
 
-        (uint64 nonce, bytes32 messageID) =
-            _initiateValidatorWeightUpdate(validationID, newValidatorWeight);
+        bytes32 messageID;
+        (messageID, validator) = _initiateValidatorWeightUpdate(validationID, newValidatorWeight);
 
-        bytes32 delegationID = keccak256(abi.encodePacked(validationID, nonce));
+        bytes32 delegationID = keccak256(abi.encodePacked(validationID, validator.sentNonce));
 
         // 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
@@ -572,14 +569,14 @@ abstract contract PoSValidatorManager is
         $._delegatorStakes[delegationID].validationID = validationID;
         $._delegatorStakes[delegationID].weight = weight;
         $._delegatorStakes[delegationID].startTime = 0;
-        $._delegatorStakes[delegationID].startingNonce = nonce;
+        $._delegatorStakes[delegationID].startingNonce = validator.sentNonce;
         $._delegatorStakes[delegationID].endingNonce = 0;
 
         emit DelegatorAdded({
             delegationID: delegationID,
             validationID: validationID,
             delegatorAddress: delegatorAddress,
-            nonce: nonce,
+            nonce: validator.sentNonce,
             validatorWeight: newValidatorWeight,
             delegatorWeight: weight,
             setWeightMessageID: messageID
@@ -754,9 +751,9 @@ abstract contract PoSValidatorManager is
             // initialize the removal.
             $._delegatorStakes[delegationID].status = DelegatorStatus.PendingRemoved;
 
-            ($._delegatorStakes[delegationID].endingNonce,) =
+            (, validator) =
                 _initiateValidatorWeightUpdate(validationID, validator.weight - delegator.weight);
-
+            $._delegatorStakes[delegationID].endingNonce = validator.sentNonce;
             uint256 reward =
                 _calculateAndSetDelegationReward(delegator, rewardRecipient, delegationID);
 

diff --git a/contracts/validator-manager/ValidatorManager.sol b/contracts/validator-manager/ValidatorManager.sol
@@ -252,7 +252,7 @@ abstract contract ValidatorManager is
         PChainOwner memory remainingBalanceOwner,
         PChainOwner memory disableOwner,
         uint64 weight
-    ) internal virtual override initializedValidatorSet returns (bytes32) {
+    ) internal virtual override initializedValidatorSet returns (bytes32, Validator memory) {
         ValidatorManagerStorage storage $ = _getValidatorManagerStorage();
 
         if (
@@ -302,19 +302,21 @@ abstract contract ValidatorManager is
 
         // Submit the message to the Warp precompile.
         bytes32 messageID = WARP_MESSENGER.sendWarpMessage(registerL1ValidatorMessage);
-        $._validationPeriods[validationID].status = ValidatorStatus.PendingAdded;
-        $._validationPeriods[validationID].nodeID = nodeID;
-        $._validationPeriods[validationID].startingWeight = weight;
-        $._validationPeriods[validationID].sentNonce = 0;
-        $._validationPeriods[validationID].weight = weight;
-        $._validationPeriods[validationID].startTime = 0; // The validation period only starts once the registration is acknowledged.
-        $._validationPeriods[validationID].endTime = 0;
+        Validator memory validator = $._validationPeriods[validationID];
+        validator.status = ValidatorStatus.PendingAdded;
+        validator.nodeID = nodeID;
+        validator.startingWeight = weight;
+        validator.sentNonce = 0;
+        validator.weight = weight;
+        validator.startTime = 0; // The validation period only starts once the registration is acknowledged.
+        validator.endTime = 0;
+        $._validationPeriods[validationID] = validator;
 
         emit InitiatedValidatorRegistration(
             validationID, nodeID, messageID, registrationExpiry, weight
         );
 
-        return validationID;
+        return (validationID, validator);
     }
 
     /**
@@ -433,7 +435,12 @@ abstract contract ValidatorManager is
      * Any rewards for this validation period will stop accruing when this function is called.
      * @param validationID The ID of the validation period being ended.
      */
-    function _initiateValidatorRemoval(bytes32 validationID) internal virtual override {
+    function _initiateValidatorRemoval(bytes32 validationID)
+        internal
+        virtual
+        override
+        returns (Validator memory)
+    {
         ValidatorManagerStorage storage $ = _getValidatorManagerStorage();
 
         // Ensure the validation period is active.
@@ -451,15 +458,20 @@ abstract contract ValidatorManager is
         // on the P-Chain.
         validator.endTime = uint64(block.timestamp);
 
+        uint64 endingWeight = validator.weight;
+
         // Save the validator updates.
         $._validationPeriods[validationID] = validator;
 
-        (, bytes32 messageID) = _initiateValidatorWeightUpdate(validationID, 0);
+        bytes32 messageID;
+        (messageID, validator) = _initiateValidatorWeightUpdate(validationID, 0);
 
         // Emit the event to signal the start of the validator removal process.
         emit InitiatedValidatorRemoval(
-            validationID, messageID, validator.weight, uint64(block.timestamp)
+            validationID, messageID, endingWeight, uint64(block.timestamp)
         );
+
+        return validator;
     }
 
     /**
@@ -558,7 +570,7 @@ abstract contract ValidatorManager is
     function _initiateValidatorWeightUpdate(
         bytes32 validationID,
         uint64 newWeight
-    ) internal virtual override returns (uint64, bytes32) {
+    ) internal virtual override returns (bytes32, Validator memory) {
         ValidatorManagerStorage storage $ = _getValidatorManagerStorage();
         uint64 validatorWeight = $._validationPeriods[validationID].weight;
 
@@ -581,7 +593,7 @@ abstract contract ValidatorManager is
             weight: newWeight
         });
 
-        return (nonce, messageID);
+        return (messageID, $._validationPeriods[validationID]);
     }
 
     function _getChurnPeriodSeconds() internal view returns (uint64) {