refactor: PrecompileEnvironment.{,Use,Refund}Gas() in lieu of args

core/vm/contracts.go

@@ -177,7 +177,9 @@ func (args *evmCallArgs) RunPrecompiledContract(p PrecompiledContract, input []b
 		return nil, 0, ErrOutOfGas
 	}
 	suppliedGas -= gasCost
-	return args.run(p, input, suppliedGas)
+	args.gasRemaining = suppliedGas
+	output, err := args.run(p, input)
+	return output, args.gasRemaining, err
 }
 
 // ECRECOVER implemented as a native contract.

core/vm/contracts.libevm.go

@@ -45,11 +45,11 @@ type evmCallArgs struct {
 	callType CallType
 
 	// args:start
-	caller ContractRef
-	addr   common.Address
-	input  []byte
-	gas    uint64
-	value  *uint256.Int
+	caller       ContractRef
+	addr         common.Address
+	input        []byte
+	gasRemaining uint64
+	value        *uint256.Int
 	// args:end
 }
 
@@ -89,20 +89,26 @@ func (t CallType) OpCode() OpCode {
 }
 
 // run runs the [PrecompiledContract], differentiating between stateful and
-// regular types.
-func (args *evmCallArgs) run(p PrecompiledContract, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error) {
-	if p, ok := p.(statefulPrecompile); ok {
-		return p(args.env(), input, suppliedGas)
+// regular types, updating `gasRemaining` in the stateful case.
+func (args *evmCallArgs) run(p PrecompiledContract, input []byte) (ret []byte, err error) {
+	switch p := p.(type) {
+	default:
+		return p.Run(input)
+	case statefulPrecompile:
+		env := args.env()
+		ret, err := p(env, input)
+		args.gasRemaining = env.Gas()
+		return ret, err
 	}
-	// Gas consumption for regular precompiles was already handled by the native
-	// RunPrecompiledContract(), which called this method.
-	ret, err = p.Run(input)
-	return ret, suppliedGas, err
 }
 
 // PrecompiledStatefulContract is the stateful equivalent of a
 // [PrecompiledContract].
-type PrecompiledStatefulContract func(env PrecompileEnvironment, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error)
+//
+// Instead of receiving and returning gas arguments, stateful precompiles use
+// the respective methods on [PrecompileEnvironment]. If a call to UseGas()
+// returns false, a stateful precompile SHOULD return [ErrOutOfGas].
+type PrecompiledStatefulContract func(env PrecompileEnvironment, input []byte) (ret []byte, err error)
 
 // NewStatefulPrecompile constructs a new PrecompiledContract that can be used
 // via an [EVM] instance but MUST NOT be called directly; a direct call to Run()
@@ -135,13 +141,18 @@ func (p statefulPrecompile) Run([]byte) ([]byte, error) {
 type PrecompileEnvironment interface {
 	ChainConfig() *params.ChainConfig
 	Rules() params.Rules
-	ReadOnly() bool
 	// StateDB will be non-nil i.f.f !ReadOnly().
 	StateDB() StateDB
 	// ReadOnlyState will always be non-nil.
 	ReadOnlyState() libevm.StateReader
-	Addresses() *libevm.AddressContext
+
 	IncomingCallType() CallType
+	Addresses() *libevm.AddressContext
+	ReadOnly() bool
+	// Equivalent to respective methods on [Contract].
+	Gas() uint64
+	UseGas(uint64) bool
+	Value() *uint256.Int
 
 	BlockHeader() (types.Header, error)
 	BlockNumber() *big.Int
@@ -150,7 +161,7 @@ type PrecompileEnvironment interface {
 	// Call is equivalent to [EVM.Call] except that the `caller` argument is
 	// removed and automatically determined according to the type of call that
 	// invoked the precompile.
-	Call(addr common.Address, input []byte, gas uint64, value *uint256.Int, _ ...CallOption) (ret []byte, gasRemaining uint64, _ error)
+	Call(addr common.Address, input []byte, gas uint64, value *uint256.Int, _ ...CallOption) (ret []byte, _ error)
 }
 
 func (args *evmCallArgs) env() *environment {
@@ -174,7 +185,7 @@ func (args *evmCallArgs) env() *environment {
 
 	// This is equivalent to the `contract` variables created by evm.*Call*()
 	// methods, for non precompiles, to pass to [EVMInterpreter.Run].
-	contract := NewContract(args.caller, AccountRef(self), value, args.gas)
+	contract := NewContract(args.caller, AccountRef(self), value, args.gasRemaining)
 	if args.callType == DelegateCall {
 		contract = contract.AsDelegate()
 	}

core/vm/contracts.libevm_test.go

@@ -39,6 +39,7 @@ import (
 	"github.com/ava-labs/libevm/libevm"
 	"github.com/ava-labs/libevm/libevm/ethtest"
 	"github.com/ava-labs/libevm/libevm/hookstest"
+	"github.com/ava-labs/libevm/libevm/legacy"
 	"github.com/ava-labs/libevm/params"
 )
 
@@ -106,6 +107,7 @@ type statefulPrecompileOutput struct {
 	ChainID                 *big.Int
 	Addresses               *libevm.AddressContext
 	StateValue              common.Hash
+	ValueReceived           *uint256.Int
 	ReadOnly                bool
 	BlockNumber, Difficulty *big.Int
 	BlockTime               uint64
@@ -159,6 +161,7 @@ func TestNewStatefulPrecompile(t *testing.T) {
 			ChainID:          env.ChainConfig().ChainID,
 			Addresses:        env.Addresses(),
 			StateValue:       env.ReadOnlyState().GetState(precompile, slot),
+			ValueReceived:    env.Value(),
 			ReadOnly:         env.ReadOnly(),
 			BlockNumber:      env.BlockNumber(),
 			BlockTime:        env.BlockTime(),
@@ -170,7 +173,11 @@ func TestNewStatefulPrecompile(t *testing.T) {
 	}
 	hooks := &hookstest.Stub{
 		PrecompileOverrides: map[common.Address]libevm.PrecompiledContract{
-			precompile: vm.NewStatefulPrecompile(run),
+			precompile: vm.NewStatefulPrecompile(
+				// In production, the new function signature should be used, but
+				// this just exercises the converter.
+				legacy.PrecompiledStatefulContract(run).Upgrade(),
+			),
 		},
 	}
 	hooks.Register(t)
@@ -181,7 +188,8 @@ func TestNewStatefulPrecompile(t *testing.T) {
 		Difficulty: rng.BigUint64(),
 	}
 	input := rng.Bytes(8)
-	value := rng.Hash()
+	stateValue := rng.Hash()
+	transferValue := rng.Uint256()
 	chainID := rng.BigUint64()
 
 	caller := common.HexToAddress("CA11E12") // caller of the precompile
@@ -197,13 +205,15 @@ func TestNewStatefulPrecompile(t *testing.T) {
 			&params.ChainConfig{ChainID: chainID},
 		),
 	)
-	state.SetState(precompile, slot, value)
+	state.SetState(precompile, slot, stateValue)
+	state.SetBalance(caller, new(uint256.Int).Not(uint256.NewInt(0)))
 	evm.Origin = eoa
 
 	tests := []struct {
-		name          string
-		call          func() ([]byte, uint64, error)
-		wantAddresses *libevm.AddressContext
+		name              string
+		call              func() ([]byte, uint64, error)
+		wantAddresses     *libevm.AddressContext
+		wantTransferValue *uint256.Int
 		// Note that this only covers evm.readOnly being true because of the
 		// precompile's call. See TestInheritReadOnly for alternate case.
 		wantReadOnly bool
@@ -212,28 +222,30 @@ func TestNewStatefulPrecompile(t *testing.T) {
 		{
 			name: "EVM.Call()",
 			call: func() ([]byte, uint64, error) {
-				return evm.Call(callerContract, precompile, input, gasLimit, uint256.NewInt(0))
+				return evm.Call(callerContract, precompile, input, gasLimit, transferValue)
 			},
 			wantAddresses: &libevm.AddressContext{
 				Origin: eoa,
 				Caller: caller,
 				Self:   precompile,
 			},
-			wantReadOnly: false,
-			wantCallType: vm.Call,
+			wantReadOnly:      false,
+			wantTransferValue: transferValue,
+			wantCallType:      vm.Call,
 		},
 		{
 			name: "EVM.CallCode()",
 			call: func() ([]byte, uint64, error) {
-				return evm.CallCode(callerContract, precompile, input, gasLimit, uint256.NewInt(0))
+				return evm.CallCode(callerContract, precompile, input, gasLimit, transferValue)
 			},
 			wantAddresses: &libevm.AddressContext{
 				Origin: eoa,
 				Caller: caller,
 				Self:   caller,
 			},
-			wantReadOnly: false,
-			wantCallType: vm.CallCode,
+			wantReadOnly:      false,
+			wantTransferValue: transferValue,
+			wantCallType:      vm.CallCode,
 		},
 		{
 			name: "EVM.DelegateCall()",
@@ -245,8 +257,9 @@ func TestNewStatefulPrecompile(t *testing.T) {
 				Caller: eoa, // inherited from caller
 				Self:   caller,
 			},
-			wantReadOnly: false,
-			wantCallType: vm.DelegateCall,
+			wantReadOnly:      false,
+			wantTransferValue: uint256.NewInt(0),
+			wantCallType:      vm.DelegateCall,
 		},
 		{
 			name: "EVM.StaticCall()",
@@ -258,8 +271,9 @@ func TestNewStatefulPrecompile(t *testing.T) {
 				Caller: caller,
 				Self:   precompile,
 			},
-			wantReadOnly: true,
-			wantCallType: vm.StaticCall,
+			wantReadOnly:      true,
+			wantTransferValue: uint256.NewInt(0),
+			wantCallType:      vm.StaticCall,
 		},
 	}
 
@@ -268,7 +282,8 @@ func TestNewStatefulPrecompile(t *testing.T) {
 			wantOutput := statefulPrecompileOutput{
 				ChainID:          chainID,
 				Addresses:        tt.wantAddresses,
-				StateValue:       value,
+				StateValue:       stateValue,
+				ValueReceived:    tt.wantTransferValue,
 				ReadOnly:         tt.wantReadOnly,
 				BlockNumber:      header.Number,
 				BlockTime:        header.Time,
@@ -318,11 +333,11 @@ func TestInheritReadOnly(t *testing.T) {
 	hooks := &hookstest.Stub{
 		PrecompileOverrides: map[common.Address]libevm.PrecompiledContract{
 			precompile: vm.NewStatefulPrecompile(
-				func(env vm.PrecompileEnvironment, input []byte, suppliedGas uint64) ([]byte, uint64, error) {
+				func(env vm.PrecompileEnvironment, input []byte) ([]byte, error) {
 					if env.ReadOnly() {
-						return []byte{ifReadOnly}, suppliedGas, nil
+						return []byte{ifReadOnly}, nil
 					}
-					return []byte{ifNotReadOnly}, suppliedGas, nil
+					return []byte{ifNotReadOnly}, nil
 				},
 			),
 		},
@@ -535,21 +550,21 @@ func TestPrecompileMakeCall(t *testing.T) {
 
 	hooks := &hookstest.Stub{
 		PrecompileOverrides: map[common.Address]libevm.PrecompiledContract{
-			sut: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error) {
+			sut: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte) (ret []byte, err error) {
 				var opts []vm.CallOption
 				if bytes.Equal(input, unsafeCallerProxyOptSentinel) {
 					opts = append(opts, vm.WithUNSAFECallerAddressProxying())
 				}
 				// We are ultimately testing env.Call(), hence why this is the SUT.
-				return env.Call(dest, precompileCallData, suppliedGas, uint256.NewInt(0), opts...)
+				return env.Call(dest, precompileCallData, env.Gas(), uint256.NewInt(0), opts...)
 			}),
-			dest: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error) {
+			dest: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte) (ret []byte, err error) {
 				out := &statefulPrecompileOutput{
 					Addresses: env.Addresses(),
 					ReadOnly:  env.ReadOnly(),
 					Input:     input, // expected to be callData
 				}
-				return out.Bytes(), suppliedGas, nil
+				return out.Bytes(), nil
 			}),
 		},
 	}
@@ -696,8 +711,8 @@ func TestPrecompileCallWithTracer(t *testing.T) {
 
 	hooks := &hookstest.Stub{
 		PrecompileOverrides: map[common.Address]libevm.PrecompiledContract{
-			precompile: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error) {
-				return env.Call(contract, nil, suppliedGas, uint256.NewInt(0))
+			precompile: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte) (ret []byte, err error) {
+				return env.Call(contract, nil, env.Gas(), uint256.NewInt(0))
 			}),
 		},
 	}

core/vm/environment.libevm.go

@@ -23,6 +23,7 @@ import (
 	"github.com/holiman/uint256"
 
 	"github.com/ava-labs/libevm/common"
+	"github.com/ava-labs/libevm/common/math"
 	"github.com/ava-labs/libevm/core/types"
 	"github.com/ava-labs/libevm/libevm"
 	"github.com/ava-labs/libevm/params"
@@ -36,13 +37,26 @@ type environment struct {
 	callType CallType
 }
 
+func (e *environment) Gas() uint64            { return e.self.Gas }
+func (e *environment) UseGas(gas uint64) bool { return e.self.UseGas(gas) }
+func (e *environment) Value() *uint256.Int    { return new(uint256.Int).Set(e.self.Value()) }
+
 func (e *environment) ChainConfig() *params.ChainConfig  { return e.evm.chainConfig }
 func (e *environment) Rules() params.Rules               { return e.evm.chainRules }
 func (e *environment) ReadOnlyState() libevm.StateReader { return e.evm.StateDB }
 func (e *environment) IncomingCallType() CallType        { return e.callType }
 func (e *environment) BlockNumber() *big.Int             { return new(big.Int).Set(e.evm.Context.BlockNumber) }
 func (e *environment) BlockTime() uint64                 { return e.evm.Context.Time }
 
+func (e *environment) refundGas(add uint64) error {
+	gas, overflow := math.SafeAdd(e.self.Gas, add)
+	if overflow {
+		return ErrGasUintOverflow
+	}
+	e.self.Gas = gas
+	return nil
+}
+
 func (e *environment) ReadOnly() bool {
 	// A switch statement provides clearer code coverage for difficult-to-test
 	// cases.
@@ -86,11 +100,11 @@ func (e *environment) BlockHeader() (types.Header, error) {
 	return *hdr, nil
 }
 
-func (e *environment) Call(addr common.Address, input []byte, gas uint64, value *uint256.Int, opts ...CallOption) ([]byte, uint64, error) {
+func (e *environment) Call(addr common.Address, input []byte, gas uint64, value *uint256.Int, opts ...CallOption) ([]byte, error) {
 	return e.callContract(Call, addr, input, gas, value, opts...)
 }
 
-func (e *environment) callContract(typ CallType, addr common.Address, input []byte, gas uint64, value *uint256.Int, opts ...CallOption) (retData []byte, retGas uint64, retErr error) {
+func (e *environment) callContract(typ CallType, addr common.Address, input []byte, gas uint64, value *uint256.Int, opts ...CallOption) (retData []byte, retErr error) {
 	// Depth and read-only setting are handled by [EVMInterpreter.Run], which
 	// isn't used for precompiles, so we need to do it ourselves to maintain the
 	// expected invariants.
@@ -118,13 +132,17 @@ func (e *environment) callContract(typ CallType, addr common.Address, input []by
 			}
 		case nil:
 		default:
-			return nil, gas, fmt.Errorf("unsupported option %T", o)
+			return nil, fmt.Errorf("unsupported option %T", o)
 		}
 	}
 
 	if in.readOnly && value != nil && !value.IsZero() {
-		return nil, gas, ErrWriteProtection
+		return nil, ErrWriteProtection
 	}
+	if !e.UseGas(gas) {
+		return nil, ErrOutOfGas
+	}
+
 	if t := e.evm.Config.Tracer; t != nil {
 		var bigVal *big.Int
 		if value != nil {
@@ -134,13 +152,17 @@ func (e *environment) callContract(typ CallType, addr common.Address, input []by
 
 		startGas := gas
 		defer func() {
-			t.CaptureEnd(retData, startGas-retGas, retErr)
+			t.CaptureEnd(retData, startGas-e.Gas(), retErr)
 		}()
 	}
 
 	switch typ {
 	case Call:
-		return e.evm.Call(caller, addr, input, gas, value)
+		ret, returnGas, err := e.evm.Call(caller, addr, input, gas, value)
+		if err := e.refundGas(returnGas); err != nil {
+			return nil, err
+		}
+		return ret, err
 	case CallCode, DelegateCall, StaticCall:
 		// TODO(arr4n): these cases should be very similar to CALL, hence the
 		// early abstraction, to signal to future maintainers. If implementing
@@ -149,6 +171,6 @@ func (e *environment) callContract(typ CallType, addr common.Address, input []by
 		// compatibility.
 		fallthrough
 	default:
-		return nil, gas, fmt.Errorf("unimplemented precompile call type %v", typ)
+		return nil, fmt.Errorf("unimplemented precompile call type %v", typ)
 	}
 }

core/vm/libevm_test.go

@@ -18,5 +18,5 @@ package vm
 // The original RunPrecompiledContract was migrated to being a method on
 // [evmCallArgs]. We need to replace it for use by regular geth tests.
 func RunPrecompiledContract(p PrecompiledContract, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error) {
-	return (*evmCallArgs)(nil).RunPrecompiledContract(p, input, suppliedGas)
+	return new(evmCallArgs).RunPrecompiledContract(p, input, suppliedGas)
 }