feat: override EVM.Reset() args (#36)

core/vm/contracts.libevm.go

@@ -13,6 +13,7 @@
 // You should have received a copy of the GNU Lesser General Public License
 // along with the go-ethereum library. If not, see
 // <http://www.gnu.org/licenses/>.
+
 package vm
 
 import (

core/vm/evm.go

@@ -159,8 +159,7 @@ func NewEVM(blockCtx BlockContext, txCtx TxContext, statedb StateDB, chainConfig
 // Reset resets the EVM with a new transaction context.Reset
 // This is not threadsafe and should only be done very cautiously.
 func (evm *EVM) Reset(txCtx TxContext, statedb StateDB) {
-	evm.TxContext = txCtx
-	evm.StateDB = statedb
+	evm.TxContext, evm.StateDB = overrideEVMResetArgs(txCtx, statedb)
 }
 
 // Cancel cancels any running EVM operation. This may be called concurrently and

core/vm/evm.libevm_test.go

@@ -24,26 +24,62 @@ import (
 	"github.com/ethereum/go-ethereum/params"
 )
 
-type chainIDOverrider struct {
-	chainID int64
+type evmArgOverrider struct {
+	newEVMchainID int64
+
+	resetTxCtx   TxContext
+	resetStateDB StateDB
 }
 
-func (o chainIDOverrider) OverrideNewEVMArgs(args *NewEVMArgs) *NewEVMArgs {
-	args.ChainConfig = &params.ChainConfig{ChainID: big.NewInt(o.chainID)}
+func (o evmArgOverrider) OverrideNewEVMArgs(args *NewEVMArgs) *NewEVMArgs {
+	args.ChainConfig = &params.ChainConfig{ChainID: big.NewInt(o.newEVMchainID)}
 	return args
 }
 
+func (o evmArgOverrider) OverrideEVMResetArgs(*EVMResetArgs) *EVMResetArgs {
+	return &EVMResetArgs{
+		TxContext: o.resetTxCtx,
+		StateDB:   o.resetStateDB,
+	}
+}
+
+func (o evmArgOverrider) register(t *testing.T) {
+	t.Helper()
+	libevmHooks = nil
+	RegisterHooks(o)
+	t.Cleanup(func() {
+		libevmHooks = nil
+	})
+}
+
 func TestOverrideNewEVMArgs(t *testing.T) {
 	// The overrideNewEVMArgs function accepts and returns all arguments to
 	// NewEVM(), in order. Here we lock in our assumption of that order. If this
 	// breaks then all functionality overriding the args MUST be updated.
 	var _ func(BlockContext, TxContext, StateDB, *params.ChainConfig, Config) *EVM = NewEVM
 
 	const chainID = 13579
-	libevmHooks = nil
-	RegisterHooks(chainIDOverrider{chainID: chainID})
-	defer func() { libevmHooks = nil }()
+	hooks := evmArgOverrider{newEVMchainID: chainID}
+	hooks.register(t)
+
+	evm := NewEVM(BlockContext{}, TxContext{}, nil, nil, Config{})
+	got := evm.ChainConfig().ChainID
+	require.Equalf(t, big.NewInt(chainID), got, "%T.ChainConfig().ChainID set by NewEVM() hook", evm)
+}
+
+func TestOverrideEVMResetArgs(t *testing.T) {
+	// Equivalent to rationale for TestOverrideNewEVMArgs above.
+	var _ func(TxContext, StateDB) = (*EVM)(nil).Reset
+
+	const gasPrice = 1357924680
+	hooks := evmArgOverrider{
+		resetTxCtx: TxContext{
+			GasPrice: big.NewInt(gasPrice),
+		},
+	}
+	hooks.register(t)
 
-	got := NewEVM(BlockContext{}, TxContext{}, nil, nil, Config{}).ChainConfig().ChainID
-	require.Equal(t, big.NewInt(chainID), got)
+	evm := NewEVM(BlockContext{}, TxContext{}, nil, nil, Config{})
+	evm.Reset(TxContext{}, nil)
+	require.Equalf(t, big.NewInt(gasPrice), evm.GasPrice, "%T.GasPrice set by Reset() hook", evm)
 }

core/vm/hooks.libevm.go

@@ -13,6 +13,7 @@
 // You should have received a copy of the GNU Lesser General Public License
 // along with the go-ethereum library. If not, see
 // <http://www.gnu.org/licenses/>.
+
 package vm
 
 import "github.com/ethereum/go-ethereum/params"
@@ -29,11 +30,14 @@ func RegisterHooks(h Hooks) {
 var libevmHooks Hooks
 
 // Hooks are arbitrary configuration functions to modify default VM behaviour.
+// See [RegisterHooks].
 type Hooks interface {
 	OverrideNewEVMArgs(*NewEVMArgs) *NewEVMArgs
+	OverrideEVMResetArgs(*EVMResetArgs) *EVMResetArgs
 }
 
-// NewEVMArgs are the arguments received by [NewEVM], available for override.
+// NewEVMArgs are the arguments received by [NewEVM], available for override
+// via [Hooks].
 type NewEVMArgs struct {
 	BlockContext BlockContext
 	TxContext    TxContext
@@ -42,6 +46,13 @@ type NewEVMArgs struct {
 	Config       Config
 }
 
+// EVMResetArgs are the arguments received by [EVM.Reset], available for
+// override via [Hooks].
+type EVMResetArgs struct {
+	TxContext TxContext
+	StateDB   StateDB
+}
+
 func overrideNewEVMArgs(
 	blockCtx BlockContext,
 	txCtx TxContext,
@@ -55,3 +66,11 @@ func overrideNewEVMArgs(
 	args := libevmHooks.OverrideNewEVMArgs(&NewEVMArgs{blockCtx, txCtx, statedb, chainConfig, config})
 	return args.BlockContext, args.TxContext, args.StateDB, args.ChainConfig, args.Config
 }
+
+func overrideEVMResetArgs(txCtx TxContext, statedb StateDB) (TxContext, StateDB) {
+	if libevmHooks == nil {
+		return txCtx, statedb
+	}
+	args := libevmHooks.OverrideEVMResetArgs(&EVMResetArgs{txCtx, statedb})
+	return args.TxContext, args.StateDB
+}

core/vm/stack.libevm.go

@@ -13,6 +13,7 @@
 // You should have received a copy of the GNU Lesser General Public License
 // along with the go-ethereum library. If not, see
 // <http://www.gnu.org/licenses/>.
+
 package vm
 
 import "github.com/holiman/uint256"