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

CHANGELOG.md

@@ -117,6 +117,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.
 
 ### API Breaking Changes
 

baseapp/abci.go

@@ -1245,15 +1245,37 @@ func (app *BaseApp) CreateQueryContext(height int64, prove bool) (sdk.Context, e
 			)
 	}
 
-	// branch the commit multi-store for safety
+	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 h.Height == height || 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,
+			)
+	}
+
 	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 {

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,50 @@ func TestABCI_CreateQueryContext(t *testing.T) {
 	}
 }
 
+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()))