Add GRPCBlockHeightHeader in clientCtx gRPC requests (#7021)

* Add header in client context grpc

* Second try

* Make headers work

* Use QueryABCI

* Put header in types

Co-authored-by: Alexander Bezobchuk <alexanderbez@users.noreply.github.com>

baseapp/grpcrouter.go

@@ -3,14 +3,13 @@ package baseapp
 import (
 	"fmt"
 
-	"github.com/cosmos/cosmos-sdk/codec/types"
-
 	gogogrpc "github.com/gogo/protobuf/grpc"
 	abci "github.com/tendermint/tendermint/abci/types"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/encoding"
 	"google.golang.org/grpc/encoding/proto"
 
+	"github.com/cosmos/cosmos-sdk/codec/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 

baseapp/grpcserver.go

@@ -10,8 +10,8 @@ import (
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 
-	servergrpc "github.com/cosmos/cosmos-sdk/server/grpc"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	grpctypes "github.com/cosmos/cosmos-sdk/types/grpc"
 )
 
 // GRPCQueryRouter returns the GRPCQueryRouter of a BaseApp.
@@ -30,7 +30,7 @@ func (app *BaseApp) RegisterGRPCServer(server gogogrpc.Server) {
 
 		// Get height header from the request context, if present.
 		var height int64
-		if heightHeaders := md.Get(servergrpc.GRPCBlockHeightHeader); len(heightHeaders) > 0 {
+		if heightHeaders := md.Get(grpctypes.GRPCBlockHeightHeader); len(heightHeaders) > 0 {
 			height, err = strconv.ParseInt(heightHeaders[0], 10, 64)
 			if err != nil {
 				return nil, err
@@ -51,7 +51,7 @@ func (app *BaseApp) RegisterGRPCServer(server gogogrpc.Server) {
 		if height == 0 {
 			height = sdkCtx.BlockHeight() // If height was not set in the request, set it to the latest
 		}
-		md = metadata.Pairs(servergrpc.GRPCBlockHeightHeader, strconv.FormatInt(height, 10))
+		md = metadata.Pairs(grpctypes.GRPCBlockHeightHeader, strconv.FormatInt(height, 10))
 		grpc.SetHeader(grpcCtx, md)
 
 		return handler(grpcCtx, req)

client/grpc_query.go

@@ -3,11 +3,15 @@ package client
 import (
 	gocontext "context"
 	"fmt"
+	"strconv"
 
+	grpctypes "github.com/cosmos/cosmos-sdk/types/grpc"
 	gogogrpc "github.com/gogo/protobuf/grpc"
+	abci "github.com/tendermint/tendermint/abci/types"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/encoding"
 	"google.golang.org/grpc/encoding/proto"
+	"google.golang.org/grpc/metadata"
 
 	"github.com/cosmos/cosmos-sdk/codec/types"
 )
@@ -17,21 +21,41 @@ var _ gogogrpc.ClientConn = Context{}
 var protoCodec = encoding.GetCodec(proto.Name)
 
 // Invoke implements the grpc ClientConn.Invoke method
-func (ctx Context) Invoke(_ gocontext.Context, method string, args, reply interface{}, _ ...grpc.CallOption) error {
+func (ctx Context) Invoke(grpcCtx gocontext.Context, method string, args, reply interface{}, opts ...grpc.CallOption) error {
 	reqBz, err := protoCodec.Marshal(args)
 	if err != nil {
 		return err
 	}
-	resBz, _, err := ctx.QueryWithData(method, reqBz)
+	req := abci.RequestQuery{
+		Path: method,
+		Data: reqBz,
+	}
+
+	res, err := ctx.QueryABCI(req)
 	if err != nil {
 		return err
 	}
 
-	err = protoCodec.Unmarshal(resBz, reply)
+	err = protoCodec.Unmarshal(res.Value, reply)
 	if err != nil {
 		return err
 	}
 
+	// Create header metadata. For now the headers contain:
+	// - block height
+	// We then parse all the call options, if the call option is a
+	// HeaderCallOption, then we manually set the value of that header to the
+	// metadata.
+	md := metadata.Pairs(grpctypes.GRPCBlockHeightHeader, strconv.FormatInt(res.Height, 10))
+	for _, callOpt := range opts {
+		header, ok := callOpt.(grpc.HeaderCallOption)
+		if !ok {
+			continue
+		}
+
+		*header.HeaderAddr = md
+	}
+
 	if ctx.InterfaceRegistry != nil {
 		return types.UnpackInterfaces(reply, ctx.InterfaceRegistry)
 	}

client/grpc_query_test.go

@@ -0,0 +1,79 @@
+package client_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/cosmos/cosmos-sdk/testutil/network"
+	"github.com/cosmos/cosmos-sdk/testutil/testdata"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	grpctypes "github.com/cosmos/cosmos-sdk/types/grpc"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
+	"github.com/stretchr/testify/suite"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+type IntegrationTestSuite struct {
+	suite.Suite
+
+	network *network.Network
+}
+
+func (s *IntegrationTestSuite) SetupSuite() {
+	s.T().Log("setting up integration test suite")
+
+	s.network = network.New(s.T(), network.DefaultConfig())
+	s.Require().NotNil(s.network)
+
+	_, err := s.network.WaitForHeight(2)
+	s.Require().NoError(err)
+}
+
+func (s *IntegrationTestSuite) TearDownSuite() {
+	s.T().Log("tearing down integration test suite")
+	s.network.Cleanup()
+}
+
+func (s *IntegrationTestSuite) TestGRPCQuery() {
+	val0 := s.network.Validators[0]
+
+	// gRPC query to test service should work
+	testClient := testdata.NewTestServiceClient(val0.ClientCtx)
+	testRes, err := testClient.Echo(context.Background(), &testdata.EchoRequest{Message: "hello"})
+	s.Require().NoError(err)
+	s.Require().Equal("hello", testRes.Message)
+
+	// gRPC query to bank service should work
+	denom := fmt.Sprintf("%stoken", val0.Moniker)
+	bankClient := banktypes.NewQueryClient(val0.ClientCtx)
+	var header metadata.MD
+	bankRes, err := bankClient.Balance(
+		context.Background(),
+		&banktypes.QueryBalanceRequest{Address: val0.Address, Denom: denom},
+		grpc.Header(&header), // Also fetch grpc header
+	)
+	s.Require().NoError(err)
+	s.Require().Equal(
+		sdk.NewCoin(denom, s.network.Config.AccountTokens),
+		*bankRes.GetBalance(),
+	)
+	blockHeight := header.Get(grpctypes.GRPCBlockHeightHeader)
+	s.Require().NotEmpty(blockHeight[0]) // Should contain the block height
+
+	// Request metadata should work
+	val0.ClientCtx = val0.ClientCtx.WithHeight(1) // We set clientCtx to height 1
+	bankClient = banktypes.NewQueryClient(val0.ClientCtx)
+	bankRes, err = bankClient.Balance(
+		context.Background(),
+		&banktypes.QueryBalanceRequest{Address: val0.Address, Denom: denom},
+		grpc.Header(&header),
+	)
+	blockHeight = header.Get(grpctypes.GRPCBlockHeightHeader)
+	s.Require().Equal([]string{"1"}, blockHeight)
+}
+
+func TestIntegrationTestSuite(t *testing.T) {
+	suite.Run(t, new(IntegrationTestSuite))
+}

server/grpc/server.go

@@ -11,11 +11,6 @@ import (
 	"github.com/cosmos/cosmos-sdk/server/types"
 )
 
-const (
-	// GRPCBlockHeightHeader is the gRPC header for block height.
-	GRPCBlockHeightHeader = "x-cosmos-block-height"
-)
-
 // StartGRPCServer starts a gRPC server on the given address.
 func StartGRPCServer(app types.Application, address string) (*grpc.Server, error) {
 	grpcSrv := grpc.NewServer()

server/grpc/server_test.go

@@ -11,10 +11,10 @@ import (
 
 	rpb "google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
 
-	servergrpc "github.com/cosmos/cosmos-sdk/server/grpc"
 	"github.com/cosmos/cosmos-sdk/testutil/network"
 	"github.com/cosmos/cosmos-sdk/testutil/testdata"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	grpctypes "github.com/cosmos/cosmos-sdk/types/grpc"
 	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
 )
 
@@ -39,7 +39,7 @@ func (s *IntegrationTestSuite) TearDownSuite() {
 	s.network.Cleanup()
 }
 
-func (s *IntegrationTestSuite) TestGRPC() {
+func (s *IntegrationTestSuite) TestGRPCServer() {
 	val0 := s.network.Validators[0]
 	conn, err := grpc.Dial(
 		val0.AppConfig.GRPC.Address,
@@ -67,16 +67,16 @@ func (s *IntegrationTestSuite) TestGRPC() {
 		sdk.NewCoin(denom, s.network.Config.AccountTokens),
 		*bankRes.GetBalance(),
 	)
-	blockHeight := header.Get(servergrpc.GRPCBlockHeightHeader)
+	blockHeight := header.Get(grpctypes.GRPCBlockHeightHeader)
 	s.Require().NotEmpty(blockHeight[0]) // Should contain the block height
 
 	// Request metadata should work
 	bankRes, err = bankClient.Balance(
-		metadata.AppendToOutgoingContext(context.Background(), servergrpc.GRPCBlockHeightHeader, "1"), // Add metadata to request
+		metadata.AppendToOutgoingContext(context.Background(), grpctypes.GRPCBlockHeightHeader, "1"), // Add metadata to request
 		&banktypes.QueryBalanceRequest{Address: val0.Address, Denom: denom},
 		grpc.Header(&header),
 	)
-	blockHeight = header.Get(servergrpc.GRPCBlockHeightHeader)
+	blockHeight = header.Get(grpctypes.GRPCBlockHeightHeader)
 	s.Require().Equal([]string{"1"}, blockHeight)
 
 	// Test server reflection

types/grpc/headers.go

@@ -0,0 +1,6 @@
+package grpc
+
+const (
+	// GRPCBlockHeightHeader is the gRPC header for block height.
+	GRPCBlockHeightHeader = "x-cosmos-block-height"
+)