feat: override vm.NewEVM() args (#35)

* feat: override `vm.NewEVM()` args

* test: `vm.Hooks.OverrideNewEVMArgs`

* refactor: use `NewEVMArgs struct` in hook

* chore: `gci`

* fix: clear `vm.Hooks` at end of test

Author: ARR4N

core/vm/evm.go

@@ -143,6 +143,7 @@ func NewEVM(blockCtx BlockContext, txCtx TxContext, statedb StateDB, chainConfig
 			blockCtx.BlobBaseFee = new(big.Int)
 		}
 	}
+	blockCtx, txCtx, statedb, chainConfig, config = overrideNewEVMArgs(blockCtx, txCtx, statedb, chainConfig, config)
 	evm := &EVM{
 		Context:     blockCtx,
 		TxContext:   txCtx,

core/vm/evm.libevm_test.go

@@ -0,0 +1,34 @@
+package vm
+
+import (
+	"math/big"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ethereum/go-ethereum/params"
+)
+
+type chainIDOverrider struct {
+	chainID int64
+}
+
+func (o chainIDOverrider) OverrideNewEVMArgs(args *NewEVMArgs) *NewEVMArgs {
+	args.ChainConfig = &params.ChainConfig{ChainID: big.NewInt(o.chainID)}
+	return args
+}
+
+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 }()
+
+	got := NewEVM(BlockContext{}, TxContext{}, nil, nil, Config{}).ChainConfig().ChainID
+	require.Equal(t, big.NewInt(chainID), got)
+}

core/vm/hooks.libevm.go

@@ -0,0 +1,42 @@
+package vm
+
+import "github.com/ethereum/go-ethereum/params"
+
+// RegisterHooks registers the Hooks. It is expected to be called in an `init()`
+// function and MUST NOT be called more than once.
+func RegisterHooks(h Hooks) {
+	if libevmHooks != nil {
+		panic("already registered")
+	}
+	libevmHooks = h
+}
+
+var libevmHooks Hooks
+
+// Hooks are arbitrary configuration functions to modify default VM behaviour.
+type Hooks interface {
+	OverrideNewEVMArgs(*NewEVMArgs) *NewEVMArgs
+}
+
+// NewEVMArgs are the arguments received by [NewEVM], available for override.
+type NewEVMArgs struct {
+	BlockContext BlockContext
+	TxContext    TxContext
+	StateDB      StateDB
+	ChainConfig  *params.ChainConfig
+	Config       Config
+}
+
+func overrideNewEVMArgs(
+	blockCtx BlockContext,
+	txCtx TxContext,
+	statedb StateDB,
+	chainConfig *params.ChainConfig,
+	config Config,
+) (BlockContext, TxContext, StateDB, *params.ChainConfig, Config) {
+	if libevmHooks == nil {
+		return blockCtx, txCtx, statedb, chainConfig, config
+	}
+	args := libevmHooks.OverrideNewEVMArgs(&NewEVMArgs{blockCtx, txCtx, statedb, chainConfig, config})
+	return args.BlockContext, args.TxContext, args.StateDB, args.ChainConfig, args.Config
+}