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 <519948+ARR4N@users.noreply.github.com>
Co-authored-by: Quentin McGaw <quentin.mcgaw@gmail.com>

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

@@ -48,7 +48,7 @@ type evmCallArgs struct {
 	caller       ContractRef
 	addr         common.Address
 	input        []byte
-	gas    uint64
+	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.
+// regular types, updating `args.gasRemaining` in the stateful case.
-func (args *evmCallArgs) run(p PrecompiledContract, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error) {
+func (args *evmCallArgs) run(p PrecompiledContract, input []byte) (ret []byte, err error) {
-	if p, ok := p.(statefulPrecompile); ok {
+	switch p := p.(type) {
-		return p(args.env(), input, suppliedGas)
+	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
+		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,7 +222,7 @@ 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,
@@ -220,12 +230,13 @@ func TestNewStatefulPrecompile(t *testing.T) {
 				Self:   precompile,
 			},
 			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,
@@ -233,6 +244,7 @@ func TestNewStatefulPrecompile(t *testing.T) {
 				Self:   caller,
 			},
 			wantReadOnly:      false,
+			wantTransferValue: transferValue,
 			wantCallType:      vm.CallCode,
 		},
 		{
@@ -246,6 +258,7 @@ func TestNewStatefulPrecompile(t *testing.T) {
 				Self:   caller,
 			},
 			wantReadOnly:      false,
+			wantTransferValue: uint256.NewInt(0),
 			wantCallType:      vm.DelegateCall,
 		},
 		{
@@ -259,6 +272,7 @@ func TestNewStatefulPrecompile(t *testing.T) {
 				Self:   precompile,
 			},
 			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) {
+			precompile: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte) (ret []byte, err error) {
-				return env.Call(contract, nil, suppliedGas, uint256.NewInt(0))
+				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,6 +38,10 @@ 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 }
@@ -44,6 +49,15 @@ 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)
 }

libevm/legacy/legacy.go

@@ -0,0 +1,39 @@
+// Copyright 2024 the libevm authors.
+//
+// The libevm additions to go-ethereum are free software: you can redistribute
+// them and/or modify them under the terms of the GNU Lesser General Public License
+// as published by the Free Software Foundation, either version 3 of the License,
+// or (at your option) any later version.
+//
+// The libevm additions are distributed in the hope that they will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser
+// General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see
+// <http://www.gnu.org/licenses/>.
+
+// Package legacy provides converters between legacy types and their refactored
+// equivalents.
+package legacy
+
+import "github.com/ava-labs/libevm/core/vm"
+
+// PrecompiledStatefulContract is the legacy signature of
+// [vm.PrecompiledStatefulContract], which explicitly accepts and returns gas
+// values. Instances SHOULD NOT use the [vm.PrecompileEnvironment]
+// gas-management methods as this may result in unexpected behaviour.
+type PrecompiledStatefulContract func(env vm.PrecompileEnvironment, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error)
+
+// Upgrade converts the legacy precompile signature into the now-required form.
+func (c PrecompiledStatefulContract) Upgrade() vm.PrecompiledStatefulContract {
+	return func(env vm.PrecompileEnvironment, input []byte) ([]byte, error) {
+		gas := env.Gas()
+		ret, remainingGas, err := c(env, input, gas)
+		if used := gas - remainingGas; used < gas {
+			env.UseGas(used)
+		}
+		return ret, err
+	}
+}