From 1a6dd02a181b4238cbbb33e70f8205460e0d5e20 Mon Sep 17 00:00:00 2001
From: Arran Schlosberg <519948+ARR4N@users.noreply.github.com>
Date: Thu, 19 Sep 2024 16:54:52 -0400
Subject: [PATCH] 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
---
 core/vm/evm.go             |  1 +
 core/vm/evm.libevm_test.go | 34 ++++++++++++++++++++++++++++++
 core/vm/hooks.libevm.go    | 42 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 77 insertions(+)
 create mode 100644 core/vm/evm.libevm_test.go
 create mode 100644 core/vm/hooks.libevm.go

diff --git a/core/vm/evm.go b/core/vm/evm.go
index 0cca927264a9..d61f19f3cdc1 100644
--- a/core/vm/evm.go
+++ b/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,
diff --git a/core/vm/evm.libevm_test.go b/core/vm/evm.libevm_test.go
new file mode 100644
index 000000000000..3987ab1f2fe4
--- /dev/null
+++ b/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)
+}
diff --git a/core/vm/hooks.libevm.go b/core/vm/hooks.libevm.go
new file mode 100644
index 000000000000..bd5668b593d2
--- /dev/null
+++ b/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
+}