feat: vm.OperationEnvironment for custom instructions

Author: ARR4N

core/vm/jump_table.libevm.go

@@ -18,7 +18,7 @@ func overrideJumpTable(r params.Rules, jt *JumpTable) *JumpTable {
 // An OperationBuilder is a factory for a new operations to include in a
 // [JumpTable].
 type OperationBuilder struct {
-	Execute            func(pc *uint64, interpreter *EVMInterpreter, callContext *ScopeContext) ([]byte, error)
+	Execute            OperationFunc
 	ConstantGas        uint64
 	DynamicGas         func(_ *EVM, _ *Contract, _ *Stack, _ *Memory, requestedMemorySize uint64) (uint64, error)
 	MinStack, MaxStack int
@@ -28,7 +28,7 @@ type OperationBuilder struct {
 // Build constructs the operation.
 func (b OperationBuilder) Build() *operation {
 	o := &operation{
-		execute:     b.Execute,
+		execute:     b.Execute.internal(),
 		constantGas: b.ConstantGas,
 		dynamicGas:  b.DynamicGas,
 		minStack:    b.MinStack,
@@ -38,6 +38,27 @@ func (b OperationBuilder) Build() *operation {
 	return o
 }
 
+// An OperationFunc is the execution function of a custom instruction.
+type OperationFunc func(_ *OperationEnvironment, pc *uint64, _ *EVMInterpreter, _ *ScopeContext) ([]byte, error)
+
+// An OperationEnvironment provides information about the context in which a
+// custom instruction is being executed.
+type OperationEnvironment struct {
+	StateDB StateDB
+}
+
+// internal converts an exported [OperationFunc] into an un-exported
+// [executionFunc] as required to build an [operation].
+func (fn OperationFunc) internal() executionFunc {
+	return func(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([]byte, error) {
+		return fn(
+			&OperationEnvironment{
+				StateDB: interpreter.evm.StateDB,
+			}, pc, interpreter, scope,
+		)
+	}
+}
+
 // Hooks are arbitrary configuration functions to modify default VM behaviour.
 type Hooks interface {
 	// OverrideJumpTable will only be called if

core/vm/jump_table.libevm_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/vm"
 	"github.com/ethereum/go-ethereum/libevm/ethtest"
 	"github.com/ethereum/go-ethereum/libevm/hookstest"
@@ -35,6 +36,16 @@ func (s *vmHooksStub) OverrideJumpTable(_ params.Rules, jt *vm.JumpTable) *vm.Ju
 	return jt
 }
 
+// An opRecorder is an instruction that records its inputs.
+type opRecorder struct {
+	stateVal common.Hash
+}
+
+func (op *opRecorder) execute(env *vm.OperationEnvironment, pc *uint64, interpreter *vm.EVMInterpreter, scope *vm.ScopeContext) ([]byte, error) {
+	op.stateVal = env.StateDB.GetState(scope.Contract.Address(), common.Hash{})
+	return nil, nil
+}
+
 func TestOverrideJumpTable(t *testing.T) {
 	override := new(bool)
 	hooks := &hookstest.Stub{
@@ -50,15 +61,12 @@ func TestOverrideJumpTable(t *testing.T) {
 	)
 	rng := ethtest.NewPseudoRand(142857)
 	gasCost := 1 + rng.Uint64n(gasLimit)
-	executed := false
+	spy := &opRecorder{}
 
 	vmHooks := &vmHooksStub{
 		replacement: &vm.JumpTable{
 			opcode: vm.OperationBuilder{
-				Execute: func(pc *uint64, interpreter *vm.EVMInterpreter, callContext *vm.ScopeContext) ([]byte, error) {
-					executed = true
-					return nil, nil
-				},
+				Execute:     spy.execute,
 				ConstantGas: gasCost,
 				MemorySize: func(s *vm.Stack) (size uint64, overflow bool) {
 					return 0, false
@@ -91,11 +99,13 @@ func TestOverrideJumpTable(t *testing.T) {
 		contract := rng.Address()
 		state.CreateAccount(contract)
 		state.SetCode(contract, []byte{opcode})
+		value := rng.Hash()
+		state.SetState(contract, common.Hash{}, value)
 
 		_, gasRemaining, err := evm.Call(vm.AccountRef(rng.Address()), contract, []byte{}, gasLimit, uint256.NewInt(0))
 		require.NoError(t, err, "evm.Call([contract with overridden opcode])")
-		assert.True(t, executed, "executionFunc was called")
 		assert.Equal(t, gasLimit-gasCost, gasRemaining, "gas remaining")
+		assert.Equal(t, spy.stateVal, value, "StateDB propagated")
 	})
 }