[api] fix estimateGas for migrateStake (#4377)

iotexproject · Sep 3, 2024 · ef0af12 · ef0af12
1 parent d2f813b
commit ef0af12
Show file tree

Hide file tree

Showing 7 changed files with 97 additions and 20 deletions.
diff --git a/action/protocol/managers.go b/action/protocol/managers.go
@@ -93,3 +93,16 @@ type (
 		Reset()
 	}
 )
+
+type (
+	SimulateOption       func(*SimulateOptionConfig)
+	SimulateOptionConfig struct {
+		PreOpt func(StateManager) error
+	}
+)
+
+func WithSimulatePreOpt(fn func(StateManager) error) SimulateOption {
+	return func(so *SimulateOptionConfig) {
+		so.PreOpt = fn
+	}
+}
diff --git a/api/coreservice.go b/api/coreservice.go
@@ -138,7 +138,7 @@ type (
 		// EstimateGasForNonExecution  estimates action gas except execution
 		EstimateGasForNonExecution(action.Action) (uint64, error)
 		// EstimateExecutionGasConsumption estimate gas consumption for execution action
-		EstimateExecutionGasConsumption(ctx context.Context, sc *action.Execution, callerAddr address.Address) (uint64, error)
+		EstimateExecutionGasConsumption(ctx context.Context, sc *action.Execution, callerAddr address.Address, opts ...protocol.SimulateOption) (uint64, error)
 		// LogsInBlockByHash filter logs in the block by hash
 		LogsInBlockByHash(filter *logfilter.LogFilter, blockHash hash.Hash256) ([]*action.Log, error)
 		// LogsInRange filter logs among [start, end] blocks
@@ -1513,11 +1513,22 @@ func (core *coreService) EstimateMigrateStakeGasConsumption(ctx context.Context,
 		GasLimit:       g.BlockGasLimitByHeight(header.Height() + 1),
 		Producer:       zeroAddr,
 	})
+
 	exec, err := staking.FindProtocol(core.registry).ConstructExecution(ctx, ms, core.sf)
 	if err != nil {
 		return 0, err
 	}
-	gas, err := core.EstimateExecutionGasConsumption(ctx, exec, caller)
+	gas, err := core.EstimateExecutionGasConsumption(ctx, exec, caller, protocol.WithSimulatePreOpt(func(sm protocol.StateManager) error {
+		// add amount to the sender account
+		sender, err := accountutil.LoadAccount(sm, caller)
+		if err != nil {
+			return err
+		}
+		if err = sender.AddBalance(exec.Amount()); err != nil {
+			return err
+		}
+		return accountutil.StoreAccount(sm, caller, sender)
+	}))
 	if err != nil {
 		return 0, err
 	}
@@ -1529,7 +1540,7 @@ func (core *coreService) EstimateMigrateStakeGasConsumption(ctx context.Context,
 }
 
 // EstimateExecutionGasConsumption estimate gas consumption for execution action
-func (core *coreService) EstimateExecutionGasConsumption(ctx context.Context, sc *action.Execution, callerAddr address.Address) (uint64, error) {
+func (core *coreService) EstimateExecutionGasConsumption(ctx context.Context, sc *action.Execution, callerAddr address.Address, opts ...protocol.SimulateOption) (uint64, error) {
 	ctx = genesis.WithGenesisContext(ctx, core.bc.Genesis())
 	state, err := accountutil.AccountState(ctx, core.sf, callerAddr)
 	if err != nil {
@@ -1552,7 +1563,7 @@ func (core *coreService) EstimateExecutionGasConsumption(ctx context.Context, sc
 		blockGasLimit = g.BlockGasLimitByHeight(core.bc.TipHeight())
 	)
 	sc.SetGasLimit(blockGasLimit)
-	enough, receipt, err := core.isGasLimitEnough(ctx, callerAddr, sc)
+	enough, receipt, err := core.isGasLimitEnough(ctx, callerAddr, sc, opts...)
 	if err != nil {
 		return 0, status.Error(codes.Internal, err.Error())
 	}
@@ -1564,7 +1575,7 @@ func (core *coreService) EstimateExecutionGasConsumption(ctx context.Context, sc
 	}
 	estimatedGas := receipt.GasConsumed
 	sc.SetGasLimit(estimatedGas)
-	enough, _, err = core.isGasLimitEnough(ctx, callerAddr, sc)
+	enough, _, err = core.isGasLimitEnough(ctx, callerAddr, sc, opts...)
 	if err != nil && err != action.ErrInsufficientFunds {
 		return 0, status.Error(codes.Internal, err.Error())
 	}
@@ -1574,7 +1585,7 @@ func (core *coreService) EstimateExecutionGasConsumption(ctx context.Context, sc
 		for low <= high {
 			mid := (low + high) / 2
 			sc.SetGasLimit(mid)
-			enough, _, err = core.isGasLimitEnough(ctx, callerAddr, sc)
+			enough, _, err = core.isGasLimitEnough(ctx, callerAddr, sc, opts...)
 			if err != nil && err != action.ErrInsufficientFunds {
 				return 0, status.Error(codes.Internal, err.Error())
 			}
@@ -1594,6 +1605,7 @@ func (core *coreService) isGasLimitEnough(
 	ctx context.Context,
 	caller address.Address,
 	sc *action.Execution,
+	opts ...protocol.SimulateOption,
 ) (bool, *action.Receipt, error) {
 	ctx, span := tracer.NewSpan(ctx, "Server.isGasLimitEnough")
 	defer span.End()
@@ -1602,7 +1614,7 @@ func (core *coreService) isGasLimitEnough(
 		return false, nil, err
 	}
 
-	_, receipt, err := core.simulateExecution(ctx, caller, sc, core.dao.GetBlockHash, core.getBlockTime)
+	_, receipt, err := core.simulateExecution(ctx, caller, sc, core.dao.GetBlockHash, core.getBlockTime, opts...)
 	if err != nil {
 		return false, nil, err
 	}
@@ -1908,14 +1920,14 @@ func (core *coreService) traceTx(ctx context.Context, txctx *tracers.Context, co
 	return retval, receipt, tracer, err
 }
 
-func (core *coreService) simulateExecution(ctx context.Context, addr address.Address, exec *action.Execution, getBlockHash evm.GetBlockHash, getBlockTime evm.GetBlockTime) ([]byte, *action.Receipt, error) {
+func (core *coreService) simulateExecution(ctx context.Context, addr address.Address, exec *action.Execution, getBlockHash evm.GetBlockHash, getBlockTime evm.GetBlockTime, opts ...protocol.SimulateOption) ([]byte, *action.Receipt, error) {
 	ctx = evm.WithHelperCtx(ctx, evm.HelperContext{
 		GetBlockHash:   getBlockHash,
 		GetBlockTime:   getBlockTime,
 		DepositGasFunc: rewarding.DepositGasWithSGD,
 		Sgd:            core.sgdIndexer,
 	})
-	return core.sf.SimulateExecution(ctx, addr, exec)
+	return core.sf.SimulateExecution(ctx, addr, exec, opts...)
 }
 
 func filterReceipts(receipts []*action.Receipt, actHash hash.Hash256) *action.Receipt {

diff --git a/e2etest/native_staking_test.go b/e2etest/native_staking_test.go
@@ -851,6 +851,7 @@ func TestCandidateTransferOwnership(t *testing.T) {
 			return data
 		}
 		deployCode := append(bytecode, mustCallData("", minAmount)...)
+		poorID := 30
 		test.run([]*testcase{
 			{
 				name: "deploy staking contract",
@@ -1020,6 +1021,28 @@ func TestCandidateTransferOwnership(t *testing.T) {
 					})
 					require.NoError(err)
 					require.Equal(uint64(194912), resp.Gas)
+					require.Len(receipt.Logs(), 1)
+					topic := receipt.Logs()[0].Topics[1][:]
+					bktIdx := byteutil.BytesToUint64BigEndian(topic[len(topic)-8:])
+					require.Equal(uint64(6), bktIdx)
+				}}},
+			},
+			{
+				name: "estimateGasPoorAcc",
+				act:  &actionWithTime{mustNoErr(action.SignedTransferStake(test.nonceMgr.pop(identityset.Address(stakerID).String()), identityset.Address(poorID).String(), 6, nil, gasLimit, gasPrice, identityset.PrivateKey(stakerID), action.WithChainID(chainID))), time.Now()},
+				expect: []actionExpect{&functionExpect{func(test *e2etest, act *action.SealedEnvelope, receipt *action.Receipt, err error) {
+					resp1, err := test.api.GetAccount(context.Background(), &iotexapi.GetAccountRequest{Address: identityset.Address(poorID).String()})
+					require.NoError(err)
+					require.Equal("0", resp1.GetAccountMeta().Balance)
+					ms, err := action.NewMigrateStake(0, 6, gasLimit, gasPrice)
+					require.NoError(err)
+					resp, err := test.api.EstimateActionGasConsumption(context.Background(), &iotexapi.EstimateActionGasConsumptionRequest{
+						Action:        &iotexapi.EstimateActionGasConsumptionRequest_StakeMigrate{StakeMigrate: ms.Proto()},
+						CallerAddress: identityset.Address(poorID).String(),
+						GasPrice:      gasPrice.String(),
+					})
+					require.NoError(err)
+					require.Equal(uint64(194912), resp.Gas)
 				}}},
 			},
 		})

diff --git a/state/factory/factory.go b/state/factory/factory.go
@@ -86,7 +86,7 @@ type (
 		Validate(context.Context, *block.Block) error
 		// NewBlockBuilder creates block builder
 		NewBlockBuilder(context.Context, actpool.ActPool, func(action.Envelope) (*action.SealedEnvelope, error)) (*block.Builder, error)
-		SimulateExecution(context.Context, address.Address, *action.Execution) ([]byte, *action.Receipt, error)
+		SimulateExecution(context.Context, address.Address, *action.Execution, ...protocol.SimulateOption) ([]byte, *action.Receipt, error)
 		ReadContractStorage(context.Context, address.Address, []byte) ([]byte, error)
 		PutBlock(context.Context, *block.Block) error
 		DeleteTipBlock(context.Context, *block.Block) error
@@ -387,6 +387,7 @@ func (sf *factory) SimulateExecution(
 	ctx context.Context,
 	caller address.Address,
 	ex *action.Execution,
+	opts ...protocol.SimulateOption,
 ) ([]byte, *action.Receipt, error) {
 	ctx, span := tracer.NewSpan(ctx, "factory.SimulateExecution")
 	defer span.End()
@@ -397,7 +398,15 @@ func (sf *factory) SimulateExecution(
 	if err != nil {
 		return nil, nil, errors.Wrap(err, "failed to obtain working set from state factory")
 	}
-
+	cfg := &protocol.SimulateOptionConfig{}
+	for _, opt := range opts {
+		opt(cfg)
+	}
+	if cfg.PreOpt != nil {
+		if err := cfg.PreOpt(ws); err != nil {
+			return nil, nil, err
+		}
+	}
 	return evm.SimulateExecution(ctx, ws, caller, ex)
 }
 

diff --git a/state/factory/statedb.go b/state/factory/statedb.go
@@ -267,6 +267,7 @@ func (sdb *stateDB) SimulateExecution(
 	ctx context.Context,
 	caller address.Address,
 	ex *action.Execution,
+	opts ...protocol.SimulateOption,
 ) ([]byte, *action.Receipt, error) {
 	ctx, span := tracer.NewSpan(ctx, "stateDB.SimulateExecution")
 	defer span.End()
@@ -278,7 +279,15 @@ func (sdb *stateDB) SimulateExecution(
 	if err != nil {
 		return nil, nil, err
 	}
-
+	cfg := &protocol.SimulateOptionConfig{}
+	for _, opt := range opts {
+		opt(cfg)
+	}
+	if cfg.PreOpt != nil {
+		if err := cfg.PreOpt(ws); err != nil {
+			return nil, nil, err
+		}
+	}
 	return evm.SimulateExecution(ctx, ws, caller, ex)
 }
 

diff --git a/test/mock/mock_apicoreservice/mock_apicoreservice.go b/test/mock/mock_apicoreservice/mock_apicoreservice.go
diff --git a/test/mock/mock_factory/mock_factory.go b/test/mock/mock_factory/mock_factory.go