eip7251: Bugfix and more withdrawal tests

CHANGELOG.md

@@ -43,6 +43,7 @@ The format is based on Keep a Changelog, and this project adheres to Semantic Ve
 - Fixed mesh size by appending `gParams.Dhi = gossipSubDhi`
 - Fix skipping partial withdrawals count.
 - wait for the async StreamEvent writer to exit before leaving the http handler, avoiding race condition panics [pr](https://github.com/prysmaticlabs/prysm/pull/14557)
+- EIP7521 - fixing withdrawal bug
 
 ### Security
 

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:
+//		    # [Modified in Electra:EIP7251]
+//	   	expected_withdrawals, processed_partial_withdrawals_count = get_expected_withdrawals(state)
 //
-//	    assert len(payload.withdrawals) == len(expected_withdrawals)
+//		    assert len(payload.withdrawals) == len(expected_withdrawals)
 //
-//	    for expected_withdrawal, withdrawal in zip(expected_withdrawals, payload.withdrawals):
-//	        assert withdrawal == expected_withdrawal
-//	        decrease_balance(state, withdrawal.validator_index, withdrawal.amount)
+//		    for expected_withdrawal, withdrawal in zip(expected_withdrawals, payload.withdrawals):
+//		        assert withdrawal == expected_withdrawal
+//		        decrease_balance(state, withdrawal.validator_index, withdrawal.amount)
 //
-//	    # Update pending partial withdrawals [New in Electra:EIP7251]
-//	    state.pending_partial_withdrawals = state.pending_partial_withdrawals[partial_withdrawals_count:]
+//		    # Update pending partial withdrawals [New in Electra:EIP7251]
+//		    state.pending_partial_withdrawals = state.pending_partial_withdrawals[processed_partial_withdrawals_count:]
 //
-//	    # 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 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
+//		    # 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)
 		}
 	}

beacon-chain/state/state-native/getters_withdrawal.go

@@ -48,57 +48,60 @@ 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] = []
+//		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] = []
+//			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) {

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])
+	})
 }