From 4c5eaea4d0cc3d863048fcf3672e46401e1436be Mon Sep 17 00:00:00 2001 From: BZO95 Date: Thu, 5 Mar 2026 17:02:15 +0800 Subject: [PATCH] trie/bintrie: fix code chunk key encoding --- trie/bintrie/key_encoding.go | 4 +-- trie/bintrie/key_encoding_test.go | 59 +++++++++++++++++++++++++++++++ trie/bintrie/trie.go | 6 ++-- trie/bintrie/trie_test.go | 53 +++++++++++++++++++++++++++ 4 files changed, 117 insertions(+), 5 deletions(-) create mode 100644 trie/bintrie/key_encoding_test.go diff --git a/trie/bintrie/key_encoding.go b/trie/bintrie/key_encoding.go index 94a22d52d0..10a4743d2a 100644 --- a/trie/bintrie/key_encoding.go +++ b/trie/bintrie/key_encoding.go @@ -90,8 +90,8 @@ func GetBinaryTreeKeyStorageSlot(address common.Address, key []byte) []byte { } func GetBinaryTreeKeyCodeChunk(address common.Address, chunknr *uint256.Int) []byte { - chunkOffset := new(uint256.Int).Add(codeOffset, chunknr).Bytes() - return GetBinaryTreeKey(address, chunkOffset) + chunkOffset := new(uint256.Int).Add(codeOffset, chunknr).Bytes32() + return GetBinaryTreeKey(address, chunkOffset[:]) } func StorageIndex(storageKey []byte) (*uint256.Int, byte) { diff --git a/trie/bintrie/key_encoding_test.go b/trie/bintrie/key_encoding_test.go new file mode 100644 index 0000000000..497cb9aed3 --- /dev/null +++ b/trie/bintrie/key_encoding_test.go @@ -0,0 +1,59 @@ +// Copyright 2025 The go-ethereum Authors +// This file is part of the go-ethereum library. +// +// The go-ethereum library is free software: you can redistribute it and/or modify +// it under the terms of the GNU Lesser General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// The go-ethereum library is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Lesser General Public License for more details. +// +// You should have received a copy of the GNU Lesser General Public License +// along with the go-ethereum library. If not, see . + +package bintrie + +import ( + "bytes" + "encoding/binary" + "testing" + + "github.com/ethereum/go-ethereum/common" + "github.com/holiman/uint256" +) + +func TestGetBinaryTreeKeyCodeChunkBoundaries(t *testing.T) { + addr := common.HexToAddress("0x1234567890abcdef1234567890abcdef12345678") + + for _, chunknr := range []uint64{0, 1, 127, 128, 255, 256, 257, 1024} { + got := GetBinaryTreeKeyCodeChunk(addr, uint256.NewInt(chunknr)) + + var offset [HashSize]byte + binary.BigEndian.PutUint64(offset[24:], chunknr+128) + want := GetBinaryTreeKey(addr, offset[:]) + + if !bytes.Equal(got, want) { + t.Fatalf("wrong code chunk key for chunk=%d: got=%x want=%x", chunknr, got, want) + } + } +} + +func TestGetBinaryTreeKeyCodeChunkLargeIndex(t *testing.T) { + addr := common.HexToAddress("0x1234567890abcdef1234567890abcdef12345678") + + var chunknr uint256.Int + chunknr.SetBytes(common.FromHex("0x0102030405060708090a0b0c0d0e0f10")) + got := GetBinaryTreeKeyCodeChunk(addr, &chunknr) + + var offset uint256.Int + offset.Add(codeOffset, &chunknr) + offsetBytes := offset.Bytes32() + want := GetBinaryTreeKey(addr, offsetBytes[:]) + + if !bytes.Equal(got, want) { + t.Fatalf("wrong code chunk key for large chunk index: got=%x want=%x", got, want) + } +} diff --git a/trie/bintrie/trie.go b/trie/bintrie/trie.go index a509c471b8..861ed08a0a 100644 --- a/trie/bintrie/trie.go +++ b/trie/bintrie/trie.go @@ -383,9 +383,9 @@ func (t *BinaryTrie) UpdateContractCode(addr common.Address, codeHash common.Has groupOffset := (chunknr + 128) % StemNodeWidth if groupOffset == 0 /* start of new group */ || chunknr == 0 /* first chunk in header group */ { values = make([][]byte, StemNodeWidth) - var offset [HashSize]byte - binary.BigEndian.PutUint64(offset[24:], chunknr+128) - key = GetBinaryTreeKey(addr, offset[:]) + var chunknrInt uint256.Int + chunknrInt.SetUint64(chunknr) + key = GetBinaryTreeKeyCodeChunk(addr, &chunknrInt) } values[groupOffset] = chunks[i : i+HashSize] diff --git a/trie/bintrie/trie_test.go b/trie/bintrie/trie_test.go index 256fd218e2..3e878e9c58 100644 --- a/trie/bintrie/trie_test.go +++ b/trie/bintrie/trie_test.go @@ -267,6 +267,59 @@ func TestStorageRoundTrip(t *testing.T) { } } +// TestContractCodeRoundTrip verifies that UpdateContractCode stores the first +// code chunk under the key returned by GetBinaryTreeKeyCodeChunk. +func TestContractCodeRoundTrip(t *testing.T) { + tracer := trie.NewPrevalueTracer() + tr := &BinaryTrie{ + root: NewBinaryNode(), + tracer: tracer, + } + addr := common.HexToAddress("0x1234567890abcdef1234567890abcdef12345678") + + acc := &types.StateAccount{ + Nonce: 1, + Balance: uint256.NewInt(1000), + CodeHash: common.HexToHash("c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470").Bytes(), + } + if err := tr.UpdateAccount(addr, acc, 0); err != nil { + t.Fatalf("UpdateAccount error: %v", err) + } + + code := common.FromHex("0x6001600055") + if err := tr.UpdateContractCode(addr, common.BytesToHash(acc.CodeHash), code); err != nil { + t.Fatalf("UpdateContractCode error: %v", err) + } + + key := GetBinaryTreeKeyCodeChunk(addr, uint256.NewInt(0)) + var ( + got []byte + err error + ) + switch root := tr.root.(type) { + case *InternalNode: + got, err = root.Get(key, tr.nodeResolver) + if err != nil { + t.Fatalf("Get error: %v", err) + } + case *StemNode: + values, err := root.GetValuesAtStem(key[:StemSize], nil) + if err != nil { + t.Fatalf("GetValuesAtStem error: %v", err) + } + if values == nil { + t.Fatal("expected stem values for code chunk key") + } + got = values[key[StemSize]] + default: + t.Fatalf("unexpected root type %T", tr.root) + } + want := ChunkifyCode(code)[:HashSize] + if !bytes.Equal(got, want) { + t.Fatalf("wrong chunk for chunk #0 key: got=%x want=%x", got, want) + } +} + func TestBinaryTrieWitness(t *testing.T) { tracer := trie.NewPrevalueTracer()