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>

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))
+}