fix(baseapp): align block header when query with latest height

CHANGELOG.md

@@ -119,6 +119,7 @@ Every module contains its own CHANGELOG.md. Please refer to the module you are i
 * [#19851](https://github.com/cosmos/cosmos-sdk/pull/19851) Fix some places in which we call Remove inside a Walk (x/staking and x/gov).
 * [#20939](https://github.com/cosmos/cosmos-sdk/pull/20939) Fix collection reverse iterator to include `pagination.key` in the result.
 * (client/grpc) [#20969](https://github.com/cosmos/cosmos-sdk/pull/20969) Fix `node.NewQueryServer` method not setting `cfg`.
+* (baseapp) [#21003](https://github.com/cosmos/cosmos-sdk/pull/21003) Align block header when query with latest height.
 * (testutil/integration) [#21006](https://github.com/cosmos/cosmos-sdk/pull/21006) Fix `NewIntegrationApp` method not writing default genesis to state
 
 ### API Breaking Changes

baseapp/abci.go

@@ -139,7 +139,7 @@ func (app *BaseApp) Info(_ *abci.InfoRequest) (*abci.InfoResponse, error) {
 	lastCommitID := app.cms.LastCommitID()
 	appVersion := InitialAppVersion
 	if lastCommitID.Version > 0 {
-		ctx, err := app.CreateQueryContext(lastCommitID.Version, false)
+		ctx, err := app.CreateQueryContextWithCheckHeader(lastCommitID.Version, false, false)
 		if err != nil {
 			return nil, fmt.Errorf("failed creating query context: %w", err)
 		}
@@ -1200,7 +1200,13 @@ func checkNegativeHeight(height int64) error {
 // CreateQueryContext creates a new sdk.Context for a query, taking as args
 // the block height and whether the query needs a proof or not.
 func (app *BaseApp) CreateQueryContext(height int64, prove bool) (sdk.Context, error) {
-	if err := checkNegativeHeight(height); err != nil {
+	return app.CreateQueryContextWithCheckHeader(height, prove, true)
+}
+
+// CreateQueryContextWithCheckHeader creates a new sdk.Context for a query, taking as args
+// the block height, whether the query needs a proof or not, and whether to check the header or not.
+func (app *BaseApp) CreateQueryContextWithCheckHeader(height int64, prove, checkHeader bool) (ctx sdk.Context, err error) {
+	if err = checkNegativeHeight(height); err != nil {
 		return sdk.Context{}, err
 	}
 
@@ -1223,19 +1229,57 @@ func (app *BaseApp) CreateQueryContext(height int64, prove bool) (sdk.Context, e
 			)
 	}
 
-	// when a client did not provide a query height, manually inject the latest
-	if height == 0 {
-		height = lastBlockHeight
-	}
-
-	if height <= 1 && prove {
+	if height > 0 && height <= 1 && prove {
 		return sdk.Context{},
 			errorsmod.Wrap(
 				sdkerrors.ErrInvalidRequest,
 				"cannot query with proof when height <= 1; please provide a valid height",
 			)
 	}
 
+	var header *cmtproto.Header
+	for _, state := range []*state{
+		app.checkState,
+		app.finalizeBlockState,
+	} {
+		if state != nil {
+			// branch the commit multi-store for safety
+			h := state.Context().BlockHeader()
+			if height == 0 {
+				lastBlockHeight = qms.LatestVersion()
+			}
+			if !checkHeader || height != 0 || height == 0 && h.Height == lastBlockHeight {
+				header = &h
+				break
+			}
+		}
+	}
+
+	if header == nil {
+		return sdk.Context{},
+			errorsmod.Wrapf(
+				sdkerrors.ErrInvalidHeight,
+				"header height in all state context is not latest height (%d)", lastBlockHeight,
+			)
+	}
+
+	// when a client did not provide a query height, manually inject the latest
+	if height == 0 {
+		height = lastBlockHeight
+	} else {
+		defer func() {
+			if err == nil {
+				rms, ok := app.cms.(*rootmulti.Store)
+				if ok {
+					cInfo, err := rms.GetCommitInfo(height)
+					if cInfo != nil && err == nil {
+						ctx = ctx.WithHeaderInfo(coreheader.Info{Height: height, Time: cInfo.Timestamp})
+					}
+				}
+			}
+		}()
+	}
+
 	cacheMS, err := qms.CacheMultiStoreWithVersion(height)
 	if err != nil {
 		return sdk.Context{},
@@ -1245,27 +1289,15 @@ func (app *BaseApp) CreateQueryContext(height int64, prove bool) (sdk.Context, e
 			)
 	}
 
-	// branch the commit multi-store for safety
-	ctx := sdk.NewContext(cacheMS, true, app.logger).
+	ctx = sdk.NewContext(cacheMS, true, app.logger).
 		WithMinGasPrices(app.minGasPrices).
 		WithGasMeter(storetypes.NewGasMeter(app.queryGasLimit)).
 		WithHeaderInfo(coreheader.Info{
 			ChainID: app.chainID,
 			Height:  height,
 		}).
-		WithBlockHeader(app.checkState.Context().BlockHeader()).
+		WithBlockHeader(*header).
 		WithBlockHeight(height)
-
-	if height != lastBlockHeight {
-		rms, ok := app.cms.(*rootmulti.Store)
-		if ok {
-			cInfo, err := rms.GetCommitInfo(height)
-			if cInfo != nil && err == nil {
-				ctx = ctx.WithHeaderInfo(coreheader.Info{Height: height, Time: cInfo.Timestamp})
-			}
-		}
-	}
-
 	return ctx, nil
 }
 

baseapp/baseapp_test.go

@@ -688,26 +688,26 @@ func TestBaseAppPostHandler(t *testing.T) {
 	require.NotContains(t, suite.logBuffer.String(), "panic recovered in runTx")
 }
 
+type mockABCIListener struct {
+	ListenCommitFn func(context.Context, abci.CommitResponse, []*storetypes.StoreKVPair) error
+}
+
+func (m mockABCIListener) ListenFinalizeBlock(_ context.Context, _ abci.FinalizeBlockRequest, _ abci.FinalizeBlockResponse) error {
+	return nil
+}
+
+func (m *mockABCIListener) ListenCommit(ctx context.Context, commit abci.CommitResponse, pairs []*storetypes.StoreKVPair) error {
+	return m.ListenCommitFn(ctx, commit, pairs)
+}
+
 // Test and ensure that invalid block heights always cause errors.
 // See issues:
 // - https://github.com/cosmos/cosmos-sdk/issues/11220
 // - https://github.com/cosmos/cosmos-sdk/issues/7662
 func TestABCI_CreateQueryContext(t *testing.T) {
 	t.Parallel()
+	app := getQueryBaseapp(t)
 
-	db := dbm.NewMemDB()
-	name := t.Name()
-	app := baseapp.NewBaseApp(name, log.NewTestLogger(t), db, nil)
-
-	_, err := app.FinalizeBlock(&abci.FinalizeBlockRequest{Height: 1})
-	require.NoError(t, err)
-	_, err = app.Commit()
-	require.NoError(t, err)
-
-	_, err = app.FinalizeBlock(&abci.FinalizeBlockRequest{Height: 2})
-	require.NoError(t, err)
-	_, err = app.Commit()
-	require.NoError(t, err)
 	testCases := []struct {
 		name         string
 		height       int64
@@ -716,7 +716,7 @@ func TestABCI_CreateQueryContext(t *testing.T) {
 		expErr       bool
 	}{
 		{"valid height", 2, 2, true, false},
-		{"valid height with different initial height", 2, 1, true, false},
+		{"valid height with different initial height", 2, 1, true, true},
 		{"future height", 10, 10, true, true},
 		{"negative height, prove=true", -1, -1, true, true},
 		{"negative height, prove=false", -1, -1, false, true},
@@ -741,6 +741,81 @@ func TestABCI_CreateQueryContext(t *testing.T) {
 	}
 }
 
+func TestABCI_CreateQueryContextWithCheckHeader(t *testing.T) {
+	t.Parallel()
+	app := getQueryBaseapp(t)
+	var height int64 = 2
+	var headerHeight int64 = 1
+
+	testCases := []struct {
+		checkHeader bool
+		expErr      bool
+	}{
+		{true, true},
+		{false, false},
+	}
+
+	for _, tc := range testCases {
+		t.Run("valid height with different initial height", func(t *testing.T) {
+			_, err := app.InitChain(&abci.InitChainRequest{
+				InitialHeight: headerHeight,
+			})
+			require.NoError(t, err)
+			ctx, err := app.CreateQueryContextWithCheckHeader(0, true, tc.checkHeader)
+			if tc.expErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, height, ctx.BlockHeight())
+			}
+		})
+	}
+}
+
+func TestABCI_CreateQueryContext_Before_Set_CheckState(t *testing.T) {
+	t.Parallel()
+
+	db := dbm.NewMemDB()
+	name := t.Name()
+	var height int64 = 2
+	var headerHeight int64 = 1
+
+	t.Run("valid height with different initial height", func(t *testing.T) {
+		app := baseapp.NewBaseApp(name, log.NewTestLogger(t), db, nil)
+
+		_, err := app.FinalizeBlock(&abci.FinalizeBlockRequest{Height: 1})
+		require.NoError(t, err)
+		_, err = app.Commit()
+		require.NoError(t, err)
+
+		_, err = app.FinalizeBlock(&abci.FinalizeBlockRequest{Height: 2})
+		require.NoError(t, err)
+
+		var queryCtx *sdk.Context
+		var queryCtxErr error
+		app.SetStreamingManager(storetypes.StreamingManager{
+			ABCIListeners: []storetypes.ABCIListener{
+				&mockABCIListener{
+					ListenCommitFn: func(context.Context, abci.CommitResponse, []*storetypes.StoreKVPair) error {
+						qCtx, qErr := app.CreateQueryContext(height, true)
+						queryCtx = &qCtx
+						queryCtxErr = qErr
+						return nil
+					},
+				},
+			},
+		})
+		_, err = app.Commit()
+		require.NoError(t, err)
+		require.NoError(t, queryCtxErr)
+		require.Equal(t, height, queryCtx.BlockHeight())
+		_, err = app.InitChain(&abci.InitChainRequest{
+			InitialHeight: headerHeight,
+		})
+		require.NoError(t, err)
+	})
+}
+
 func TestSetMinGasPrices(t *testing.T) {
 	minGasPrices := sdk.DecCoins{sdk.NewInt64DecCoin("stake", 5000)}
 	suite := NewBaseAppSuite(t, baseapp.SetMinGasPrices(minGasPrices.String()))