Skip to content

Commit 6170c14

Browse files
ARR4NARR4N
committed
refactor: PrecompileEnvironment.UseGas() in lieu of return value
1 parent 594abd9 commit 6170c14

File tree

3 files changed

+65
-25
lines changed

3 files changed

+65
-25
lines changed

‎core/vm/contracts.libevm.go‎

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,8 @@ func (t CallType) OpCode() OpCode {
9292
// regular types.
9393
func (args *evmCallArgs) run(p PrecompiledContract, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error) {
9494
if p, ok := p.(statefulPrecompile); ok {
95-
return p(args.env(), input, suppliedGas)
95+
// `suppliedGas` is already held by the args.
96+
return p.run(args.env(), input)
9697
}
9798
// Gas consumption for regular precompiles was already handled by the native
9899
// RunPrecompiledContract(), which called this method.
@@ -102,7 +103,11 @@ func (args *evmCallArgs) run(p PrecompiledContract, input []byte, suppliedGas ui
102103

103104
// PrecompiledStatefulContract is the stateful equivalent of a
104105
// [PrecompiledContract].
105-
type PrecompiledStatefulContract func(env PrecompileEnvironment, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error)
106+
//
107+
// Instead of receiving and returning gas limits, stateful precompiles use the
108+
// respective methods on [PrecompileEnvironment]. If a call to UseGas() returns
109+
// false, a stateful precompile SHOULD return [ErrOutOfGas].
110+
type PrecompiledStatefulContract func(env PrecompileEnvironment, input []byte) (ret []byte, err error)
106111

107112
// NewStatefulPrecompile constructs a new PrecompiledContract that can be used
108113
// via an [EVM] instance but MUST NOT be called directly; a direct call to Run()
@@ -118,6 +123,11 @@ func NewStatefulPrecompile(run PrecompiledStatefulContract) PrecompiledContract
118123
// [PrecompiledStatefulContract] to hide implementation details.
119124
type statefulPrecompile PrecompiledStatefulContract
120125

126+
func (p statefulPrecompile) run(env *environment, input []byte) ([]byte, uint64, error) {
127+
ret, err := p(env, input)
128+
return ret, env.self.Gas, err
129+
}
130+
121131
// RequiredGas always returns zero as this gas is consumed by native geth code
122132
// before the contract is run.
123133
func (statefulPrecompile) RequiredGas([]byte) uint64 { return 0 }
@@ -135,13 +145,18 @@ func (p statefulPrecompile) Run([]byte) ([]byte, error) {
135145
type PrecompileEnvironment interface {
136146
ChainConfig() *params.ChainConfig
137147
Rules() params.Rules
138-
ReadOnly() bool
139148
// StateDB will be non-nil i.f.f !ReadOnly().
140149
StateDB() StateDB
141150
// ReadOnlyState will always be non-nil.
142151
ReadOnlyState() libevm.StateReader
143-
Addresses() *libevm.AddressContext
152+
144153
IncomingCallType() CallType
154+
Addresses() *libevm.AddressContext
155+
ReadOnly() bool
156+
// Equivalent to respective methods on [Contract].
157+
Gas() uint64
158+
UseGas(uint64) bool
159+
Value() *uint256.Int
145160

146161
BlockHeader() (types.Header, error)
147162
BlockNumber() *big.Int
@@ -150,7 +165,7 @@ type PrecompileEnvironment interface {
150165
// Call is equivalent to [EVM.Call] except that the `caller` argument is
151166
// removed and automatically determined according to the type of call that
152167
// invoked the precompile.
153-
Call(addr common.Address, input []byte, gas uint64, value *uint256.Int, _ ...CallOption) (ret []byte, gasRemaining uint64, _ error)
168+
Call(addr common.Address, input []byte, gas uint64, value *uint256.Int, _ ...CallOption) (ret []byte, _ error)
154169
}
155170

156171
func (args *evmCallArgs) env() *environment {

‎core/vm/contracts.libevm_test.go‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -146,13 +146,16 @@ func TestNewStatefulPrecompile(t *testing.T) {
146146
const gasLimit = 1e6
147147
gasCost := rng.Uint64n(gasLimit)
148148

149-
run := func(env vm.PrecompileEnvironment, input []byte, suppliedGas uint64) ([]byte, uint64, error) {
149+
run := func(env vm.PrecompileEnvironment, input []byte) ([]byte, error) {
150150
if got, want := env.StateDB() != nil, !env.ReadOnly(); got != want {
151-
return nil, 0, fmt.Errorf("PrecompileEnvironment().StateDB() must be non-nil i.f.f. not read-only; got non-nil? %t; want %t", got, want)
151+
return nil, fmt.Errorf("PrecompileEnvironment().StateDB() must be non-nil i.f.f. not read-only; got non-nil? %t; want %t", got, want)
152152
}
153153
hdr, err := env.BlockHeader()
154154
if err != nil {
155-
return nil, 0, err
155+
return nil, err
156+
}
157+
if !env.UseGas(gasCost) {
158+
return nil, vm.ErrOutOfGas
156159
}
157160

158161
out := &statefulPrecompileOutput{
@@ -166,7 +169,7 @@ func TestNewStatefulPrecompile(t *testing.T) {
166169
Input: input,
167170
IncomingCallType: env.IncomingCallType(),
168171
}
169-
return out.Bytes(), suppliedGas - gasCost, nil
172+
return out.Bytes(), nil
170173
}
171174
hooks := &hookstest.Stub{
172175
PrecompileOverrides: map[common.Address]libevm.PrecompiledContract{
@@ -318,11 +321,11 @@ func TestInheritReadOnly(t *testing.T) {
318321
hooks := &hookstest.Stub{
319322
PrecompileOverrides: map[common.Address]libevm.PrecompiledContract{
320323
precompile: vm.NewStatefulPrecompile(
321-
func(env vm.PrecompileEnvironment, input []byte, suppliedGas uint64) ([]byte, uint64, error) {
324+
func(env vm.PrecompileEnvironment, input []byte) ([]byte, error) {
322325
if env.ReadOnly() {
323-
return []byte{ifReadOnly}, suppliedGas, nil
326+
return []byte{ifReadOnly}, nil
324327
}
325-
return []byte{ifNotReadOnly}, suppliedGas, nil
328+
return []byte{ifNotReadOnly}, nil
326329
},
327330
),
328331
},
@@ -535,21 +538,21 @@ func TestPrecompileMakeCall(t *testing.T) {
535538

536539
hooks := &hookstest.Stub{
537540
PrecompileOverrides: map[common.Address]libevm.PrecompiledContract{
538-
sut: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error) {
541+
sut: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte) (ret []byte, err error) {
539542
var opts []vm.CallOption
540543
if bytes.Equal(input, unsafeCallerProxyOptSentinel) {
541544
opts = append(opts, vm.WithUNSAFECallerAddressProxying())
542545
}
543546
// We are ultimately testing env.Call(), hence why this is the SUT.
544-
return env.Call(dest, precompileCallData, suppliedGas, uint256.NewInt(0), opts...)
547+
return env.Call(dest, precompileCallData, env.Gas(), uint256.NewInt(0), opts...)
545548
}),
546-
dest: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error) {
549+
dest: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte) (ret []byte, err error) {
547550
out := &statefulPrecompileOutput{
548551
Addresses: env.Addresses(),
549552
ReadOnly: env.ReadOnly(),
550553
Input: input, // expected to be callData
551554
}
552-
return out.Bytes(), suppliedGas, nil
555+
return out.Bytes(), nil
553556
}),
554557
},
555558
}
@@ -696,8 +699,8 @@ func TestPrecompileCallWithTracer(t *testing.T) {
696699

697700
hooks := &hookstest.Stub{
698701
PrecompileOverrides: map[common.Address]libevm.PrecompiledContract{
699-
precompile: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error) {
700-
return env.Call(contract, nil, suppliedGas, uint256.NewInt(0))
702+
precompile: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte) (ret []byte, err error) {
703+
return env.Call(contract, nil, env.Gas(), uint256.NewInt(0))
701704
}),
702705
},
703706
}

‎core/vm/environment.libevm.go‎

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import (
2323
"github.com/holiman/uint256"
2424

2525
"github.com/ava-labs/libevm/common"
26+
"github.com/ava-labs/libevm/common/math"
2627
"github.com/ava-labs/libevm/core/types"
2728
"github.com/ava-labs/libevm/libevm"
2829
"github.com/ava-labs/libevm/params"
@@ -36,13 +37,26 @@ type environment struct {
3637
callType CallType
3738
}
3839

40+
func (e *environment) Gas() uint64 { return e.self.Gas }
41+
func (e *environment) UseGas(gas uint64) bool { return e.self.UseGas(gas) }
42+
func (e *environment) Value() *uint256.Int { return e.self.Value() }
43+
3944
func (e *environment) ChainConfig() *params.ChainConfig { return e.evm.chainConfig }
4045
func (e *environment) Rules() params.Rules { return e.evm.chainRules }
4146
func (e *environment) ReadOnlyState() libevm.StateReader { return e.evm.StateDB }
4247
func (e *environment) IncomingCallType() CallType { return e.callType }
4348
func (e *environment) BlockNumber() *big.Int { return new(big.Int).Set(e.evm.Context.BlockNumber) }
4449
func (e *environment) BlockTime() uint64 { return e.evm.Context.Time }
4550

51+
func (e *environment) refundGas(add uint64) error {
52+
gas, overflow := math.SafeAdd(e.self.Gas, add)
53+
if overflow {
54+
return ErrGasUintOverflow
55+
}
56+
e.self.Gas = gas
57+
return nil
58+
}
59+
4660
func (e *environment) ReadOnly() bool {
4761
// A switch statement provides clearer code coverage for difficult-to-test
4862
// cases.
@@ -86,11 +100,11 @@ func (e *environment) BlockHeader() (types.Header, error) {
86100
return *hdr, nil
87101
}
88102

89-
func (e *environment) Call(addr common.Address, input []byte, gas uint64, value *uint256.Int, opts ...CallOption) ([]byte, uint64, error) {
103+
func (e *environment) Call(addr common.Address, input []byte, gas uint64, value *uint256.Int, opts ...CallOption) ([]byte, error) {
90104
return e.callContract(Call, addr, input, gas, value, opts...)
91105
}
92106

93-
func (e *environment) callContract(typ CallType, addr common.Address, input []byte, gas uint64, value *uint256.Int, opts ...CallOption) (retData []byte, retGas uint64, retErr error) {
107+
func (e *environment) callContract(typ CallType, addr common.Address, input []byte, gas uint64, value *uint256.Int, opts ...CallOption) (retData []byte, retErr error) {
94108
// Depth and read-only setting are handled by [EVMInterpreter.Run], which
95109
// isn't used for precompiles, so we need to do it ourselves to maintain the
96110
// expected invariants.
@@ -118,13 +132,17 @@ func (e *environment) callContract(typ CallType, addr common.Address, input []by
118132
}
119133
case nil:
120134
default:
121-
return nil, gas, fmt.Errorf("unsupported option %T", o)
135+
return nil, fmt.Errorf("unsupported option %T", o)
122136
}
123137
}
124138

125139
if in.readOnly && value != nil && !value.IsZero() {
126-
return nil, gas, ErrWriteProtection
140+
return nil, ErrWriteProtection
127141
}
142+
if !e.UseGas(gas) {
143+
return nil, ErrOutOfGas
144+
}
145+
128146
if t := e.evm.Config.Tracer; t != nil {
129147
var bigVal *big.Int
130148
if value != nil {
@@ -134,13 +152,17 @@ func (e *environment) callContract(typ CallType, addr common.Address, input []by
134152

135153
startGas := gas
136154
defer func() {
137-
t.CaptureEnd(retData, startGas-retGas, retErr)
155+
t.CaptureEnd(retData, startGas-e.Gas(), retErr)
138156
}()
139157
}
140158

141159
switch typ {
142160
case Call:
143-
return e.evm.Call(caller, addr, input, gas, value)
161+
ret, returnGas, err := e.evm.Call(caller, addr, input, gas, value)
162+
if err := e.refundGas(returnGas); err != nil {
163+
return nil, err
164+
}
165+
return ret, err
144166
case CallCode, DelegateCall, StaticCall:
145167
// TODO(arr4n): these cases should be very similar to CALL, hence the
146168
// early abstraction, to signal to future maintainers. If implementing
@@ -149,6 +171,6 @@ func (e *environment) callContract(typ CallType, addr common.Address, input []by
149171
// compatibility.
150172
fallthrough
151173
default:
152-
return nil, gas, fmt.Errorf("unimplemented precompile call type %v", typ)
174+
return nil, fmt.Errorf("unimplemented precompile call type %v", typ)
153175
}
154176
}

0 commit comments

Comments
 (0)