From 295a2d92eb63b49ecf30e017f9b1cc4058b6cae2 Mon Sep 17 00:00:00 2001 From: Unique-Divine Date: Fri, 25 Oct 2024 03:06:40 -0500 Subject: [PATCH] statedb: add cacheing for multistore before precompile runs --- x/evm/keeper/erc20.go | 6 +-- x/evm/keeper/precompiles.go | 4 ++ x/evm/precompile/funtoken.go | 3 +- x/evm/precompile/precompile.go | 38 ++++++++++--- x/evm/precompile/test/export.go | 3 +- x/evm/precompile/wasm.go | 4 +- x/evm/precompile/wasm_test.go | 4 +- x/evm/statedb/interfaces.go | 2 + x/evm/statedb/journal.go | 41 ++++++++++++++ x/evm/statedb/journal_test.go | 61 +++++++++++++-------- x/evm/statedb/state_object.go | 5 +- x/evm/statedb/statedb.go | 94 +++++++++++++++++++++++++++++++-- 12 files changed, 219 insertions(+), 46 deletions(-) diff --git a/x/evm/keeper/erc20.go b/x/evm/keeper/erc20.go index 10404bea4..328452ecf 100644 --- a/x/evm/keeper/erc20.go +++ b/x/evm/keeper/erc20.go @@ -218,11 +218,8 @@ func (k Keeper) CallContractWithInput( // sent by a user txConfig := k.TxConfig(ctx, gethcommon.BigToHash(big.NewInt(0))) - // Using tmp context to not modify the state in case of evm revert - tmpCtx, commitCtx := ctx.CacheContext() - evmResp, evmObj, err = k.ApplyEvmMsg( - tmpCtx, evmMsg, evm.NewNoOpTracer(), commit, evmCfg, txConfig, + ctx, evmMsg, evm.NewNoOpTracer(), commit, evmCfg, txConfig, ) if err != nil { // We don't know the actual gas used, so consuming the gas limit @@ -245,7 +242,6 @@ func (k Keeper) CallContractWithInput( } else { // Success, committing the state to ctx if commit { - commitCtx() totalGasUsed, err := k.AddToBlockGasUsed(ctx, evmResp.GasUsed) if err != nil { k.ResetGasMeterAndConsumeGas(ctx, ctx.GasMeter().Limit()) diff --git a/x/evm/keeper/precompiles.go b/x/evm/keeper/precompiles.go index 965866660..3b1a0e480 100644 --- a/x/evm/keeper/precompiles.go +++ b/x/evm/keeper/precompiles.go @@ -21,3 +21,7 @@ func (k *Keeper) AddPrecompiles( } } } + +func (k *Keeper) IsPrecompile(addr gethcommon.Address) bool { + return k.precompiles.Has(addr) +} diff --git a/x/evm/precompile/funtoken.go b/x/evm/precompile/funtoken.go index 5c585c2e9..b95fd0084 100644 --- a/x/evm/precompile/funtoken.go +++ b/x/evm/precompile/funtoken.go @@ -72,8 +72,7 @@ func (p precompileFunToken) Run( if err != nil { return nil, err } - // Dirty journal entries in `StateDB` must be committed - return bz, start.StateDB.Commit() + return bz, OnRunEnd(start.StateDB, start.SnapshotBeforeRun, p.Address()) } func PrecompileFunToken(keepers keepers.PublicKeepers) vm.PrecompiledContract { diff --git a/x/evm/precompile/precompile.go b/x/evm/precompile/precompile.go index a6bbfefc4..428af9775 100644 --- a/x/evm/precompile/precompile.go +++ b/x/evm/precompile/precompile.go @@ -147,6 +147,11 @@ type OnRunStartResult struct { Method *gethabi.Method StateDB *statedb.StateDB + + // SnapshotBeforeRun captures the state before precompile execution to enable + // proper state reversal if the call fails or if [statedb.JournalChange] + // is reverted in general. + SnapshotBeforeRun statedb.PrecompileSnapshotBeforeRun } // OnRunStart prepares the execution environment for a precompiled contract call. @@ -188,19 +193,40 @@ func OnRunStart( err = fmt.Errorf("failed to load the sdk.Context from the EVM StateDB") return } - ctx := stateDB.GetContext() - if err = stateDB.Commit(); err != nil { + cacheCtx, snapshot := stateDB.CacheCtxForPrecompile(contract.Address()) + if err = stateDB.CommitCacheCtx(); err != nil { return res, fmt.Errorf("error committing dirty journal entries: %w", err) } return OnRunStartResult{ - Args: args, - Ctx: ctx, - Method: method, - StateDB: stateDB, + Args: args, + Ctx: cacheCtx, + Method: method, + StateDB: stateDB, + SnapshotBeforeRun: snapshot, }, nil } +// OnRunEnd finalizes a precompile execution by saving its state snapshot to the +// journal. This ensures that any state changes can be properly reverted if needed. +// +// Args: +// - stateDB: The EVM state database +// - snapshot: The state snapshot taken before the precompile executed +// - precompileAddr: The address of the precompiled contract +// +// The snapshot is critical for maintaining state consistency when: +// - The operation gets reverted ([statedb.JournalChange] Revert). +// - The precompile modifies state in other modules (e.g., bank, wasm) +// - Multiple precompiles are called within a single transaction +func OnRunEnd( + stateDB *statedb.StateDB, + snapshot statedb.PrecompileSnapshotBeforeRun, + precompileAddr gethcommon.Address, +) error { + return stateDB.SavePrecompileSnapshotToJournal(precompileAddr, snapshot) +} + var precompileMethodIsTxMap map[PrecompileMethod]bool = map[PrecompileMethod]bool{ WasmMethod_execute: true, WasmMethod_instantiate: true, diff --git a/x/evm/precompile/test/export.go b/x/evm/precompile/test/export.go index 966dd3359..05405c730 100644 --- a/x/evm/precompile/test/export.go +++ b/x/evm/precompile/test/export.go @@ -267,6 +267,7 @@ func IncrementWasmCounterWithExecuteMulti( deps *evmtest.TestDeps, wasmContract sdk.AccAddress, times uint, + finalizeTx bool, ) (evmObj *vm.EVM) { msgArgsBz := []byte(` { @@ -308,7 +309,7 @@ func IncrementWasmCounterWithExecuteMulti( s.Require().NoError(err) ethTxResp, evmObj, err := deps.EvmKeeper.CallContractWithInput( - deps.Ctx, deps.Sender.EthAddr, &precompile.PrecompileAddr_Wasm, true, input, + deps.Ctx, deps.Sender.EthAddr, &precompile.PrecompileAddr_Wasm, finalizeTx, input, ) s.Require().NoError(err) s.Require().NotEmpty(ethTxResp.Ret) diff --git a/x/evm/precompile/wasm.go b/x/evm/precompile/wasm.go index 10817c673..49fe40f87 100644 --- a/x/evm/precompile/wasm.go +++ b/x/evm/precompile/wasm.go @@ -62,9 +62,7 @@ func (p precompileWasm) Run( if err != nil { return nil, err } - - // Dirty journal entries in `StateDB` must be committed - return bz, start.StateDB.Commit() + return bz, OnRunEnd(start.StateDB, start.SnapshotBeforeRun, p.Address()) } type precompileWasm struct { diff --git a/x/evm/precompile/wasm_test.go b/x/evm/precompile/wasm_test.go index d796f8b89..dbe35b839 100644 --- a/x/evm/precompile/wasm_test.go +++ b/x/evm/precompile/wasm_test.go @@ -100,13 +100,13 @@ func (s *WasmSuite) TestExecuteMultiHappy() { test.AssertWasmCounterState(&s.Suite, deps, wasmContract, 0) // count += 2 test.IncrementWasmCounterWithExecuteMulti( - &s.Suite, &deps, wasmContract, 2) + &s.Suite, &deps, wasmContract, 2, true) // count = 2 test.AssertWasmCounterState(&s.Suite, deps, wasmContract, 2) s.assertWasmCounterStateRaw(deps, wasmContract, 2) // count += 67 test.IncrementWasmCounterWithExecuteMulti( - &s.Suite, &deps, wasmContract, 67) + &s.Suite, &deps, wasmContract, 67, true) // count = 69 test.AssertWasmCounterState(&s.Suite, deps, wasmContract, 69) s.assertWasmCounterStateRaw(deps, wasmContract, 69) diff --git a/x/evm/statedb/interfaces.go b/x/evm/statedb/interfaces.go index a4c1c3b59..a9a0f37b4 100644 --- a/x/evm/statedb/interfaces.go +++ b/x/evm/statedb/interfaces.go @@ -38,4 +38,6 @@ type Keeper interface { // DeleteAccount handles contract's suicide call, clearing the balance, // contract bytecode, contract state, and its native account. DeleteAccount(ctx sdk.Context, addr common.Address) error + + IsPrecompile(addr common.Address) bool } diff --git a/x/evm/statedb/journal.go b/x/evm/statedb/journal.go index ac041b617..40a8cc3af 100644 --- a/x/evm/statedb/journal.go +++ b/x/evm/statedb/journal.go @@ -21,6 +21,8 @@ import ( "math/big" "sort" + store "github.com/cosmos/cosmos-sdk/store/types" + sdk "github.com/cosmos/cosmos-sdk/types" "github.com/ethereum/go-ethereum/common" ) @@ -335,3 +337,42 @@ func (ch accessListAddSlotChange) Revert(s *StateDB) { func (ch accessListAddSlotChange) Dirtied() *common.Address { return nil } + +// ------------------------------------------------------ +// PrecompileSnapshotBeforeRun + +// PrecompileSnapshotBeforeRun: Precompiles can alter persistent storage of other +// modules. These changes to persistent storage are not reverted by a `Revert` of +// [JournalChange] by default, as it generally manages only changes to accounts +// and Bank balances for ether (NIBI). +// +// As a workaround to make state changes from precompiles reversible, we store +// [PrecompileSnapshotBeforeRun] snapshots that sync and record the prior state +// of the other modules, allowing precompile calls to truly be reverted. +// +// As a simple example, suppose that a transaction calls a precompile. +// 1. If the precompile changes the state in the Bank Module or Wasm module +// 2. The call gets reverted (`revert()` in Solidity), which shoud restore the +// state to a in-memory snapshot recorded on the StateDB journal. +// 3. This could cause a problem where changes to the rest of the blockchain state +// are still in effect following the reversion in the EVM state DB. +type PrecompileSnapshotBeforeRun struct { + MultiStore store.CacheMultiStore + Events sdk.Events + Precompile common.Address +} + +var _ JournalChange = PrecompileSnapshotBeforeRun{} + +func (ch PrecompileSnapshotBeforeRun) Revert(s *StateDB) { + s.cacheCtx = s.cacheCtx.WithMultiStore(ch.MultiStore) + // Rewrite the `writeCacheCtxFn` using the same logic as sdk.Context.CacheCtx + s.writeToCommitCtxFromCacheCtx = func() { + s.ctx.EventManager().EmitEvents(ch.Events) + ch.MultiStore.Write() + } +} + +func (ch PrecompileSnapshotBeforeRun) Dirtied() *common.Address { + return &ch.Precompile +} diff --git a/x/evm/statedb/journal_test.go b/x/evm/statedb/journal_test.go index 5863face5..0297e8999 100644 --- a/x/evm/statedb/journal_test.go +++ b/x/evm/statedb/journal_test.go @@ -18,7 +18,7 @@ import ( "github.com/NibiruChain/nibiru/v2/x/evm/statedb" ) -func (s *Suite) TestPrecompileSnapshots() { +func (s *Suite) TestComplexJournalChanges() { deps := evmtest.NewTestDeps() bankDenom := evm.EVMBankDenom s.Require().NoError(testapp.FundAccount( @@ -32,25 +32,11 @@ func (s *Suite) TestPrecompileSnapshots() { wasmContract := test.SetupWasmContracts(&deps, &s.Suite)[1] fmt.Printf("wasmContract: %s\n", wasmContract) - assertionsBeforeRun := func(deps *evmtest.TestDeps) { - test.AssertWasmCounterState( - &s.Suite, *deps, wasmContract, 0, - ) - } - run := func(deps *evmtest.TestDeps) *vm.EVM { - return test.IncrementWasmCounterWithExecuteMulti( - &s.Suite, deps, wasmContract, 7, - ) - } - assertionsAfterRun := func(deps *evmtest.TestDeps) { - test.AssertWasmCounterState( - &s.Suite, *deps, wasmContract, 7, - ) - } s.T().Log("Assert before transition") - - assertionsBeforeRun(&deps) + test.AssertWasmCounterState( + &s.Suite, deps, wasmContract, 0, + ) deployArgs := []any{"name", "SYMBOL", uint8(18)} deployResp, err := evmtest.DeployContract( @@ -136,15 +122,48 @@ func (s *Suite) TestPrecompileSnapshots() { s.Require().ErrorContains(err, vm.ErrExecutionReverted.Error()) }) - s.Run("Precompile calls also start and end clean (no dirty changes)", func() { - evmObj = run(&deps) - assertionsAfterRun(&deps) + s.Run("Precompile calls populate snapshots", func() { + s.T().Log("commitEvmTx=true, expect 0 dirty journal entries") + commitEvmTx := true + evmObj = test.IncrementWasmCounterWithExecuteMulti( + &s.Suite, &deps, wasmContract, 7, commitEvmTx, + ) + // assertions after run + test.AssertWasmCounterState( + &s.Suite, deps, wasmContract, 7, + ) stateDB, ok := evmObj.StateDB.(*statedb.StateDB) s.Require().True(ok, "error retrieving StateDB from the EVM") if stateDB.DirtiesCount() != 0 { debugDirtiesCountMismatch(stateDB, s.T()) s.FailNow("expected 0 dirty journal changes") } + + s.T().Log("commitEvmTx=false, expect dirty journal entries") + commitEvmTx = false + evmObj = test.IncrementWasmCounterWithExecuteMulti( + &s.Suite, &deps, wasmContract, 5, commitEvmTx, + ) + stateDB, ok = evmObj.StateDB.(*statedb.StateDB) + s.Require().True(ok, "error retrieving StateDB from the EVM") + + s.T().Log("Expect exactly 1 dirty journal entry for the precompile snapshot") + if stateDB.DirtiesCount() != 1 { + debugDirtiesCountMismatch(stateDB, s.T()) + s.FailNow("expected 1 dirty journal changes") + } + + s.T().Log("Expect no change since the StateDB has not been committed") + test.AssertWasmCounterState( + &s.Suite, deps, wasmContract, 7, // 7 = 7 + 0 + ) + + s.T().Log("Expect change after the StateDB gets committed") + err = stateDB.Commit() + s.Require().NoError(err) + test.AssertWasmCounterState( + &s.Suite, deps, wasmContract, 12, // 12 = 7 + 5 + ) }) } diff --git a/x/evm/statedb/state_object.go b/x/evm/statedb/state_object.go index e371beae0..3e546362b 100644 --- a/x/evm/statedb/state_object.go +++ b/x/evm/statedb/state_object.go @@ -121,8 +121,9 @@ type stateObject struct { address common.Address // flags - DirtyCode bool - Suicided bool + DirtyCode bool + Suicided bool + IsPrecompile bool } // newObject creates a state object. diff --git a/x/evm/statedb/statedb.go b/x/evm/statedb/statedb.go index 223e92edb..839b40815 100644 --- a/x/evm/statedb/statedb.go +++ b/x/evm/statedb/statedb.go @@ -43,6 +43,21 @@ type StateDB struct { txConfig TxConfig + // cacheCtx: An sdk.Context produced from the [StateDB.ctx] with the + // multi-store cached and a new event manager. The cached context + // (`cacheCtx`) is written to the persistent context (`ctx`) when + // `writeCacheCtx` is called. + cacheCtx sdk.Context + + // writeToCommitCtxFromCacheCtx is the "write" function received from + // `s.ctx.CacheContext()`. It saves mutations on s.cacheCtx to the StateDB's + // commit context (s.ctx). This synchronizes the multistore and event manager + // of the two contexts. + writeToCommitCtxFromCacheCtx func() + + // The number of precompiled contract calls within the current transaction + precompileSnapshotsCount uint8 + // The refund counter, also used by state transitioning. refund uint64 @@ -212,6 +227,16 @@ func (s *StateDB) getStateObject(addr common.Address) *stateObject { if obj := s.stateObjects[addr]; obj != nil { return obj } + + if s.keeper.IsPrecompile(addr) { + obj := newObject(s, addr, Account{ + Nonce: 0, + }) + obj.IsPrecompile = true + s.setStateObject(obj) + return obj + } + // If no live objects are available, load it from keeper account := s.keeper.GetAccount(s.ctx, addr) if account == nil { @@ -461,13 +486,38 @@ func errorf(format string, args ...any) error { // StateDB object cannot be reused after [Commit] has completed. A new // object needs to be created from the EVM. func (s *StateDB) Commit() error { - ctx := s.GetContext() + if s.writeToCommitCtxFromCacheCtx != nil { + // cacheCtxSyncNeeded: If a precompile was called, a [JournalChange] + // of type [PrecompileSnapshotBeforeRun] gets added and we branch off a + // cache of the commit context (s.ctx). + s.writeToCommitCtxFromCacheCtx() + } + return s.commitCtx(s.GetContext()) +} + +// CommitCacheCtx is identical to [StateDB.Commit], except it: +// (1) uses the cacheCtx of the [StateDB] and +// (2) does not save mutations of the cacheCtx to the commit context (s.ctx). +// The reason for (2) is that the overall EVM transaction (block, not internal) +// is only finalized when [Commit] is called, not when [CommitCacheCtx] is +// called. +func (s *StateDB) CommitCacheCtx() error { + return s.commitCtx(s.cacheCtx) +} + +// commitCtx writes the dirty journal state changes to the EVM Keeper. The +// StateDB object cannot be reused after [commitCtx] has completed. A new +// object needs to be created from the EVM. +func (s *StateDB) commitCtx(ctx sdk.Context) error { for _, addr := range s.Journal.sortedDirties() { obj := s.getStateObject(addr) if obj == nil { continue } - if obj.Suicided { + if obj.IsPrecompile { + s.Journal.dirties[addr] = 0 + continue + } else if obj.Suicided { // Invariant: After [StateDB.Suicide] for some address, the // corresponding account's state object is only available until the // state is committed. @@ -493,13 +543,49 @@ func (s *StateDB) Commit() error { obj.OriginStorage[key] = dirtyVal } } - // Clear the dirty counts because all state changes have been - // committed. + // Reset the dirty count to 0 because all state changes for this dirtied + // address in the journal have been committed. s.Journal.dirties[addr] = 0 } return nil } +func (s *StateDB) CacheCtxForPrecompile(precompileAddr common.Address) ( + sdk.Context, PrecompileSnapshotBeforeRun, +) { + if s.writeToCommitCtxFromCacheCtx == nil { + s.cacheCtx, s.writeToCommitCtxFromCacheCtx = s.ctx.CacheContext() + } + return s.cacheCtx, PrecompileSnapshotBeforeRun{ + MultiStore: s.cacheCtx.MultiStore().CacheMultiStore(), + Events: s.cacheCtx.EventManager().Events(), + Precompile: precompileAddr, + } +} + +// SavePrecompileSnapshotToJournal adds a snapshot of the commit multistore +// ([PrecompileSnapshotBeforeRun]) to the [StateDB] journal at the end of +// successful invocation of a precompiled contract. This is necessary to revert +// intermediate states where an EVM contract augments the multistore with a +// precompile and an inconsistency occurs between the EVM module and other +// modules. +// +// See [PrecompileSnapshotBeforeRun] for more info. +func (s *StateDB) SavePrecompileSnapshotToJournal( + precompileAddr common.Address, + snapshot PrecompileSnapshotBeforeRun, +) error { + obj := s.getOrNewStateObject(precompileAddr) + obj.db.Journal.append(snapshot) + s.precompileSnapshotsCount++ + if s.precompileSnapshotsCount > maxPrecompileCalls { + return fmt.Errorf("exceeded maximum allowed number of precompiled contract calls in one transaction (%d)", maxPrecompileCalls) + } + return nil +} + +const maxPrecompileCalls uint8 = 10 + // StateObjects: Returns a copy of the [StateDB.stateObjects] map. func (s *StateDB) StateObjects() map[common.Address]*stateObject { copyOfMap := make(map[common.Address]*stateObject)