refactor(precompiles): apply journal-based revert approach (#205)

* refactor(precompiles): apply journal-based revert approach

* refactor: remove unused Snapshot type

* chore: fix lint

Author: cloudgray

precompiles/bank/bank.go

@@ -104,47 +104,41 @@ func (p Precompile) RequiredGas(input []byte) uint64 {
 
 // Run executes the precompiled contract bank query methods defined in the ABI.
 func (p Precompile) Run(evm *vm.EVM, contract *vm.Contract, readOnly bool) (bz []byte, err error) {
-	ctx, stateDB, snapshot, method, initialGas, args, err := p.RunSetup(evm, contract, readOnly, p.IsTransaction)
+	ctx, stateDB, method, initialGas, args, err := p.RunSetup(evm, contract, readOnly, p.IsTransaction)
 	if err != nil {
 		return nil, err
 	}
 
 	// This handles any out of gas errors that may occur during the execution of a precompile query.
 	// It avoids panics and returns the out of gas error so the EVM can continue gracefully.
-	defer cmn.HandleGasError(ctx, contract, initialGas, &err, stateDB, snapshot)()
-
-	return p.RunAtomic(
-		snapshot,
-		stateDB,
-		func() ([]byte, error) {
-			switch method.Name {
-			// Bank queries
-			case BalancesMethod:
-				bz, err = p.Balances(ctx, contract, method, args)
-			case TotalSupplyMethod:
-				bz, err = p.TotalSupply(ctx, contract, method, args)
-			case SupplyOfMethod:
-				bz, err = p.SupplyOf(ctx, contract, method, args)
-			default:
-				return nil, fmt.Errorf(cmn.ErrUnknownMethod, method.Name)
-			}
-
-			if err != nil {
-				return nil, err
-			}
-
-			cost := ctx.GasMeter().GasConsumed() - initialGas
-
-			if !contract.UseGas(cost, nil, tracing.GasChangeCallPrecompiledContract) {
-				return nil, vm.ErrOutOfGas
-			}
-			if err := p.AddJournalEntries(stateDB, snapshot); err != nil {
-				return nil, err
-			}
-
-			return bz, nil
-		},
-	)
+	defer cmn.HandleGasError(ctx, contract, initialGas, &err)()
+
+	switch method.Name {
+	// Bank queries
+	case BalancesMethod:
+		bz, err = p.Balances(ctx, contract, method, args)
+	case TotalSupplyMethod:
+		bz, err = p.TotalSupply(ctx, contract, method, args)
+	case SupplyOfMethod:
+		bz, err = p.SupplyOf(ctx, contract, method, args)
+	default:
+		return nil, fmt.Errorf(cmn.ErrUnknownMethod, method.Name)
+	}
+
+	if err != nil {
+		return nil, err
+	}
+
+	cost := ctx.GasMeter().GasConsumed() - initialGas
+
+	if !contract.UseGas(cost, nil, tracing.GasChangeCallPrecompiledContract) {
+		return nil, vm.ErrOutOfGas
+	}
+	if err = p.AddJournalEntries(stateDB); err != nil {
+		return nil, err
+	}
+
+	return bz, nil
 }
 
 // IsTransaction checks if the given method name corresponds to a transaction or query.

precompiles/common/precompile.go

@@ -45,14 +45,6 @@ func NewBalanceChangeEntry(acc common.Address, amt *uint256.Int, op Operation) B
 	return BalanceChangeEntry{acc, amt, op}
 }
 
-// Snapshot contains all state and events previous to the precompile call
-// This is needed to allow us to revert the changes
-// during the EVM execution
-type Snapshot struct {
-	MultiStore storetypes.CacheMultiStore
-	Events     sdk.Events
-}
-
 // RequiredGas calculates the base minimum required gas for a transaction or a query.
 // It uses the method ID to determine if the input is a transaction or a query and
 // uses the Cosmos SDK gas config flat cost and the flat per byte cost * len(argBz) to calculate the gas.
@@ -66,46 +58,41 @@ func (p Precompile) RequiredGas(input []byte, isTransaction bool) uint64 {
 	return p.KvGasConfig.ReadCostFlat + (p.KvGasConfig.ReadCostPerByte * uint64(len(argsBz)))
 }
 
-// RunAtomic is used within the Run function of each Precompile implementation.
-// It handles rolling back to the provided snapshot if an error is returned from the core precompile logic.
-// Note: This is only required for stateful precompiles.
-func (p Precompile) RunAtomic(s Snapshot, stateDB *statedb.StateDB, fn func() ([]byte, error)) ([]byte, error) {
-	bz, err := fn()
-	if err != nil {
-		// revert to snapshot on error
-		stateDB.RevertMultiStore(s.MultiStore, s.Events)
-	}
-	return bz, err
-}
-
 // RunSetup runs the initial setup required to run a transaction or a query.
 // It returns the sdk Context, EVM stateDB, ABI method, initial gas and calling arguments.
 func (p Precompile) RunSetup(
 	evm *vm.EVM,
 	contract *vm.Contract,
 	readOnly bool,
 	isTransaction func(name *abi.Method) bool,
-) (ctx sdk.Context, stateDB *statedb.StateDB, s Snapshot, method *abi.Method, gasConfig storetypes.Gas, args []interface{}, err error) {
+) (ctx sdk.Context, stateDB *statedb.StateDB, method *abi.Method, gasConfig storetypes.Gas, args []interface{}, err error) {
 	stateDB, ok := evm.StateDB.(*statedb.StateDB)
 	if !ok {
-		return sdk.Context{}, nil, s, nil, uint64(0), nil, errors.New(ErrNotRunInEvm)
+		return sdk.Context{}, nil, nil, uint64(0), nil, errors.New(ErrNotRunInEvm)
 	}
 
 	// get the stateDB cache ctx
 	ctx, err = stateDB.GetCacheContext()
 	if err != nil {
-		return sdk.Context{}, nil, s, nil, uint64(0), nil, err
+		return sdk.Context{}, nil, nil, uint64(0), nil, err
 	}
 
 	// take a snapshot of the current state before any changes
 	// to be able to revert the changes
-	s.MultiStore = stateDB.MultiStoreSnapshot()
-	s.Events = ctx.EventManager().Events()
+	snapshot := stateDB.MultiStoreSnapshot()
+	events := ctx.EventManager().Events()
+
+	// add precompileCall entry on the stateDB journal
+	// this allows to revert the changes within an evm txAdd commentMore actions
+	err = stateDB.AddPrecompileFn(p.Address(), snapshot, events)
+	if err != nil {
+		return sdk.Context{}, nil, nil, uint64(0), nil, err
+	}
 
 	// commit the current changes in the cache ctx
 	// to get the updated state for the precompile call
 	if err := stateDB.CommitWithCacheCtx(); err != nil {
-		return sdk.Context{}, nil, s, nil, uint64(0), nil, err
+		return sdk.Context{}, nil, nil, uint64(0), nil, err
 	}
 
 	// NOTE: This is a special case where the calling transaction does not specify a function name.
@@ -131,26 +118,26 @@ func (p Precompile) RunSetup(
 	}
 
 	if err != nil {
-		return sdk.Context{}, nil, s, nil, uint64(0), nil, err
+		return sdk.Context{}, nil, nil, uint64(0), nil, err
 	}
 
 	// return error if trying to write to state during a read-only call
 	if readOnly && isTransaction(method) {
-		return sdk.Context{}, nil, s, nil, uint64(0), nil, vm.ErrWriteProtection
+		return sdk.Context{}, nil, nil, uint64(0), nil, vm.ErrWriteProtection
 	}
 
 	// if the method type is `function` continue looking for arguments
 	if method.Type == abi.Function {
 		argsBz := contract.Input[4:]
 		args, err = method.Inputs.Unpack(argsBz)
 		if err != nil {
-			return sdk.Context{}, nil, s, nil, uint64(0), nil, err
+			return sdk.Context{}, nil, nil, uint64(0), nil, err
 		}
 	}
 
 	initialGas := ctx.GasMeter().GasConsumed()
 
-	defer HandleGasError(ctx, contract, initialGas, &err, stateDB, s)()
+	defer HandleGasError(ctx, contract, initialGas, &err)()
 
 	// set the default SDK gas configuration to track gas usage
 	// we are changing the gas meter type, so it panics gracefully when out of gas
@@ -160,20 +147,16 @@ func (p Precompile) RunSetup(
 	// we need to consume the gas that was already used by the EVM
 	ctx.GasMeter().ConsumeGas(initialGas, "creating a new gas meter")
 
-	return ctx, stateDB, s, method, initialGas, args, nil
+	return ctx, stateDB, method, initialGas, args, nil
 }
 
 // HandleGasError handles the out of gas panic by resetting the gas meter and returning an error.
 // This is used in order to avoid panics and to allow for the EVM to continue cleanup if the tx or query run out of gas.
-func HandleGasError(ctx sdk.Context, contract *vm.Contract, initialGas storetypes.Gas, err *error, stateDB *statedb.StateDB, snapshot Snapshot) func() {
+func HandleGasError(ctx sdk.Context, contract *vm.Contract, initialGas storetypes.Gas, err *error) func() {
 	return func() {
 		if r := recover(); r != nil {
 			switch r.(type) {
 			case storetypes.ErrorOutOfGas:
-
-				// revert to snapshot on error
-				stateDB.RevertMultiStore(snapshot.MultiStore, snapshot.Events)
-
 				// update contract gas
 				usedGas := ctx.GasMeter().GasConsumed() - initialGas
 				_ = contract.UseGas(usedGas, nil, tracing.GasChangeCallFailedExecution)
@@ -190,9 +173,7 @@ func HandleGasError(ctx sdk.Context, contract *vm.Contract, initialGas storetype
 }
 
 // AddJournalEntries adds the balanceChange (if corresponds)
-// and precompileCall entries on the stateDB journal
-// This allows to revert the call changes within an evm tx
-func (p Precompile) AddJournalEntries(stateDB *statedb.StateDB, s Snapshot) error {
+func (p Precompile) AddJournalEntries(stateDB *statedb.StateDB) error {
 	for _, entry := range p.journalEntries {
 		switch entry.Op {
 		case Sub:
@@ -204,7 +185,7 @@ func (p Precompile) AddJournalEntries(stateDB *statedb.StateDB, s Snapshot) erro
 		}
 	}
 
-	return stateDB.AddPrecompileFn(p.Address(), s.MultiStore, s.Events)
+	return nil
 }
 
 // SetBalanceChangeEntries sets the balanceChange entries

precompiles/distribution/distribution.go

@@ -85,68 +85,66 @@ func (p Precompile) RequiredGas(input []byte) uint64 {
 
 // Run executes the precompiled contract distribution methods defined in the ABI.
 func (p Precompile) Run(evm *vm.EVM, contract *vm.Contract, readOnly bool) (bz []byte, err error) {
-	ctx, stateDB, snapshot, method, initialGas, args, err := p.RunSetup(evm, contract, readOnly, p.IsTransaction)
+	ctx, stateDB, method, initialGas, args, err := p.RunSetup(evm, contract, readOnly, p.IsTransaction)
 	if err != nil {
 		return nil, err
 	}
 
 	// This handles any out of gas errors that may occur during the execution of a precompile tx or query.
 	// It avoids panics and returns the out of gas error so the EVM can continue gracefully.
-	defer cmn.HandleGasError(ctx, contract, initialGas, &err, stateDB, snapshot)()
-
-	return p.RunAtomic(snapshot, stateDB, func() ([]byte, error) {
-		switch method.Name {
-		// Custom transactions
-		case ClaimRewardsMethod:
-			bz, err = p.ClaimRewards(ctx, contract, stateDB, method, args)
-		// Distribution transactions
-		case SetWithdrawAddressMethod:
-			bz, err = p.SetWithdrawAddress(ctx, contract, stateDB, method, args)
-		case WithdrawDelegatorRewardMethod:
-			bz, err = p.WithdrawDelegatorReward(ctx, contract, stateDB, method, args)
-		case WithdrawValidatorCommissionMethod:
-			bz, err = p.WithdrawValidatorCommission(ctx, contract, stateDB, method, args)
-		case FundCommunityPoolMethod:
-			bz, err = p.FundCommunityPool(ctx, contract, stateDB, method, args)
-		case DepositValidatorRewardsPoolMethod:
-			bz, err = p.DepositValidatorRewardsPool(ctx, contract, stateDB, method, args)
-		// Distribution queries
-		case ValidatorDistributionInfoMethod:
-			bz, err = p.ValidatorDistributionInfo(ctx, contract, method, args)
-		case ValidatorOutstandingRewardsMethod:
-			bz, err = p.ValidatorOutstandingRewards(ctx, contract, method, args)
-		case ValidatorCommissionMethod:
-			bz, err = p.ValidatorCommission(ctx, contract, method, args)
-		case ValidatorSlashesMethod:
-			bz, err = p.ValidatorSlashes(ctx, contract, method, args)
-		case DelegationRewardsMethod:
-			bz, err = p.DelegationRewards(ctx, contract, method, args)
-		case DelegationTotalRewardsMethod:
-			bz, err = p.DelegationTotalRewards(ctx, contract, method, args)
-		case DelegatorValidatorsMethod:
-			bz, err = p.DelegatorValidators(ctx, contract, method, args)
-		case DelegatorWithdrawAddressMethod:
-			bz, err = p.DelegatorWithdrawAddress(ctx, contract, method, args)
-		case CommunityPoolMethod:
-			bz, err = p.CommunityPool(ctx, contract, method, args)
-		}
-
-		if err != nil {
-			return nil, err
-		}
-
-		cost := ctx.GasMeter().GasConsumed() - initialGas
-
-		if !contract.UseGas(cost, nil, tracing.GasChangeCallPrecompiledContract) {
-			return nil, vm.ErrOutOfGas
-		}
-
-		if err := p.AddJournalEntries(stateDB, snapshot); err != nil {
-			return nil, err
-		}
-
-		return bz, nil
-	})
+	defer cmn.HandleGasError(ctx, contract, initialGas, &err)()
+
+	switch method.Name {
+	// Custom transactions
+	case ClaimRewardsMethod:
+		bz, err = p.ClaimRewards(ctx, contract, stateDB, method, args)
+	// Distribution transactions
+	case SetWithdrawAddressMethod:
+		bz, err = p.SetWithdrawAddress(ctx, contract, stateDB, method, args)
+	case WithdrawDelegatorRewardMethod:
+		bz, err = p.WithdrawDelegatorReward(ctx, contract, stateDB, method, args)
+	case WithdrawValidatorCommissionMethod:
+		bz, err = p.WithdrawValidatorCommission(ctx, contract, stateDB, method, args)
+	case FundCommunityPoolMethod:
+		bz, err = p.FundCommunityPool(ctx, contract, stateDB, method, args)
+	case DepositValidatorRewardsPoolMethod:
+		bz, err = p.DepositValidatorRewardsPool(ctx, contract, stateDB, method, args)
+	// Distribution queries
+	case ValidatorDistributionInfoMethod:
+		bz, err = p.ValidatorDistributionInfo(ctx, contract, method, args)
+	case ValidatorOutstandingRewardsMethod:
+		bz, err = p.ValidatorOutstandingRewards(ctx, contract, method, args)
+	case ValidatorCommissionMethod:
+		bz, err = p.ValidatorCommission(ctx, contract, method, args)
+	case ValidatorSlashesMethod:
+		bz, err = p.ValidatorSlashes(ctx, contract, method, args)
+	case DelegationRewardsMethod:
+		bz, err = p.DelegationRewards(ctx, contract, method, args)
+	case DelegationTotalRewardsMethod:
+		bz, err = p.DelegationTotalRewards(ctx, contract, method, args)
+	case DelegatorValidatorsMethod:
+		bz, err = p.DelegatorValidators(ctx, contract, method, args)
+	case DelegatorWithdrawAddressMethod:
+		bz, err = p.DelegatorWithdrawAddress(ctx, contract, method, args)
+	case CommunityPoolMethod:
+		bz, err = p.CommunityPool(ctx, contract, method, args)
+	}
+
+	if err != nil {
+		return nil, err
+	}
+
+	cost := ctx.GasMeter().GasConsumed() - initialGas
+
+	if !contract.UseGas(cost, nil, tracing.GasChangeCallPrecompiledContract) {
+		return nil, vm.ErrOutOfGas
+	}
+
+	if err = p.AddJournalEntries(stateDB); err != nil {
+		return nil, err
+	}
+
+	return bz, nil
 }
 
 // IsTransaction checks if the given method name corresponds to a transaction or query.

precompiles/erc20/erc20.go

@@ -144,31 +144,29 @@ func (p Precompile) Run(evm *vm.EVM, contract *vm.Contract, readOnly bool) (bz [
 		return nil, fmt.Errorf(ErrCannotReceiveFunds, contract.Value().String())
 	}
 
-	ctx, stateDB, snapshot, method, initialGas, args, err := p.RunSetup(evm, contract, readOnly, p.IsTransaction)
+	ctx, stateDB, method, initialGas, args, err := p.RunSetup(evm, contract, readOnly, p.IsTransaction)
 	if err != nil {
 		return nil, err
 	}
 
 	// This handles any out of gas errors that may occur during the execution of a precompile tx or query.
 	// It avoids panics and returns the out of gas error so the EVM can continue gracefully.
-	defer cmn.HandleGasError(ctx, contract, initialGas, &err, stateDB, snapshot)()
-
-	return p.RunAtomic(snapshot, stateDB, func() ([]byte, error) {
-		bz, err = p.HandleMethod(ctx, contract, stateDB, method, args)
-		if err != nil {
-			return nil, err
-		}
-
-		cost := ctx.GasMeter().GasConsumed() - initialGas
-
-		if !contract.UseGas(cost, nil, tracing.GasChangeCallPrecompiledContract) {
-			return nil, vm.ErrOutOfGas
-		}
-		if err := p.AddJournalEntries(stateDB, snapshot); err != nil {
-			return nil, err
-		}
-		return bz, nil
-	})
+	defer cmn.HandleGasError(ctx, contract, initialGas, &err)()
+
+	bz, err = p.HandleMethod(ctx, contract, stateDB, method, args)
+	if err != nil {
+		return nil, err
+	}
+
+	cost := ctx.GasMeter().GasConsumed() - initialGas
+
+	if !contract.UseGas(cost, nil, tracing.GasChangeCallPrecompiledContract) {
+		return nil, vm.ErrOutOfGas
+	}
+	if err = p.AddJournalEntries(stateDB); err != nil {
+		return nil, err
+	}
+	return bz, nil
 }
 
 // IsTransaction checks if the given method name corresponds to a transaction or query.