Skip to content

Commit 46cd527

Browse files
cloudgraycloudgray
authored
Merge pull request #1 from zsystm/poc/precompiles-balance-handler
refactoring: balance handling
2 parents b532fd5 + 7c5ff7b commit 46cd527

File tree

9 files changed

+247
-398
lines changed

9 files changed

+247
-398
lines changed

‎precompiles/common/precompile.go

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -231,3 +231,41 @@ func (p *Precompile) GetBalanceHandler() *BalanceHandler {
231231
}
232232
return p.balanceHandler
233233
}
234+
235+
// ExecuteWithBalanceHandling wraps the execution of a precompile method with
236+
// the necessary balance handling and gas accounting logic. Precompile
237+
// implementations can use this helper to avoid repetitive boilerplate in their
238+
// `Run` methods.
239+
func (p Precompile) ExecuteWithBalanceHandling(
240+
evm *vm.EVM,
241+
contract *vm.Contract,
242+
readOnly bool,
243+
isTransaction func(*abi.Method) bool,
244+
handle func(ctx sdk.Context, contract *vm.Contract, stateDB *statedb.StateDB, method *abi.Method, args []interface{}) ([]byte, error),
245+
) (bz []byte, err error) {
246+
ctx, stateDB, method, initialGas, args, err := p.RunSetup(evm, contract, readOnly, isTransaction)
247+
if err != nil {
248+
return nil, err
249+
}
250+
251+
p.GetBalanceHandler().BeforeBalanceChange(ctx)
252+
253+
defer HandleGasError(ctx, contract, initialGas, &err)()
254+
255+
bz, err = handle(ctx, contract, stateDB, method, args)
256+
if err != nil {
257+
return nil, err
258+
}
259+
260+
cost := ctx.GasMeter().GasConsumed() - initialGas
261+
262+
if !contract.UseGas(cost, nil, tracing.GasChangeCallPrecompiledContract) {
263+
return nil, vm.ErrOutOfGas
264+
}
265+
266+
if err = p.GetBalanceHandler().AfterBalanceChange(ctx, stateDB); err != nil {
267+
return nil, err
268+
}
269+
270+
return bz, nil
271+
}

‎precompiles/distribution/distribution.go

Lines changed: 44 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,11 @@ package distribution
33
import (
44
"embed"
55
"fmt"
6+
sdk "github.com/cosmos/cosmos-sdk/types"
7+
"github.com/cosmos/evm/x/vm/statedb"
68

79
"github.com/ethereum/go-ethereum/accounts/abi"
810
"github.com/ethereum/go-ethereum/common"
9-
"github.com/ethereum/go-ethereum/core/tracing"
1011
"github.com/ethereum/go-ethereum/core/vm"
1112

1213
cmn "github.com/cosmos/evm/precompiles/common"
@@ -84,71 +85,48 @@ func (p Precompile) RequiredGas(input []byte) uint64 {
8485
}
8586

8687
// Run executes the precompiled contract distribution methods defined in the ABI.
87-
func (p Precompile) Run(evm *vm.EVM, contract *vm.Contract, readOnly bool) (bz []byte, err error) {
88-
ctx, stateDB, method, initialGas, args, err := p.RunSetup(evm, contract, readOnly, p.IsTransaction)
89-
if err != nil {
90-
return nil, err
91-
}
92-
93-
// Start the balance change handler before executing the precompile.
94-
p.GetBalanceHandler().BeforeBalanceChange(ctx)
95-
96-
// This handles any out of gas errors that may occur during the execution of a precompile tx or query.
97-
// It avoids panics and returns the out of gas error so the EVM can continue gracefully.
98-
defer cmn.HandleGasError(ctx, contract, initialGas, &err)()
99-
100-
switch method.Name {
101-
// Custom transactions
102-
case ClaimRewardsMethod:
103-
bz, err = p.ClaimRewards(ctx, contract, stateDB, method, args)
104-
// Distribution transactions
105-
case SetWithdrawAddressMethod:
106-
bz, err = p.SetWithdrawAddress(ctx, contract, stateDB, method, args)
107-
case WithdrawDelegatorRewardMethod:
108-
bz, err = p.WithdrawDelegatorReward(ctx, contract, stateDB, method, args)
109-
case WithdrawValidatorCommissionMethod:
110-
bz, err = p.WithdrawValidatorCommission(ctx, contract, stateDB, method, args)
111-
case FundCommunityPoolMethod:
112-
bz, err = p.FundCommunityPool(ctx, contract, stateDB, method, args)
113-
case DepositValidatorRewardsPoolMethod:
114-
bz, err = p.DepositValidatorRewardsPool(ctx, contract, stateDB, method, args)
115-
// Distribution queries
116-
case ValidatorDistributionInfoMethod:
117-
bz, err = p.ValidatorDistributionInfo(ctx, contract, method, args)
118-
case ValidatorOutstandingRewardsMethod:
119-
bz, err = p.ValidatorOutstandingRewards(ctx, contract, method, args)
120-
case ValidatorCommissionMethod:
121-
bz, err = p.ValidatorCommission(ctx, contract, method, args)
122-
case ValidatorSlashesMethod:
123-
bz, err = p.ValidatorSlashes(ctx, contract, method, args)
124-
case DelegationRewardsMethod:
125-
bz, err = p.DelegationRewards(ctx, contract, method, args)
126-
case DelegationTotalRewardsMethod:
127-
bz, err = p.DelegationTotalRewards(ctx, contract, method, args)
128-
case DelegatorValidatorsMethod:
129-
bz, err = p.DelegatorValidators(ctx, contract, method, args)
130-
case DelegatorWithdrawAddressMethod:
131-
bz, err = p.DelegatorWithdrawAddress(ctx, contract, method, args)
132-
case CommunityPoolMethod:
133-
bz, err = p.CommunityPool(ctx, contract, method, args)
134-
}
135-
136-
if err != nil {
137-
return nil, err
138-
}
139-
140-
cost := ctx.GasMeter().GasConsumed() - initialGas
141-
142-
if !contract.UseGas(cost, nil, tracing.GasChangeCallPrecompiledContract) {
143-
return nil, vm.ErrOutOfGas
144-
}
145-
146-
// Process the native balance changes after the method execution.
147-
if err = p.GetBalanceHandler().AfterBalanceChange(ctx, stateDB); err != nil {
148-
return nil, err
149-
}
150-
151-
return bz, nil
88+
func (p Precompile) Run(evm *vm.EVM, contract *vm.Contract, readOnly bool) ([]byte, error) {
89+
return p.ExecuteWithBalanceHandling(
90+
evm, contract, readOnly, p.IsTransaction,
91+
func(ctx sdk.Context, contract *vm.Contract, stateDB *statedb.StateDB, method *abi.Method, args []interface{}) ([]byte, error) {
92+
switch method.Name {
93+
// Custom transactions
94+
case ClaimRewardsMethod:
95+
return p.ClaimRewards(ctx, contract, stateDB, method, args)
96+
// Distribution transactions
97+
case SetWithdrawAddressMethod:
98+
return p.SetWithdrawAddress(ctx, contract, stateDB, method, args)
99+
case WithdrawDelegatorRewardMethod:
100+
return p.WithdrawDelegatorReward(ctx, contract, stateDB, method, args)
101+
case WithdrawValidatorCommissionMethod:
102+
return p.WithdrawValidatorCommission(ctx, contract, stateDB, method, args)
103+
case FundCommunityPoolMethod:
104+
return p.FundCommunityPool(ctx, contract, stateDB, method, args)
105+
case DepositValidatorRewardsPoolMethod:
106+
return p.DepositValidatorRewardsPool(ctx, contract, stateDB, method, args)
107+
// Distribution queries
108+
case ValidatorDistributionInfoMethod:
109+
return p.ValidatorDistributionInfo(ctx, contract, method, args)
110+
case ValidatorOutstandingRewardsMethod:
111+
return p.ValidatorOutstandingRewards(ctx, contract, method, args)
112+
case ValidatorCommissionMethod:
113+
return p.ValidatorCommission(ctx, contract, method, args)
114+
case ValidatorSlashesMethod:
115+
return p.ValidatorSlashes(ctx, contract, method, args)
116+
case DelegationRewardsMethod:
117+
return p.DelegationRewards(ctx, contract, method, args)
118+
case DelegationTotalRewardsMethod:
119+
return p.DelegationTotalRewards(ctx, contract, method, args)
120+
case DelegatorValidatorsMethod:
121+
return p.DelegatorValidators(ctx, contract, method, args)
122+
case DelegatorWithdrawAddressMethod:
123+
return p.DelegatorWithdrawAddress(ctx, contract, method, args)
124+
case CommunityPoolMethod:
125+
return p.CommunityPool(ctx, contract, method, args)
126+
}
127+
return nil, nil
128+
},
129+
)
152130
}
153131

154132
// IsTransaction checks if the given method name corresponds to a transaction or query.

‎precompiles/erc20/erc20.go

Lines changed: 8 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,9 @@ package erc20
33
import (
44
"embed"
55
"fmt"
6+
"github.com/cosmos/evm/x/vm/statedb"
67

78
"github.com/ethereum/go-ethereum/accounts/abi"
8-
"github.com/ethereum/go-ethereum/core/tracing"
99
"github.com/ethereum/go-ethereum/core/vm"
1010

1111
cmn "github.com/cosmos/evm/precompiles/common"
@@ -135,7 +135,7 @@ func (p Precompile) RequiredGas(input []byte) uint64 {
135135
}
136136

137137
// Run executes the precompiled contract ERC-20 methods defined in the ABI.
138-
func (p Precompile) Run(evm *vm.EVM, contract *vm.Contract, readOnly bool) (bz []byte, err error) {
138+
func (p Precompile) Run(evm *vm.EVM, contract *vm.Contract, readOnly bool) ([]byte, error) {
139139
// ERC20 precompiles cannot receive funds because they are not managed by an
140140
// EOA and will not be possible to recover funds sent to an instance of
141141
// them.This check is a safety measure because currently funds cannot be
@@ -144,36 +144,12 @@ func (p Precompile) Run(evm *vm.EVM, contract *vm.Contract, readOnly bool) (bz [
144144
return nil, fmt.Errorf(ErrCannotReceiveFunds, contract.Value().String())
145145
}
146146

147-
ctx, stateDB, method, initialGas, args, err := p.RunSetup(evm, contract, readOnly, p.IsTransaction)
148-
if err != nil {
149-
return nil, err
150-
}
151-
152-
// Start the balance change handler before executing the precompile.
153-
p.GetBalanceHandler().BeforeBalanceChange(ctx)
154-
155-
// This handles any out of gas errors that may occur during the execution of a precompile tx or query.
156-
// It avoids panics and returns the out of gas error so the EVM can continue gracefully.
157-
defer cmn.HandleGasError(ctx, contract, initialGas, &err)()
158-
159-
bz, err = p.HandleMethod(ctx, contract, stateDB, method, args)
160-
if err != nil {
161-
return nil, err
162-
}
163-
164-
cost := ctx.GasMeter().GasConsumed() - initialGas
165-
166-
if !contract.UseGas(cost, nil, tracing.GasChangeCallPrecompiledContract) {
167-
return nil, vm.ErrOutOfGas
168-
}
169-
170-
// Process the native balance changes after the method execution.
171-
err = p.GetBalanceHandler().AfterBalanceChange(ctx, stateDB)
172-
if err != nil {
173-
return nil, err
174-
}
175-
176-
return bz, nil
147+
return p.ExecuteWithBalanceHandling(
148+
evm, contract, readOnly, p.IsTransaction,
149+
func(ctx sdk.Context, contract *vm.Contract, stateDB *statedb.StateDB, method *abi.Method, args []interface{}) ([]byte, error) {
150+
return p.HandleMethod(ctx, contract, stateDB, method, args)
151+
},
152+
)
177153
}
178154

179155
// IsTransaction checks if the given method name corresponds to a transaction or query.

‎precompiles/evidence/evidence.go

Lines changed: 19 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,10 @@ package evidence
33
import (
44
"embed"
55
"fmt"
6+
"github.com/cosmos/evm/x/vm/statedb"
67

78
"github.com/ethereum/go-ethereum/accounts/abi"
89
"github.com/ethereum/go-ethereum/common"
9-
"github.com/ethereum/go-ethereum/core/tracing"
1010
"github.com/ethereum/go-ethereum/core/vm"
1111

1212
cmn "github.com/cosmos/evm/precompiles/common"
@@ -81,48 +81,24 @@ func (p Precompile) RequiredGas(input []byte) uint64 {
8181
}
8282

8383
// Run executes the precompiled contract evidence methods defined in the ABI.
84-
func (p Precompile) Run(evm *vm.EVM, contract *vm.Contract, readOnly bool) (bz []byte, err error) {
85-
ctx, stateDB, method, initialGas, args, err := p.RunSetup(evm, contract, readOnly, p.IsTransaction)
86-
if err != nil {
87-
return nil, err
88-
}
89-
90-
// Start the balance change handler before executing the precompile.
91-
p.GetBalanceHandler().BeforeBalanceChange(ctx)
92-
93-
// This handles any out of gas errors that may occur during the execution of a precompile tx or query.
94-
// It avoids panics and returns the out of gas error so the EVM can continue gracefully.
95-
defer cmn.HandleGasError(ctx, contract, initialGas, &err)()
96-
97-
switch method.Name {
98-
// evidence transactions
99-
case SubmitEvidenceMethod:
100-
bz, err = p.SubmitEvidence(ctx, contract, stateDB, method, args)
101-
// evidence queries
102-
case EvidenceMethod:
103-
bz, err = p.Evidence(ctx, method, args)
104-
case GetAllEvidenceMethod:
105-
bz, err = p.GetAllEvidence(ctx, method, args)
106-
default:
107-
return nil, fmt.Errorf(cmn.ErrUnknownMethod, method.Name)
108-
}
109-
110-
if err != nil {
111-
return nil, err
112-
}
113-
114-
cost := ctx.GasMeter().GasConsumed() - initialGas
115-
116-
if !contract.UseGas(cost, nil, tracing.GasChangeCallPrecompiledContract) {
117-
return nil, vm.ErrOutOfGas
118-
}
119-
120-
// Process the native balance changes after the method execution.
121-
if err = p.GetBalanceHandler().AfterBalanceChange(ctx, stateDB); err != nil {
122-
return nil, err
123-
}
124-
125-
return bz, nil
84+
func (p Precompile) Run(evm *vm.EVM, contract *vm.Contract, readOnly bool) ([]byte, error) {
85+
return p.ExecuteWithBalanceHandling(
86+
evm, contract, readOnly, p.IsTransaction,
87+
func(ctx sdk.Context, contract *vm.Contract, stateDB *statedb.StateDB, method *abi.Method, args []interface{}) ([]byte, error) {
88+
switch method.Name {
89+
// evidence transactions
90+
case SubmitEvidenceMethod:
91+
return p.SubmitEvidence(ctx, contract, stateDB, method, args)
92+
// evidence queries
93+
case EvidenceMethod:
94+
return p.Evidence(ctx, method, args)
95+
case GetAllEvidenceMethod:
96+
return p.GetAllEvidence(ctx, method, args)
97+
default:
98+
return nil, fmt.Errorf(cmn.ErrUnknownMethod, method.Name)
99+
}
100+
},
101+
)
126102
}
127103

128104
// IsTransaction checks if the given method name corresponds to a transaction or query.

0 commit comments

Comments
 (0)