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

## Why this should be merged

Aligns precompiled contracts with the pattern used in
`vm.EVMInterpreter` for regular contracts and simplifies gas accounting
by using existing mechanisms. The most notable simplification occurs
when there are multiple error paths that return early and have to
account for consumed gas (any local solution would likely just mirror
this one).

This forms part of a broader, long-term direction of feature parity
between precompiled and bytecode contracts, exposed via
`vm.PrecompileEnvironment`.

The `vm.Contract.Value()` method is also exposed as a natural
accompaniment to this change.

## How this works

The original signature for `vm.PrecompiledStatefulContract` received the
amount of gas available and then returned the amount remaining; this has
been removed in lieu of exposing the existing `vm.Contract.UseGas()`
method via the `vm.PrecompileEnvironment`. A new `legacy` package wraps
the old signature and converts it to a new one so `ava-labs/coreth`
doesn't need to be refactored.

```go
func oldPrecompile(env vm.PrecompileEnvironment, input []byte, gas uint64) ([]byte, uint64, error) {
  // ...
  if err != nil {
    return nil, gas - gasCost, err // pattern susceptible to bugs; should it be `nil, gas, err` ?
  }
  // ...
  return output, gas - gasCost, nil  
}

func newPrecompile(env vm.PrecompileEnvironment, input []byte) ([]byte, error) {
  // The original `gas` argument is still available as `env.Gas()`
  // ...
  if !env.UseGas(gasCost) { // an explicit point at which gas is consumed
    return nil, vm.ErrOutOfGas
  }
  // ...
  if err != nil {
    return nil, err
  }
  // ...
  return output, nil
}
```

## How this was tested

Existing unit test modified to use the `legacy` adaptor.

---------

Signed-off-by: Arran Schlosberg <[email protected]>
Co-authored-by: Quentin McGaw <[email protected]>

Author: ARR4N

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 `args.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) (hasEnoughGas 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/libevm/options"
@@ -37,13 +38,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.
@@ -87,11 +101,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,8 +132,12 @@ func (e *environment) callContract(typ CallType, addr common.Address, input []by
 	}
 
 	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 {
@@ -129,13 +147,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, callErr := e.evm.Call(caller, addr, input, gas, value)
+		if err := e.refundGas(returnGas); err != nil {
+			return nil, err
+		}
+		return ret, callErr
 	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
@@ -144,6 +166,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)
 }