From da476c1991ee5dabf3e8e116f390327928454d30 Mon Sep 17 00:00:00 2001 From: Matthieu Vachon Date: Mon, 9 Jun 2025 10:16:22 -0400 Subject: [PATCH] More StateDB hook --- consensus/bor/bor.go | 22 +++++++++++----------- consensus/bor/contract/client.go | 13 +++++++++---- consensus/bor/genesis.go | 9 +++++---- consensus/bor/genesis_contract_mock.go | 8 ++++---- consensus/bor/heimdall/span/spanner.go | 4 ++-- consensus/bor/span.go | 5 +++-- consensus/bor/span_mock.go | 4 ++-- consensus/bor/statefull/processor.go | 8 ++++---- core/state/statedb.go | 22 ++++++++++++++++++++-- core/state/statedb_hooked.go | 18 ++++++++++++++++++ core/vm/interface.go | 18 ++++++++++++++++++ 11 files changed, 96 insertions(+), 35 deletions(-) diff --git a/consensus/bor/bor.go b/consensus/bor/bor.go index c86bcf79ab..3bf57e03d9 100644 --- a/consensus/bor/bor.go +++ b/consensus/bor/bor.go @@ -881,7 +881,7 @@ func (c *Bor) Finalize(chain consensus.ChainHeaderReader, header *types.Header, // Extract the underlying state to access methods like `IntermediateRoot` and `Copy` // required for bor consensus operations - state := wrappedState.(*state.StateDB) + // state := wrappedState.(*state.StateDB) var ( stateSyncData []*types.StateSyncData @@ -892,24 +892,24 @@ func (c *Bor) Finalize(chain consensus.ChainHeaderReader, header *types.Header, start := time.Now() cx := statefull.ChainContext{Chain: chain, Bor: c} // check and commit span - if err := c.checkAndCommitSpan(state, header, cx, tracer); err != nil { + if err := c.checkAndCommitSpan(wrappedState, header, cx, tracer); err != nil { log.Error("Error while committing span", "error", err) return } if c.HeimdallClient != nil { // commit states - stateSyncData, err = c.CommitStates(state, header, cx, tracer) + stateSyncData, err = c.CommitStates(wrappedState, header, cx, tracer) if err != nil { log.Error("Error while committing states", "error", err) return } } - state.BorConsensusTime = time.Since(start) + wrappedState.SetBorConsensusTime(time.Since(start)) } - if err = c.changeContractCodeIfNeeded(headerNumber, state); err != nil { + if err = c.changeContractCodeIfNeeded(headerNumber, wrappedState); err != nil { log.Error("Error changing contract code", "error", err) return } @@ -934,7 +934,7 @@ func decodeGenesisAlloc(i interface{}) (types.GenesisAlloc, error) { return alloc, nil } -func (c *Bor) changeContractCodeIfNeeded(headerNumber uint64, state *state.StateDB) error { +func (c *Bor) changeContractCodeIfNeeded(headerNumber uint64, state vm.StateDB) error { for blockNumber, genesisAlloc := range c.config.BlockAlloc { if blockNumber == strconv.FormatUint(headerNumber, 10) { allocs, err := decodeGenesisAlloc(genesisAlloc) @@ -977,7 +977,7 @@ func (c *Bor) FinalizeAndAssemble(chain consensus.ChainHeaderReader, header *typ cx := statefull.ChainContext{Chain: chain, Bor: c} // check and commit span - if err = c.checkAndCommitSpan(state, header, cx, tracer); err != nil { + if err = c.checkAndCommitSpan(vm.StateDB(state), header, cx, tracer); err != nil { log.Error("Error while committing span", "error", err) return nil, err } @@ -1168,7 +1168,7 @@ func (c *Bor) Close() error { } func (c *Bor) checkAndCommitSpan( - state *state.StateDB, + state vm.StateDB, header *types.Header, chain core.ChainContext, tracer *tracing.Hooks, @@ -1217,7 +1217,7 @@ func (c *Bor) needToCommitSpan(currentSpan *span.Span, headerNumber uint64) bool func (c *Bor) FetchAndCommitSpan( ctx context.Context, newSpanID uint64, - state *state.StateDB, + state vm.StateDB, header *types.Header, chain core.ChainContext, tracer *tracing.Hooks, @@ -1255,7 +1255,7 @@ func (c *Bor) FetchAndCommitSpan( // CommitStates commit states func (c *Bor) CommitStates( - state *state.StateDB, + state vm.StateDB, header *types.Header, chain statefull.ChainContext, tracer *tracing.Hooks, @@ -1272,7 +1272,7 @@ func (c *Bor) CommitStates( if c.config.IsIndore(header.Number) { // Fetch the LastStateId from contract via current state instance - lastStateIDBig, err = c.GenesisContractsClient.LastStateId(state.Copy(), number-1, header.ParentHash) + lastStateIDBig, err = c.GenesisContractsClient.LastStateId(state.Clone().(vm.StateDB), number-1, header.ParentHash) if err != nil { return nil, err } diff --git a/consensus/bor/contract/client.go b/consensus/bor/contract/client.go index 42247f9614..4321b33e2b 100644 --- a/consensus/bor/contract/client.go +++ b/consensus/bor/contract/client.go @@ -2,11 +2,12 @@ package contract import ( "context" - "github.com/ethereum/go-ethereum/core/tracing" "math" "math/big" "strings" + "github.com/ethereum/go-ethereum/core/tracing" + "github.com/ethereum/go-ethereum/accounts/abi" "github.com/ethereum/go-ethereum/common" "github.com/ethereum/go-ethereum/common/hexutil" @@ -15,6 +16,7 @@ import ( "github.com/ethereum/go-ethereum/consensus/bor/statefull" "github.com/ethereum/go-ethereum/core/state" "github.com/ethereum/go-ethereum/core/types" + "github.com/ethereum/go-ethereum/core/vm" "github.com/ethereum/go-ethereum/internal/ethapi" "github.com/ethereum/go-ethereum/log" "github.com/ethereum/go-ethereum/params" @@ -67,7 +69,7 @@ func NewGenesisContractsClient( func (gc *GenesisContractsClient) CommitState( event *clerk.EventRecordWithTime, - state *state.StateDB, + state vm.StateDB, header *types.Header, chCtx statefull.ChainContext, tracer *tracing.Hooks, @@ -105,7 +107,7 @@ func (gc *GenesisContractsClient) CommitState( return gasUsed, nil } -func (gc *GenesisContractsClient) LastStateId(state *state.StateDB, number uint64, hash common.Hash) (*big.Int, error) { +func (gc *GenesisContractsClient) LastStateId(stateDB vm.StateDB, number uint64, hash common.Hash) (*big.Int, error) { blockNr := rpc.BlockNumber(number) const method = "lastStateId" @@ -121,13 +123,16 @@ func (gc *GenesisContractsClient) LastStateId(state *state.StateDB, number uint6 toAddress := common.HexToAddress(gc.StateReceiverContract) gas := (hexutil.Uint64)(uint64(math.MaxUint64 / 2)) + // The unhooked version always return the *state.StateDB inner object + original := stateDB.Unhooked().(*state.StateDB) + // BOR: Do a 'CallWithState' so that we can fetch the last state ID from a given (incoming) // state instead of local(canonical) chain's state. result, err := gc.ethAPI.CallWithState(context.Background(), ethapi.TransactionArgs{ Gas: &gas, To: &toAddress, Data: &msgData, - }, &rpc.BlockNumberOrHash{BlockNumber: &blockNr, BlockHash: &hash}, state, nil, nil) + }, &rpc.BlockNumberOrHash{BlockNumber: &blockNr, BlockHash: &hash}, original, nil, nil) if err != nil { return nil, err } diff --git a/consensus/bor/genesis.go b/consensus/bor/genesis.go index 7de4f4ab40..3c4c2ecb3b 100644 --- a/consensus/bor/genesis.go +++ b/consensus/bor/genesis.go @@ -1,18 +1,19 @@ package bor import ( - "github.com/ethereum/go-ethereum/core/tracing" "math/big" + "github.com/ethereum/go-ethereum/core/tracing" + "github.com/ethereum/go-ethereum/core/vm" + "github.com/ethereum/go-ethereum/common" "github.com/ethereum/go-ethereum/consensus/bor/clerk" "github.com/ethereum/go-ethereum/consensus/bor/statefull" - "github.com/ethereum/go-ethereum/core/state" "github.com/ethereum/go-ethereum/core/types" ) //go:generate mockgen -destination=./genesis_contract_mock.go -package=bor . GenesisContract type GenesisContract interface { - CommitState(event *clerk.EventRecordWithTime, state *state.StateDB, header *types.Header, chCtx statefull.ChainContext, tracer *tracing.Hooks) (uint64, error) - LastStateId(state *state.StateDB, number uint64, hash common.Hash) (*big.Int, error) + CommitState(event *clerk.EventRecordWithTime, state vm.StateDB, header *types.Header, chCtx statefull.ChainContext, tracer *tracing.Hooks) (uint64, error) + LastStateId(state vm.StateDB, number uint64, hash common.Hash) (*big.Int, error) } diff --git a/consensus/bor/genesis_contract_mock.go b/consensus/bor/genesis_contract_mock.go index 785a4a1d7b..9e3fb270cd 100644 --- a/consensus/bor/genesis_contract_mock.go +++ b/consensus/bor/genesis_contract_mock.go @@ -11,9 +11,9 @@ import ( common "github.com/ethereum/go-ethereum/common" clerk "github.com/ethereum/go-ethereum/consensus/bor/clerk" statefull "github.com/ethereum/go-ethereum/consensus/bor/statefull" - state "github.com/ethereum/go-ethereum/core/state" - types "github.com/ethereum/go-ethereum/core/types" tracing "github.com/ethereum/go-ethereum/core/tracing" + types "github.com/ethereum/go-ethereum/core/types" + vm "github.com/ethereum/go-ethereum/core/vm" gomock "github.com/golang/mock/gomock" ) @@ -41,7 +41,7 @@ func (m *MockGenesisContract) EXPECT() *MockGenesisContractMockRecorder { } // CommitState mocks base method. -func (m *MockGenesisContract) CommitState(arg0 *clerk.EventRecordWithTime, arg1 *state.StateDB, arg2 *types.Header, arg3 statefull.ChainContext, arg4 *tracing.Hooks) (uint64, error) { +func (m *MockGenesisContract) CommitState(arg0 *clerk.EventRecordWithTime, arg1 vm.StateDB, arg2 *types.Header, arg3 statefull.ChainContext, arg4 *tracing.Hooks) (uint64, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "CommitState", arg0, arg1, arg2, arg3, arg4) ret0, _ := ret[0].(uint64) @@ -56,7 +56,7 @@ func (mr *MockGenesisContractMockRecorder) CommitState(arg0, arg1, arg2, arg3, a } // LastStateId mocks base method. -func (m *MockGenesisContract) LastStateId(arg0 *state.StateDB, arg1 uint64, arg2 common.Hash) (*big.Int, error) { +func (m *MockGenesisContract) LastStateId(arg0 vm.StateDB, arg1 uint64, arg2 common.Hash) (*big.Int, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "LastStateId", arg0, arg1, arg2) ret0, _ := ret[0].(*big.Int) diff --git a/consensus/bor/heimdall/span/spanner.go b/consensus/bor/heimdall/span/spanner.go index 09433fea1d..fc1d88323d 100644 --- a/consensus/bor/heimdall/span/spanner.go +++ b/consensus/bor/heimdall/span/spanner.go @@ -13,9 +13,9 @@ import ( "github.com/ethereum/go-ethereum/consensus/bor/statefull" "github.com/ethereum/go-ethereum/consensus/bor/valset" "github.com/ethereum/go-ethereum/core" - "github.com/ethereum/go-ethereum/core/state" "github.com/ethereum/go-ethereum/core/tracing" "github.com/ethereum/go-ethereum/core/types" + "github.com/ethereum/go-ethereum/core/vm" "github.com/ethereum/go-ethereum/internal/ethapi" "github.com/ethereum/go-ethereum/log" "github.com/ethereum/go-ethereum/params" @@ -288,7 +288,7 @@ func (c *ChainSpanner) GetCurrentValidatorsByHash(ctx context.Context, headerHas const method = "commitSpan" -func (c *ChainSpanner) CommitSpan(ctx context.Context, heimdallSpan HeimdallSpan, state *state.StateDB, header *types.Header, chainContext core.ChainContext, tracer *tracing.Hooks) error { +func (c *ChainSpanner) CommitSpan(ctx context.Context, heimdallSpan HeimdallSpan, state vm.StateDB, header *types.Header, chainContext core.ChainContext, tracer *tracing.Hooks) error { // get validators bytes validators := make([]valset.MinimalVal, 0, len(heimdallSpan.ValidatorSet.Validators)) for _, val := range heimdallSpan.ValidatorSet.Validators { diff --git a/consensus/bor/span.go b/consensus/bor/span.go index bff630b33a..9b69b36c06 100644 --- a/consensus/bor/span.go +++ b/consensus/bor/span.go @@ -2,13 +2,14 @@ package bor import ( "context" + "github.com/ethereum/go-ethereum/core/tracing" + "github.com/ethereum/go-ethereum/core/vm" "github.com/ethereum/go-ethereum/common" "github.com/ethereum/go-ethereum/consensus/bor/heimdall/span" "github.com/ethereum/go-ethereum/consensus/bor/valset" "github.com/ethereum/go-ethereum/core" - "github.com/ethereum/go-ethereum/core/state" "github.com/ethereum/go-ethereum/core/types" "github.com/ethereum/go-ethereum/rpc" ) @@ -18,5 +19,5 @@ type Spanner interface { GetCurrentSpan(ctx context.Context, headerHash common.Hash) (*span.Span, error) GetCurrentValidatorsByHash(ctx context.Context, headerHash common.Hash, blockNumber uint64) ([]*valset.Validator, error) GetCurrentValidatorsByBlockNrOrHash(ctx context.Context, blockNrOrHash rpc.BlockNumberOrHash, blockNumber uint64) ([]*valset.Validator, error) - CommitSpan(ctx context.Context, heimdallSpan span.HeimdallSpan, state *state.StateDB, header *types.Header, chainContext core.ChainContext, tracer *tracing.Hooks) error + CommitSpan(ctx context.Context, heimdallSpan span.HeimdallSpan, state vm.StateDB, header *types.Header, chainContext core.ChainContext, tracer *tracing.Hooks) error } diff --git a/consensus/bor/span_mock.go b/consensus/bor/span_mock.go index f93c815d00..ba03ff3e84 100644 --- a/consensus/bor/span_mock.go +++ b/consensus/bor/span_mock.go @@ -12,9 +12,9 @@ import ( span "github.com/ethereum/go-ethereum/consensus/bor/heimdall/span" valset "github.com/ethereum/go-ethereum/consensus/bor/valset" core "github.com/ethereum/go-ethereum/core" - state "github.com/ethereum/go-ethereum/core/state" tracing "github.com/ethereum/go-ethereum/core/tracing" types "github.com/ethereum/go-ethereum/core/types" + vm "github.com/ethereum/go-ethereum/core/vm" rpc "github.com/ethereum/go-ethereum/rpc" gomock "github.com/golang/mock/gomock" ) @@ -43,7 +43,7 @@ func (m *MockSpanner) EXPECT() *MockSpannerMockRecorder { } // CommitSpan mocks base method. -func (m *MockSpanner) CommitSpan(arg0 context.Context, arg1 span.HeimdallSpan, arg2 *state.StateDB, arg3 *types.Header, arg4 core.ChainContext, arg5 *tracing.Hooks) error { +func (m *MockSpanner) CommitSpan(arg0 context.Context, arg1 span.HeimdallSpan, arg2 vm.StateDB, arg3 *types.Header, arg4 core.ChainContext, arg5 *tracing.Hooks) error { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "CommitSpan", arg0, arg1, arg2, arg3, arg4, arg5) ret0, _ := ret[0].(error) diff --git a/consensus/bor/statefull/processor.go b/consensus/bor/statefull/processor.go index 36b19420fc..79d7e0bc39 100644 --- a/consensus/bor/statefull/processor.go +++ b/consensus/bor/statefull/processor.go @@ -3,16 +3,16 @@ package statefull import ( "bytes" "context" - "github.com/ethereum/go-ethereum/core/tracing" - "github.com/ethereum/go-ethereum/crypto" "math" "math/big" + "github.com/ethereum/go-ethereum/core/tracing" + "github.com/ethereum/go-ethereum/crypto" + ethereum "github.com/ethereum/go-ethereum" "github.com/ethereum/go-ethereum/common" "github.com/ethereum/go-ethereum/consensus" "github.com/ethereum/go-ethereum/core" - "github.com/ethereum/go-ethereum/core/state" "github.com/ethereum/go-ethereum/core/types" "github.com/ethereum/go-ethereum/core/vm" "github.com/ethereum/go-ethereum/log" @@ -71,7 +71,7 @@ func GetSystemMessage(toAddress common.Address, data []byte) Callmsg { func ApplyMessage( _ context.Context, msg Callmsg, - state *state.StateDB, + state vm.StateDB, header *types.Header, chainConfig *params.ChainConfig, chainContext core.ChainContext, diff --git a/core/state/statedb.go b/core/state/statedb.go index cfb1514453..4528fca9cd 100644 --- a/core/state/statedb.go +++ b/core/state/statedb.go @@ -808,13 +808,16 @@ func (s *StateDB) SubBalance(addr common.Address, amount *uint256.Int, reason tr return stateObject.SetBalance(new(uint256.Int).Sub(stateObject.Balance(), amount)) } -func (s *StateDB) SetBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason) { +func (s *StateDB) SetBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason) uint256.Int { stateObject := s.getOrNewStateObject(addr) + var prevBalance uint256.Int + if stateObject != nil { stateObject = s.mvRecordWritten(stateObject) - stateObject.SetBalance(amount) + prevBalance = stateObject.SetBalance(amount) MVWrite(s, blockstm.NewSubpathKey(addr, BalancePath)) } + return prevBalance } func (s *StateDB) SetNonce(addr common.Address, nonce uint64, reason tracing.NonceChangeReason) { @@ -1942,3 +1945,18 @@ func (s *StateDB) Witness() *stateless.Witness { func (s *StateDB) AccessEvents() *AccessEvents { return s.accessEvents } + +// Polygon specific + +func (s *StateDB) Clone() any { + return s.Copy() +} + +func (s *StateDB) Unhooked() any { + // Already unhooked, just return self + return s +} + +func (s *StateDB) SetBorConsensusTime(borConsensusTime time.Duration) { + s.BorConsensusTime = borConsensusTime +} diff --git a/core/state/statedb_hooked.go b/core/state/statedb_hooked.go index cc50bd09d8..d69e65d8f2 100644 --- a/core/state/statedb_hooked.go +++ b/core/state/statedb_hooked.go @@ -18,6 +18,7 @@ package state import ( "math/big" + "time" "github.com/ethereum/go-ethereum/common" "github.com/ethereum/go-ethereum/core/blockstm" @@ -184,6 +185,14 @@ func (s *hookedStateDB) AddBalance(addr common.Address, amount *uint256.Int, rea return prev } +func (s *hookedStateDB) SetBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason) uint256.Int { + prev := s.inner.SetBalance(addr, amount, reason) + if s.hooks.OnBalanceChange != nil { + s.hooks.OnBalanceChange(addr, prev.ToBig(), amount.ToBig(), reason) + } + return prev +} + func (s *hookedStateDB) SetNonce(address common.Address, nonce uint64, reason tracing.NonceChangeReason) { prev := s.inner.GetNonce(address) s.inner.SetNonce(address, nonce, reason) @@ -304,3 +313,12 @@ func (s *hookedStateDB) TxIndex() int { func (s *hookedStateDB) SetTxContext(txHash common.Hash, txIndex int) { s.inner.SetTxContext(txHash, txIndex) } +func (s *hookedStateDB) Clone() any { + return NewHookedState(s.inner.Copy(), s.hooks) +} +func (s *hookedStateDB) Unhooked() any { + return s.inner +} +func (s *hookedStateDB) SetBorConsensusTime(borConsensusTime time.Duration) { + s.inner.SetBorConsensusTime(borConsensusTime) +} diff --git a/core/vm/interface.go b/core/vm/interface.go index 1921ead5f6..a9776677b5 100644 --- a/core/vm/interface.go +++ b/core/vm/interface.go @@ -17,6 +17,8 @@ package vm import ( + "time" + "github.com/ethereum/go-ethereum/common" "github.com/ethereum/go-ethereum/core/blockstm" "github.com/ethereum/go-ethereum/core/state" @@ -111,4 +113,20 @@ type StateDB interface { GetLogs(txHash common.Hash, blockNumber uint64, blockHash common.Hash) []*types.Log TxIndex() int SetTxContext(txHash common.Hash, txIndex int) + SetBalance(common.Address, *uint256.Int, tracing.BalanceChangeReason) uint256.Int + // Clone is used to create a copy of the StateDB, same as `Copy` on *state.StateDB but rename so interface has its own name + // + // state.Clone().(vm.StateDB) + // + // The `any` return type is required to avoid import cycles. + Clone() any + // Unhooked is used to return the underlying state without any hooks applied, in Polygon, some potential + // state modifying operations can be called on a vm.StateDB interface, which might be hooked but we want those + // operation to always be non-recorded, this method ensures this. + // + // state.Unhooked().(vm.StateDB) + // + // The `any` return type is required to avoid import cycles. + Unhooked() any + SetBorConsensusTime(borConsensusTime time.Duration) }