core: improve withdrawal index assignment in GenerateChain (#26756)

This fixes an issue where the withdrawal index was not calculated correctly
for multiple withdrawals in a single block.
Co-authored-by: Gary Rong <garyrong0905@gmail.com>
Co-authored-by: Felix Lange <fjl@twurst.com>

core/chain_makers.go

@@ -207,23 +207,31 @@ func (b *BlockGen) AddUncle(h *types.Header) {
 }
 
 // AddWithdrawal adds a withdrawal to the generated block.
-func (b *BlockGen) AddWithdrawal(w *types.Withdrawal) {
-	// The withdrawal will be assigned the next valid index.
-	var idx uint64
+// It returns the withdrawal index.
+func (b *BlockGen) AddWithdrawal(w *types.Withdrawal) uint64 {
+	cpy := *w
+	cpy.Index = b.nextWithdrawalIndex()
+	b.withdrawals = append(b.withdrawals, &cpy)
+	return cpy.Index
+}
+
+// nextWithdrawalIndex computes the index of the next withdrawal.
+func (b *BlockGen) nextWithdrawalIndex() uint64 {
+	if len(b.withdrawals) != 0 {
+		return b.withdrawals[len(b.withdrawals)-1].Index + 1
+	}
 	for i := b.i - 1; i >= 0; i-- {
 		if wd := b.chain[i].Withdrawals(); len(wd) != 0 {
-			idx = wd[len(wd)-1].Index + 1
-			break
+			return wd[len(wd)-1].Index + 1
 		}
 		if i == 0 {
-			// Correctly set the index if no parent had withdrawals
+			// Correctly set the index if no parent had withdrawals.
 			if wd := b.parent.Withdrawals(); len(wd) != 0 {
-				idx = wd[len(wd)-1].Index + 1
+				return wd[len(wd)-1].Index + 1
 			}
 		}
 	}
-	w.Index = idx
-	b.withdrawals = append(b.withdrawals, w)
+	return 0
 }
 
 // PrevBlock returns a previously generated block by number. It panics if

core/chain_makers_test.go

@@ -19,7 +19,10 @@ package core
 import (
 	"fmt"
 	"math/big"
+	"testing"
 
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/consensus/beacon"
 	"github.com/ethereum/go-ethereum/consensus/ethash"
 	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/core/types"
@@ -28,6 +31,112 @@ import (
 	"github.com/ethereum/go-ethereum/params"
 )
 
+func TestGenerateWithdrawalChain(t *testing.T) {
+	var (
+		keyHex  = "9c647b8b7c4e7c3490668fb6c11473619db80c93704c70893d3813af4090c39c"
+		key, _  = crypto.HexToECDSA(keyHex)
+		address = crypto.PubkeyToAddress(key.PublicKey) // 658bdf435d810c91414ec09147daa6db62406379
+		aa      = common.Address{0xaa}
+		bb      = common.Address{0xbb}
+		funds   = big.NewInt(0).Mul(big.NewInt(1337), big.NewInt(params.Ether))
+		config  = *params.AllEthashProtocolChanges
+		gspec   = &Genesis{
+			Config:     &config,
+			Alloc:      GenesisAlloc{address: {Balance: funds}},
+			BaseFee:    big.NewInt(params.InitialBaseFee),
+			Difficulty: common.Big1,
+			GasLimit:   5_000_000,
+		}
+		gendb  = rawdb.NewMemoryDatabase()
+		signer = types.LatestSigner(gspec.Config)
+		db     = rawdb.NewMemoryDatabase()
+	)
+
+	config.TerminalTotalDifficultyPassed = true
+	config.TerminalTotalDifficulty = common.Big0
+	config.ShanghaiTime = u64(0)
+
+	// init 0xaa with some storage elements
+	storage := make(map[common.Hash]common.Hash)
+	storage[common.Hash{0x00}] = common.Hash{0x00}
+	storage[common.Hash{0x01}] = common.Hash{0x01}
+	storage[common.Hash{0x02}] = common.Hash{0x02}
+	storage[common.Hash{0x03}] = common.HexToHash("0303")
+	gspec.Alloc[aa] = GenesisAccount{
+		Balance: common.Big1,
+		Nonce:   1,
+		Storage: storage,
+		Code:    common.Hex2Bytes("6042"),
+	}
+	gspec.Alloc[bb] = GenesisAccount{
+		Balance: common.Big2,
+		Nonce:   1,
+		Storage: storage,
+		Code:    common.Hex2Bytes("600154600354"),
+	}
+
+	genesis := gspec.MustCommit(gendb)
+
+	chain, _ := GenerateChain(gspec.Config, genesis, beacon.NewFaker(), gendb, 4, func(i int, gen *BlockGen) {
+		tx, _ := types.SignTx(types.NewTransaction(gen.TxNonce(address), address, big.NewInt(1000), params.TxGas, new(big.Int).Add(gen.BaseFee(), common.Big1), nil), signer, key)
+		gen.AddTx(tx)
+		if i == 1 {
+			gen.AddWithdrawal(&types.Withdrawal{
+				Validator: 42,
+				Address:   common.Address{0xee},
+				Amount:    1337,
+			})
+			gen.AddWithdrawal(&types.Withdrawal{
+				Validator: 13,
+				Address:   common.Address{0xee},
+				Amount:    1,
+			})
+		}
+		if i == 3 {
+			gen.AddWithdrawal(&types.Withdrawal{
+				Validator: 42,
+				Address:   common.Address{0xee},
+				Amount:    1337,
+			})
+			gen.AddWithdrawal(&types.Withdrawal{
+				Validator: 13,
+				Address:   common.Address{0xee},
+				Amount:    1,
+			})
+		}
+	})
+
+	// Import the chain. This runs all block validation rules.
+	blockchain, _ := NewBlockChain(db, nil, gspec, nil, beacon.NewFaker(), vm.Config{}, nil, nil)
+	defer blockchain.Stop()
+
+	if i, err := blockchain.InsertChain(chain); err != nil {
+		fmt.Printf("insert error (block %d): %v\n", chain[i].NumberU64(), err)
+		return
+	}
+
+	// enforce that withdrawal indexes are monotonically increasing from 0
+	var (
+		withdrawalIndex uint64
+		head            = blockchain.CurrentBlock().NumberU64()
+	)
+	for i := 0; i < int(head); i++ {
+		block := blockchain.GetBlockByNumber(uint64(i))
+		if block == nil {
+			t.Fatalf("block %d not found", i)
+		}
+		if len(block.Withdrawals()) == 0 {
+			continue
+		}
+		for j := 0; j < len(block.Withdrawals()); j++ {
+			if block.Withdrawals()[j].Index != withdrawalIndex {
+				t.Fatalf("withdrawal index %d does not equal expected index %d", block.Withdrawals()[j].Index, withdrawalIndex)
+			}
+			withdrawalIndex += 1
+		}
+	}
+}
+
 func ExampleGenerateChain() {
 	var (
 		key1, _ = crypto.HexToECDSA("b71c71a67e1177ad4e901695e1b4b9ee17ae16c6668d313eac2f96dbcda3f291")