eip7251: Bugfix and more withdrawal tests (#14578)

* addressing bug with withdrawals for devnet 5 * changelog * fixing if statement check * adding test * terence's review comments * attempting to fix weird comment formatting * moving back more comments * Update CHANGELOG.md Co-authored-by: Sammy Rosso <15244892+saolyn@users.noreply.github.com> --------- Co-authored-by: Sammy Rosso <15244892+saolyn@users.noreply.github.com>
prysmaticlabs · Oct 31, 2024 · 61c296e · 61c296e
1 parent f264680
commit 61c296e
Show file tree

Hide file tree

Showing 4 changed files with 136 additions and 73 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -58,6 +58,8 @@ The format is based on Keep a Changelog, and this project adheres to Semantic Ve
 - Fix keymanager API should return corrected error format for malformed tokens
 - Fix keymanager API so that get keys returns an empty response instead of a 500 error when using an unsupported keystore.
 - Small log imporvement, removing some redundant or duplicate logs
+- EIP7521 - Fixes withdrawal bug by accounting for pending partial withdrawals and deducting already withdrawn amounts from the sweep balance. [PR](https://github.com/prysmaticlabs/prysm/pull/14578)
+
 
 ### Security
 

diff --git a/beacon-chain/core/blocks/withdrawals.go b/beacon-chain/core/blocks/withdrawals.go
@@ -120,35 +120,36 @@ func ValidateBLSToExecutionChange(st state.ReadOnlyBeaconState, signed *ethpb.Si
 //
 // Spec pseudocode definition:
 //
-//	def process_withdrawals(state: BeaconState, payload: ExecutionPayload) -> None:
-//	    expected_withdrawals, partial_withdrawals_count = get_expected_withdrawals(state)  # [Modified in Electra:EIP7251]
+// def process_withdrawals(state: BeaconState, payload: ExecutionPayload) -> None:
 //
-//	    assert len(payload.withdrawals) == len(expected_withdrawals)
+//	expected_withdrawals, processed_partial_withdrawals_count = get_expected_withdrawals(state) # [Modified in Electra:EIP7251]
 //
-//	    for expected_withdrawal, withdrawal in zip(expected_withdrawals, payload.withdrawals):
-//	        assert withdrawal == expected_withdrawal
-//	        decrease_balance(state, withdrawal.validator_index, withdrawal.amount)
+//	assert len(payload.withdrawals) == len(expected_withdrawals)
 //
-//	    # Update pending partial withdrawals [New in Electra:EIP7251]
-//	    state.pending_partial_withdrawals = state.pending_partial_withdrawals[partial_withdrawals_count:]
+//	for expected_withdrawal, withdrawal in zip(expected_withdrawals, payload.withdrawals):
+//	     assert withdrawal == expected_withdrawal
+//	     decrease_balance(state, withdrawal.validator_index, withdrawal.amount)
 //
-//	    # Update the next withdrawal index if this block contained withdrawals
-//	    if len(expected_withdrawals) != 0:
-//	        latest_withdrawal = expected_withdrawals[-1]
-//	        state.next_withdrawal_index = WithdrawalIndex(latest_withdrawal.index + 1)
+//	# Update pending partial withdrawals [New in Electra:EIP7251]
+//	state.pending_partial_withdrawals = state.pending_partial_withdrawals[processed_partial_withdrawals_count:]
 //
-//	    # Update the next validator index to start the next withdrawal sweep
-//	    if len(expected_withdrawals) == MAX_WITHDRAWALS_PER_PAYLOAD:
-//	        # Next sweep starts after the latest withdrawal's validator index
-//	        next_validator_index = ValidatorIndex((expected_withdrawals[-1].validator_index + 1) % len(state.validators))
-//	        state.next_withdrawal_validator_index = next_validator_index
-//	    else:
-//	        # Advance sweep by the max length of the sweep if there was not a full set of withdrawals
-//	        next_index = state.next_withdrawal_validator_index + MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP
-//	        next_validator_index = ValidatorIndex(next_index % len(state.validators))
-//	        state.next_withdrawal_validator_index = next_validator_index
+//	# Update the next withdrawal index if this block contained withdrawals
+//	if len(expected_withdrawals) != 0:
+//	    latest_withdrawal = expected_withdrawals[-1]
+//	    state.next_withdrawal_index = WithdrawalIndex(latest_withdrawal.index + 1)
+//
+//	# Update the next validator index to start the next withdrawal sweep
+//	if len(expected_withdrawals) == MAX_WITHDRAWALS_PER_PAYLOAD:
+//	    # Next sweep starts after the latest withdrawal's validator index
+//	    next_validator_index = ValidatorIndex((expected_withdrawals[-1].validator_index + 1) % len(state.validators))
+//	    state.next_withdrawal_validator_index = next_validator_index
+//	else:
+//	    # Advance sweep by the max length of the sweep if there was not a full set of withdrawals
+//	    next_index = state.next_withdrawal_validator_index + MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP
+//	    next_validator_index = ValidatorIndex(next_index % len(state.validators))
+//	    state.next_withdrawal_validator_index = next_validator_index
 func ProcessWithdrawals(st state.BeaconState, executionData interfaces.ExecutionData) (state.BeaconState, error) {
-	expectedWithdrawals, partialWithdrawalsCount, err := st.ExpectedWithdrawals()
+	expectedWithdrawals, processedPartialWithdrawalsCount, err := st.ExpectedWithdrawals()
 	if err != nil {
 		return nil, errors.Wrap(err, "could not get expected withdrawals")
 	}
@@ -192,7 +193,7 @@ func ProcessWithdrawals(st state.BeaconState, executionData interfaces.Execution
 	}
 
 	if st.Version() >= version.Electra {
-		if err := st.DequeuePartialWithdrawals(partialWithdrawalsCount); err != nil {
+		if err := st.DequeuePartialWithdrawals(processedPartialWithdrawalsCount); err != nil {
 			return nil, fmt.Errorf("unable to dequeue partial withdrawals from state: %w", err)
 		}
 	}

diff --git a/beacon-chain/state/state-native/getters_withdrawal.go b/beacon-chain/state/state-native/getters_withdrawal.go
@@ -49,56 +49,59 @@ func (b *BeaconState) NextWithdrawalValidatorIndex() (primitives.ValidatorIndex,
 // Spec definition:
 //
 //	def get_expected_withdrawals(state: BeaconState) -> Tuple[Sequence[Withdrawal], uint64]:
-//	    epoch = get_current_epoch(state)
-//	    withdrawal_index = state.next_withdrawal_index
-//	    validator_index = state.next_withdrawal_validator_index
-//	    withdrawals: List[Withdrawal] = []
+//		epoch = get_current_epoch(state)
+//		withdrawal_index = state.next_withdrawal_index
+//		validator_index = state.next_withdrawal_validator_index
+//		withdrawals: List[Withdrawal] = []
+//		processed_partial_withdrawals_count = 0
 //
-//	    # [New in Electra:EIP7251] Consume pending partial withdrawals
-//	    for withdrawal in state.pending_partial_withdrawals:
-//	        if withdrawal.withdrawable_epoch > epoch or len(withdrawals) == MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP:
-//	            break
+//		# [New in Electra:EIP7251] Consume pending partial withdrawals
+//		for withdrawal in state.pending_partial_withdrawals:
+//			if withdrawal.withdrawable_epoch > epoch or len(withdrawals) == MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP:
+//				break
 //
-//	        validator = state.validators[withdrawal.index]
-//	        has_sufficient_effective_balance = validator.effective_balance >= MIN_ACTIVATION_BALANCE
-//	        has_excess_balance = state.balances[withdrawal.index] > MIN_ACTIVATION_BALANCE
-//	        if validator.exit_epoch == FAR_FUTURE_EPOCH and has_sufficient_effective_balance and has_excess_balance:
-//	            withdrawable_balance = min(state.balances[withdrawal.index] - MIN_ACTIVATION_BALANCE, withdrawal.amount)
-//	            withdrawals.append(Withdrawal(
-//	                index=withdrawal_index,
-//	                validator_index=withdrawal.index,
-//	                address=ExecutionAddress(validator.withdrawal_credentials[12:]),
-//	                amount=withdrawable_balance,
-//	            ))
-//	            withdrawal_index += WithdrawalIndex(1)
+//			validator = state.validators[withdrawal.index]
+//			has_sufficient_effective_balance = validator.effective_balance >= MIN_ACTIVATION_BALANCE
+//			has_excess_balance = state.balances[withdrawal.index] > MIN_ACTIVATION_BALANCE
+//			if validator.exit_epoch == FAR_FUTURE_EPOCH and has_sufficient_effective_balance and has_excess_balance:
+//				withdrawable_balance = min(state.balances[withdrawal.index] - MIN_ACTIVATION_BALANCE, withdrawal.amount)
+//				withdrawals.append(Withdrawal(
+//					index=withdrawal_index,
+//					validator_index=withdrawal.index,
+//					address=ExecutionAddress(validator.withdrawal_credentials[12:]),
+//					amount=withdrawable_balance,
+//				))
+//				withdrawal_index += WithdrawalIndex(1)
 //
-//	    partial_withdrawals_count = len(withdrawals)
+//		processed_partial_withdrawals_count += 1
 //
-//	    # Sweep for remaining.
-//	    bound = min(len(state.validators), MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP)
-//	    for _ in range(bound):
-//	        validator = state.validators[validator_index]
-//	        balance = state.balances[validator_index]
-//	        if is_fully_withdrawable_validator(validator, balance, epoch):
-//	            withdrawals.append(Withdrawal(
-//	                index=withdrawal_index,
-//	                validator_index=validator_index,
-//	                address=ExecutionAddress(validator.withdrawal_credentials[12:]),
-//	                amount=balance,
-//	            ))
-//	            withdrawal_index += WithdrawalIndex(1)
-//	        elif is_partially_withdrawable_validator(validator, balance):
-//	            withdrawals.append(Withdrawal(
-//	                index=withdrawal_index,
-//	                validator_index=validator_index,
-//	                address=ExecutionAddress(validator.withdrawal_credentials[12:]),
-//	                amount=balance - get_validator_max_effective_balance(validator),  # [Modified in Electra:EIP7251]
-//	            ))
-//	            withdrawal_index += WithdrawalIndex(1)
-//	        if len(withdrawals) == MAX_WITHDRAWALS_PER_PAYLOAD:
-//	            break
-//	        validator_index = ValidatorIndex((validator_index + 1) % len(state.validators))
-//	    return withdrawals, partial_withdrawals_count
+//		# Sweep for remaining.
+//		bound = min(len(state.validators), MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP)
+//		for _ in range(bound):
+//			validator = state.validators[validator_index]
+//			# [Modified in Electra:EIP7251]
+//			partially_withdrawn_balance = sum(withdrawal.amount for withdrawal in withdrawals if withdrawal.validator_index == validator_index)
+//			balance = state.balances[validator_index] - partially_withdrawn_balance
+//			if is_fully_withdrawable_validator(validator, balance, epoch):
+//				withdrawals.append(Withdrawal(
+//					index=withdrawal_index,
+//					validator_index=validator_index,
+//					address=ExecutionAddress(validator.withdrawal_credentials[12:]),
+//					amount=balance,
+//				))
+//				withdrawal_index += WithdrawalIndex(1)
+//			elif is_partially_withdrawable_validator(validator, balance):
+//				withdrawals.append(Withdrawal(
+//					index=withdrawal_index,
+//					validator_index=validator_index,
+//					address=ExecutionAddress(validator.withdrawal_credentials[12:]),
+//					amount=balance - get_validator_max_effective_balance(validator),  # [Modified in Electra:EIP7251]
+//				))
+//				withdrawal_index += WithdrawalIndex(1)
+//			if len(withdrawals) == MAX_WITHDRAWALS_PER_PAYLOAD:
+//				break
+//			validator_index = ValidatorIndex((validator_index + 1) % len(state.validators))
+//		return withdrawals, processed_partial_withdrawals_count
 func (b *BeaconState) ExpectedWithdrawals() ([]*enginev1.Withdrawal, uint64, error) {
 	if b.version < version.Capella {
 		return nil, 0, errNotSupported("ExpectedWithdrawals", b.version)
@@ -113,7 +116,7 @@ func (b *BeaconState) ExpectedWithdrawals() ([]*enginev1.Withdrawal, uint64, err
 	epoch := slots.ToEpoch(b.slot)
 
 	// Electra partial withdrawals functionality.
-	var partialWithdrawalsCount uint64
+	var processedPartialWithdrawalsCount uint64
 	if b.version >= version.Electra {
 		for _, w := range b.pendingPartialWithdrawals {
 			if w.WithdrawableEpoch > epoch || len(withdrawals) >= int(params.BeaconConfig().MaxPendingPartialsPerWithdrawalsSweep) {
@@ -140,7 +143,7 @@ func (b *BeaconState) ExpectedWithdrawals() ([]*enginev1.Withdrawal, uint64, err
 				})
 				withdrawalIndex++
 			}
-			partialWithdrawalsCount++
+			processedPartialWithdrawalsCount++
 		}
 	}
 
@@ -155,6 +158,15 @@ func (b *BeaconState) ExpectedWithdrawals() ([]*enginev1.Withdrawal, uint64, err
 		if err != nil {
 			return nil, 0, errors.Wrapf(err, "could not retrieve balance at index %d", validatorIndex)
 		}
+		if b.version >= version.Electra {
+			var partiallyWithdrawnBalance uint64
+			for _, w := range withdrawals {
+				if w.ValidatorIndex == validatorIndex {
+					partiallyWithdrawnBalance += w.Amount
+				}
+			}
+			balance = balance - partiallyWithdrawnBalance
+		}
 		if helpers.IsFullyWithdrawableValidator(val, balance, epoch, b.version) {
 			withdrawals = append(withdrawals, &enginev1.Withdrawal{
 				Index:          withdrawalIndex,
@@ -181,7 +193,7 @@ func (b *BeaconState) ExpectedWithdrawals() ([]*enginev1.Withdrawal, uint64, err
 		}
 	}
 
-	return withdrawals, partialWithdrawalsCount, nil
+	return withdrawals, processedPartialWithdrawalsCount, nil
 }
 
 func (b *BeaconState) PendingPartialWithdrawals() ([]*ethpb.PendingPartialWithdrawal, error) {

diff --git a/beacon-chain/state/state-native/getters_withdrawal_test.go b/beacon-chain/state/state-native/getters_withdrawal_test.go
@@ -367,4 +367,52 @@ func TestExpectedWithdrawals(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, uint64(10), partialWithdrawalsCount)
 	})
+	t.Run("electra same validator has one partially and one fully withdrawable", func(t *testing.T) {
+		s, _ := util.DeterministicGenesisStateElectra(t, 1)
+		vals := make([]*ethpb.Validator, 100)
+		balances := make([]uint64, 100)
+		for i := range vals {
+			balances[i] = params.BeaconConfig().MaxEffectiveBalance
+			val := &ethpb.Validator{
+				WithdrawalCredentials: make([]byte, 32),
+				EffectiveBalance:      params.BeaconConfig().MaxEffectiveBalance,
+				WithdrawableEpoch:     primitives.Epoch(1),
+				ExitEpoch:             params.BeaconConfig().FarFutureEpoch,
+			}
+			val.WithdrawalCredentials[0] = params.BeaconConfig().ETH1AddressWithdrawalPrefixByte
+			val.WithdrawalCredentials[31] = byte(i)
+			vals[i] = val
+		}
+		balances[1] += params.BeaconConfig().MinDepositAmount
+		vals[1].WithdrawableEpoch = primitives.Epoch(0)
+		require.NoError(t, s.SetValidators(vals))
+		require.NoError(t, s.SetBalances(balances))
+		// Give validator a pending balance to withdraw.
+		require.NoError(t, s.AppendPendingPartialWithdrawal(&ethpb.PendingPartialWithdrawal{
+			Index:             1,
+			Amount:            balances[1], // will only deduct excess even though balance is more that minimum activation
+			WithdrawableEpoch: primitives.Epoch(0),
+		}))
+		p, err := s.PendingPartialWithdrawals()
+		require.NoError(t, err)
+		require.Equal(t, 1, len(p))
+		expected, _, err := s.ExpectedWithdrawals()
+		require.NoError(t, err)
+		require.Equal(t, 2, len(expected))
+
+		withdrawalFull := &enginev1.Withdrawal{
+			Index:          1,
+			ValidatorIndex: 1,
+			Address:        vals[1].WithdrawalCredentials[12:],
+			Amount:         balances[1] - params.BeaconConfig().MinDepositAmount, // subtract the partial from this
+		}
+		withdrawalPartial := &enginev1.Withdrawal{
+			Index:          0,
+			ValidatorIndex: 1,
+			Address:        vals[1].WithdrawalCredentials[12:],
+			Amount:         params.BeaconConfig().MinDepositAmount,
+		}
+		require.DeepEqual(t, withdrawalPartial, expected[0])
+		require.DeepEqual(t, withdrawalFull, expected[1])
+	})
 }