gRPC: fix server reflection (#8945)

* add: gogoreflection

* fix: server reflection

* fix: make tests resolve imports in a transitive way

* chore: cleanup getMessageType

* chore: lint

* add: extensions reflection

* chore: update interact-node.md docs

* chore: update CHANGELOG.md

* Update server/grpc/gogoreflection/fix_registration.go

Co-authored-by: Robert Zaremba <robert@zaremba.ch>

Co-authored-by: Jonathan Gimeno <jgimeno@gmail.com>
Co-authored-by: Alessio Treglia <alessio@tendermint.com>
Co-authored-by: Robert Zaremba <robert@zaremba.ch>

server/grpc/gogoreflection/doc.go

@@ -0,0 +1,5 @@
+// Package gogoreflection implements gRPC reflection for gogoproto consumers
+// the normal reflection library does not work as it points to a different
+// singleton registry. The API and codebase is taken from the official gRPC
+// reflection repository.
+package gogoreflection

server/grpc/gogoreflection/fix_registration.go

@@ -0,0 +1,143 @@
+package gogoreflection
+
+import (
+	"bytes"
+	"compress/gzip"
+	"fmt"
+	"reflect"
+
+	_ "github.com/gogo/protobuf/gogoproto" // required so it does register the gogoproto file descriptor
+	gogoproto "github.com/gogo/protobuf/proto"
+
+	// nolint: staticcheck
+	"github.com/golang/protobuf/proto"
+	dpb "github.com/golang/protobuf/protoc-gen-go/descriptor"
+	_ "github.com/regen-network/cosmos-proto" // look above
+)
+
+var importsToFix = map[string]string{
+	"gogo.proto":   "gogoproto/gogo.proto",
+	"cosmos.proto": "cosmos_proto/cosmos.proto",
+}
+
+// fixRegistration is required because certain files register themselves in a way
+// but are imported by other files in a different way.
+// NOTE(fdymylja): This fix should not be needed and should be addressed in some CI.
+// Currently every cosmos-sdk proto file is importing gogo.proto as gogoproto/gogo.proto,
+// but gogo.proto registers itself as gogo.proto, same goes for cosmos.proto.
+func fixRegistration(registeredAs, importedAs string) error {
+	raw := gogoproto.FileDescriptor(registeredAs)
+	if len(raw) == 0 {
+		return fmt.Errorf("file descriptor not found for %s", registeredAs)
+	}
+
+	fd, err := decodeFileDesc(raw)
+	if err != nil {
+		return err
+	}
+
+	// fix name
+	*fd.Name = importedAs
+	fixedRaw, err := compress(fd)
+	if err != nil {
+		return fmt.Errorf("unable to compress: %w", err)
+	}
+	gogoproto.RegisterFile(importedAs, fixedRaw)
+	return nil
+}
+
+func init() {
+	// we need to fix the gogoproto filedesc to match the import path
+	// in theory this shouldn't be required, generally speaking
+	// proto files should be imported as their registration path
+
+	for registeredAs, importedAs := range importsToFix {
+		err := fixRegistration(registeredAs, importedAs)
+		if err != nil {
+			panic(err)
+		}
+	}
+}
+
+// compress compresses the given file descriptor
+// nolint: interfacer
+func compress(fd *dpb.FileDescriptorProto) ([]byte, error) {
+	fdBytes, err := proto.Marshal(fd)
+	if err != nil {
+		return nil, err
+	}
+	buf := new(bytes.Buffer)
+	cw := gzip.NewWriter(buf)
+	_, err = cw.Write(fdBytes)
+	if err != nil {
+		return nil, err
+	}
+	err = cw.Close()
+	if err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+func getFileDescriptor(filePath string) []byte {
+	// since we got well known descriptors which are not registered into gogoproto registry
+	// but are instead registered into the proto one, we need to check both
+	fd := gogoproto.FileDescriptor(filePath)
+	if len(fd) != 0 {
+		return fd
+	}
+	// nolint: staticcheck
+	return proto.FileDescriptor(filePath)
+}
+
+func getMessageType(name string) reflect.Type {
+	typ := gogoproto.MessageType(name)
+	if typ != nil {
+		return typ
+	}
+	// nolint: staticcheck
+	return proto.MessageType(name)
+}
+
+func getExtension(extID int32, m proto.Message) *gogoproto.ExtensionDesc {
+	// check first in gogoproto registry
+	for id, desc := range gogoproto.RegisteredExtensions(m) {
+		if id == extID {
+			return desc
+		}
+	}
+	// check into proto registry
+	// nolint: staticcheck
+	for id, desc := range proto.RegisteredExtensions(m) {
+		if id == extID {
+			return &gogoproto.ExtensionDesc{
+				ExtendedType:  desc.ExtendedType,
+				ExtensionType: desc.ExtensionType,
+				Field:         desc.Field,
+				Name:          desc.Name,
+				Tag:           desc.Tag,
+				Filename:      desc.Filename,
+			}
+		}
+	}
+
+	return nil
+}
+
+func getExtensionsNumbers(m proto.Message) []int32 {
+	gogoProtoExts := gogoproto.RegisteredExtensions(m)
+	out := make([]int32, 0, len(gogoProtoExts))
+	for id := range gogoProtoExts {
+		out = append(out, id)
+	}
+	if len(out) != 0 {
+		return out
+	}
+	// nolint: staticcheck
+	protoExts := proto.RegisteredExtensions(m)
+	out = make([]int32, 0, len(protoExts))
+	for id := range protoExts {
+		out = append(out, id)
+	}
+	return out
+}

server/grpc/gogoreflection/fix_registration_test.go

@@ -0,0 +1,22 @@
+package gogoreflection
+
+import (
+	"testing"
+
+	"google.golang.org/protobuf/runtime/protoimpl"
+)
+
+func TestRegistrationFix(t *testing.T) {
+	res := getFileDescriptor("gogoproto/gogo.proto")
+	rawDesc, err := decompress(res)
+	if err != nil {
+		t.Fatal(err)
+	}
+	fd := protoimpl.DescBuilder{
+		RawDescriptor: rawDesc,
+	}.Build()
+
+	if fd.File.Extensions().Len() == 0 {
+		t.Fatal("unexpected parsing")
+	}
+}

server/grpc/gogoreflection/serverreflection.go

@@ -0,0 +1,478 @@
+/*
+ *
+ * Copyright 2016 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+/*
+Package reflection implements server reflection service.
+
+The service implemented is defined in:
+https://github.com/grpc/grpc/blob/master/src/proto/grpc/reflection/v1alpha/reflection.proto.
+
+To register server reflection on a gRPC server:
+	import "google.golang.org/grpc/reflection"
+
+	s := grpc.NewServer()
+	pb.RegisterYourOwnServer(s, &server{})
+
+	// Register reflection service on gRPC server.
+	reflection.Register(s)
+
+	s.Serve(lis)
+
+*/
+package gogoreflection // import "google.golang.org/grpc/reflection"
+
+import (
+	"bytes"
+	"compress/gzip"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"log"
+	"reflect"
+	"sort"
+	"sync"
+
+	// nolint: staticcheck
+	"github.com/golang/protobuf/proto"
+	dpb "github.com/golang/protobuf/protoc-gen-go/descriptor"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	rpb "google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
+	"google.golang.org/grpc/status"
+)
+
+type serverReflectionServer struct {
+	rpb.UnimplementedServerReflectionServer
+	s *grpc.Server
+
+	initSymbols  sync.Once
+	serviceNames []string
+	symbols      map[string]*dpb.FileDescriptorProto // map of fully-qualified names to files
+}
+
+// Register registers the server reflection service on the given gRPC server.
+func Register(s *grpc.Server) {
+	rpb.RegisterServerReflectionServer(s, &serverReflectionServer{
+		s: s,
+	})
+}
+
+// protoMessage is used for type assertion on proto messages.
+// Generated proto message implements function Descriptor(), but Descriptor()
+// is not part of interface proto.Message. This interface is needed to
+// call Descriptor().
+type protoMessage interface {
+	Descriptor() ([]byte, []int)
+}
+
+func (s *serverReflectionServer) getSymbols() (svcNames []string, symbolIndex map[string]*dpb.FileDescriptorProto) {
+	s.initSymbols.Do(func() {
+		serviceInfo := s.s.GetServiceInfo()
+
+		s.symbols = map[string]*dpb.FileDescriptorProto{}
+		s.serviceNames = make([]string, 0, len(serviceInfo))
+		processed := map[string]struct{}{}
+		for svc, info := range serviceInfo {
+			s.serviceNames = append(s.serviceNames, svc)
+			fdenc, ok := parseMetadata(info.Metadata)
+			if !ok {
+				continue
+			}
+			fd, err := decodeFileDesc(fdenc)
+			if err != nil {
+				continue
+			}
+			s.processFile(fd, processed)
+		}
+		sort.Strings(s.serviceNames)
+	})
+
+	return s.serviceNames, s.symbols
+}
+
+func (s *serverReflectionServer) processFile(fd *dpb.FileDescriptorProto, processed map[string]struct{}) {
+	filename := fd.GetName()
+	if _, ok := processed[filename]; ok {
+		return
+	}
+	processed[filename] = struct{}{}
+
+	prefix := fd.GetPackage()
+
+	for _, msg := range fd.MessageType {
+		s.processMessage(fd, prefix, msg)
+	}
+	for _, en := range fd.EnumType {
+		s.processEnum(fd, prefix, en)
+	}
+	for _, ext := range fd.Extension {
+		s.processField(fd, prefix, ext)
+	}
+	for _, svc := range fd.Service {
+		svcName := fqn(prefix, svc.GetName())
+		s.symbols[svcName] = fd
+		for _, meth := range svc.Method {
+			name := fqn(svcName, meth.GetName())
+			s.symbols[name] = fd
+		}
+	}
+
+	for _, dep := range fd.Dependency {
+		fdenc := getFileDescriptor(dep)
+		fdDep, err := decodeFileDesc(fdenc)
+		if err != nil {
+			continue
+		}
+		s.processFile(fdDep, processed)
+	}
+}
+
+func (s *serverReflectionServer) processMessage(fd *dpb.FileDescriptorProto, prefix string, msg *dpb.DescriptorProto) {
+	msgName := fqn(prefix, msg.GetName())
+	s.symbols[msgName] = fd
+
+	for _, nested := range msg.NestedType {
+		s.processMessage(fd, msgName, nested)
+	}
+	for _, en := range msg.EnumType {
+		s.processEnum(fd, msgName, en)
+	}
+	for _, ext := range msg.Extension {
+		s.processField(fd, msgName, ext)
+	}
+	for _, fld := range msg.Field {
+		s.processField(fd, msgName, fld)
+	}
+	for _, oneof := range msg.OneofDecl {
+		oneofName := fqn(msgName, oneof.GetName())
+		s.symbols[oneofName] = fd
+	}
+}
+
+func (s *serverReflectionServer) processEnum(fd *dpb.FileDescriptorProto, prefix string, en *dpb.EnumDescriptorProto) {
+	enName := fqn(prefix, en.GetName())
+	s.symbols[enName] = fd
+
+	for _, val := range en.Value {
+		valName := fqn(enName, val.GetName())
+		s.symbols[valName] = fd
+	}
+}
+
+func (s *serverReflectionServer) processField(fd *dpb.FileDescriptorProto, prefix string, fld *dpb.FieldDescriptorProto) {
+	fldName := fqn(prefix, fld.GetName())
+	s.symbols[fldName] = fd
+}
+
+func fqn(prefix, name string) string {
+	if prefix == "" {
+		return name
+	}
+	return prefix + "." + name
+}
+
+// fileDescForType gets the file descriptor for the given type.
+// The given type should be a proto message.
+func (s *serverReflectionServer) fileDescForType(st reflect.Type) (*dpb.FileDescriptorProto, error) {
+	m, ok := reflect.Zero(reflect.PtrTo(st)).Interface().(protoMessage)
+	if !ok {
+		return nil, fmt.Errorf("failed to create message from type: %v", st)
+	}
+	enc, _ := m.Descriptor()
+
+	return decodeFileDesc(enc)
+}
+
+// decodeFileDesc does decompression and unmarshalling on the given
+// file descriptor byte slice.
+func decodeFileDesc(enc []byte) (*dpb.FileDescriptorProto, error) {
+	raw, err := decompress(enc)
+	if err != nil {
+		return nil, fmt.Errorf("failed to decompress enc: %v", err)
+	}
+
+	fd := new(dpb.FileDescriptorProto)
+	if err := proto.Unmarshal(raw, fd); err != nil {
+		return nil, fmt.Errorf("bad descriptor: %v", err)
+	}
+	return fd, nil
+}
+
+// decompress does gzip decompression.
+func decompress(b []byte) ([]byte, error) {
+	r, err := gzip.NewReader(bytes.NewReader(b))
+	if err != nil {
+		return nil, fmt.Errorf("bad gzipped descriptor: %v", err)
+	}
+	out, err := ioutil.ReadAll(r)
+	if err != nil {
+		return nil, fmt.Errorf("bad gzipped descriptor: %v", err)
+	}
+	return out, nil
+}
+
+func typeForName(name string) (reflect.Type, error) {
+	pt := getMessageType(name)
+	if pt == nil {
+		return nil, fmt.Errorf("unknown type: %q", name)
+	}
+	st := pt.Elem()
+
+	return st, nil
+}
+
+func fileDescContainingExtension(st reflect.Type, ext int32) (*dpb.FileDescriptorProto, error) {
+	m, ok := reflect.Zero(reflect.PtrTo(st)).Interface().(proto.Message)
+	if !ok {
+		return nil, fmt.Errorf("failed to create message from type: %v", st)
+	}
+
+	extDesc := getExtension(ext, m)
+
+	if extDesc == nil {
+		return nil, fmt.Errorf("failed to find registered extension for extension number %v", ext)
+	}
+
+	return decodeFileDesc(getFileDescriptor(extDesc.Filename))
+}
+
+func (s *serverReflectionServer) allExtensionNumbersForType(st reflect.Type) ([]int32, error) {
+	m, ok := reflect.Zero(reflect.PtrTo(st)).Interface().(proto.Message)
+	if !ok {
+		return nil, fmt.Errorf("failed to create message from type: %v", st)
+	}
+
+	out := getExtensionsNumbers(m)
+	return out, nil
+}
+
+// fileDescWithDependencies returns a slice of serialized fileDescriptors in
+// wire format ([]byte). The fileDescriptors will include fd and all the
+// transitive dependencies of fd with names not in sentFileDescriptors.
+func fileDescWithDependencies(fd *dpb.FileDescriptorProto, sentFileDescriptors map[string]bool) ([][]byte, error) {
+	r := [][]byte{}
+	queue := []*dpb.FileDescriptorProto{fd}
+	for len(queue) > 0 {
+		currentfd := queue[0]
+		queue = queue[1:]
+		if sent := sentFileDescriptors[currentfd.GetName()]; len(r) == 0 || !sent {
+			sentFileDescriptors[currentfd.GetName()] = true
+			currentfdEncoded, err := proto.Marshal(currentfd)
+			if err != nil {
+				return nil, err
+			}
+			r = append(r, currentfdEncoded)
+		}
+		for _, dep := range currentfd.Dependency {
+			fdenc := getFileDescriptor(dep)
+			fdDep, err := decodeFileDesc(fdenc)
+			if err != nil {
+				continue
+			}
+			queue = append(queue, fdDep)
+		}
+	}
+	return r, nil
+}
+
+// fileDescEncodingByFilename finds the file descriptor for given filename,
+// finds all of its previously unsent transitive dependencies, does marshalling
+// on them, and returns the marshalled result.
+func (s *serverReflectionServer) fileDescEncodingByFilename(name string, sentFileDescriptors map[string]bool) ([][]byte, error) {
+	enc := getFileDescriptor(name)
+	if enc == nil {
+		return nil, fmt.Errorf("unknown file: %v", name)
+	}
+	fd, err := decodeFileDesc(enc)
+	if err != nil {
+		return nil, err
+	}
+	return fileDescWithDependencies(fd, sentFileDescriptors)
+}
+
+// parseMetadata finds the file descriptor bytes specified meta.
+// For SupportPackageIsVersion4, m is the name of the proto file, we
+// call proto.FileDescriptor to get the byte slice.
+// For SupportPackageIsVersion3, m is a byte slice itself.
+func parseMetadata(meta interface{}) ([]byte, bool) {
+	// Check if meta is the file name.
+	if fileNameForMeta, ok := meta.(string); ok {
+		return getFileDescriptor(fileNameForMeta), true
+	}
+
+	// Check if meta is the byte slice.
+	if enc, ok := meta.([]byte); ok {
+		return enc, true
+	}
+
+	return nil, false
+}
+
+// fileDescEncodingContainingSymbol finds the file descriptor containing the
+// given symbol, finds all of its previously unsent transitive dependencies,
+// does marshalling on them, and returns the marshalled result. The given symbol
+// can be a type, a service or a method.
+func (s *serverReflectionServer) fileDescEncodingContainingSymbol(name string, sentFileDescriptors map[string]bool) ([][]byte, error) {
+	_, symbols := s.getSymbols()
+	fd := symbols[name]
+	if fd == nil {
+		// Check if it's a type name that was not present in the
+		// transitive dependencies of the registered services.
+		if st, err := typeForName(name); err == nil {
+			fd, err = s.fileDescForType(st)
+			if err != nil {
+				return nil, err
+			}
+		}
+	}
+
+	if fd == nil {
+		return nil, fmt.Errorf("unknown symbol: %v", name)
+	}
+
+	return fileDescWithDependencies(fd, sentFileDescriptors)
+}
+
+// fileDescEncodingContainingExtension finds the file descriptor containing
+// given extension, finds all of its previously unsent transitive dependencies,
+// does marshalling on them, and returns the marshalled result.
+func (s *serverReflectionServer) fileDescEncodingContainingExtension(typeName string, extNum int32, sentFileDescriptors map[string]bool) ([][]byte, error) {
+	st, err := typeForName(typeName)
+	if err != nil {
+		return nil, err
+	}
+	fd, err := fileDescContainingExtension(st, extNum)
+	if err != nil {
+		return nil, err
+	}
+	return fileDescWithDependencies(fd, sentFileDescriptors)
+}
+
+// allExtensionNumbersForTypeName returns all extension numbers for the given type.
+func (s *serverReflectionServer) allExtensionNumbersForTypeName(name string) ([]int32, error) {
+	st, err := typeForName(name)
+	if err != nil {
+		return nil, err
+	}
+	extNums, err := s.allExtensionNumbersForType(st)
+	if err != nil {
+		return nil, err
+	}
+	return extNums, nil
+}
+
+// ServerReflectionInfo is the reflection service handler.
+func (s *serverReflectionServer) ServerReflectionInfo(stream rpb.ServerReflection_ServerReflectionInfoServer) error {
+	sentFileDescriptors := make(map[string]bool)
+	for {
+		in, err := stream.Recv()
+		if err == io.EOF {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+
+		out := &rpb.ServerReflectionResponse{
+			ValidHost:       in.Host,
+			OriginalRequest: in,
+		}
+		switch req := in.MessageRequest.(type) {
+		case *rpb.ServerReflectionRequest_FileByFilename:
+			b, err := s.fileDescEncodingByFilename(req.FileByFilename, sentFileDescriptors)
+			if err != nil {
+				out.MessageResponse = &rpb.ServerReflectionResponse_ErrorResponse{
+					ErrorResponse: &rpb.ErrorResponse{
+						ErrorCode:    int32(codes.NotFound),
+						ErrorMessage: err.Error(),
+					},
+				}
+			} else {
+				out.MessageResponse = &rpb.ServerReflectionResponse_FileDescriptorResponse{
+					FileDescriptorResponse: &rpb.FileDescriptorResponse{FileDescriptorProto: b},
+				}
+			}
+		case *rpb.ServerReflectionRequest_FileContainingSymbol:
+			b, err := s.fileDescEncodingContainingSymbol(req.FileContainingSymbol, sentFileDescriptors)
+			if err != nil {
+				out.MessageResponse = &rpb.ServerReflectionResponse_ErrorResponse{
+					ErrorResponse: &rpb.ErrorResponse{
+						ErrorCode:    int32(codes.NotFound),
+						ErrorMessage: err.Error(),
+					},
+				}
+			} else {
+				out.MessageResponse = &rpb.ServerReflectionResponse_FileDescriptorResponse{
+					FileDescriptorResponse: &rpb.FileDescriptorResponse{FileDescriptorProto: b},
+				}
+			}
+		case *rpb.ServerReflectionRequest_FileContainingExtension:
+			typeName := req.FileContainingExtension.ContainingType
+			extNum := req.FileContainingExtension.ExtensionNumber
+			b, err := s.fileDescEncodingContainingExtension(typeName, extNum, sentFileDescriptors)
+			if err != nil {
+				out.MessageResponse = &rpb.ServerReflectionResponse_ErrorResponse{
+					ErrorResponse: &rpb.ErrorResponse{
+						ErrorCode:    int32(codes.NotFound),
+						ErrorMessage: err.Error(),
+					},
+				}
+			} else {
+				out.MessageResponse = &rpb.ServerReflectionResponse_FileDescriptorResponse{
+					FileDescriptorResponse: &rpb.FileDescriptorResponse{FileDescriptorProto: b},
+				}
+			}
+		case *rpb.ServerReflectionRequest_AllExtensionNumbersOfType:
+			extNums, err := s.allExtensionNumbersForTypeName(req.AllExtensionNumbersOfType)
+			if err != nil {
+				out.MessageResponse = &rpb.ServerReflectionResponse_ErrorResponse{
+					ErrorResponse: &rpb.ErrorResponse{
+						ErrorCode:    int32(codes.NotFound),
+						ErrorMessage: err.Error(),
+					},
+				}
+				log.Printf("OH NO: %s", err)
+			} else {
+				out.MessageResponse = &rpb.ServerReflectionResponse_AllExtensionNumbersResponse{
+					AllExtensionNumbersResponse: &rpb.ExtensionNumberResponse{
+						BaseTypeName:    req.AllExtensionNumbersOfType,
+						ExtensionNumber: extNums,
+					},
+				}
+			}
+		case *rpb.ServerReflectionRequest_ListServices:
+			svcNames, _ := s.getSymbols()
+			serviceResponses := make([]*rpb.ServiceResponse, len(svcNames))
+			for i, n := range svcNames {
+				serviceResponses[i] = &rpb.ServiceResponse{
+					Name: n,
+				}
+			}
+			out.MessageResponse = &rpb.ServerReflectionResponse_ListServicesResponse{
+				ListServicesResponse: &rpb.ListServiceResponse{
+					Service: serviceResponses,
+				},
+			}
+		default:
+			return status.Errorf(codes.InvalidArgument, "invalid MessageRequest: %v", in.MessageRequest)
+		}
+		if err := stream.Send(out); err != nil {
+			return err
+		}
+	}
+}

server/grpc/server.go

@@ -6,9 +6,9 @@ import (
 	"time"
 
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/reflection"
 
 	"github.com/cosmos/cosmos-sdk/client"
+	"github.com/cosmos/cosmos-sdk/server/grpc/gogoreflection"
 	"github.com/cosmos/cosmos-sdk/server/types"
 )
 
@@ -19,7 +19,7 @@ func StartGRPCServer(clientCtx client.Context, app types.Application, address st
 
 	// Reflection allows external clients to see what services and methods
 	// the gRPC server exposes.
-	reflection.Register(grpcSrv)
+	gogoreflection.Register(grpcSrv)
 
 	listener, err := net.Listen("tcp", address)
 	if err != nil {

server/grpc/server_test.go

@@ -6,6 +6,9 @@ import (
 	"context"
 	"fmt"
 	"testing"
+	"time"
+
+	"github.com/jhump/protoreflect/grpcreflect"
 
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
@@ -96,21 +99,25 @@ func (s *IntegrationTestSuite) TestGRPCServer_BankBalance() {
 
 func (s *IntegrationTestSuite) TestGRPCServer_Reflection() {
 	// Test server reflection
-	reflectClient := rpb.NewServerReflectionClient(s.conn)
-	stream, err := reflectClient.ServerReflectionInfo(context.Background(), grpc.WaitForReady(true))
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+	defer cancel()
+	stub := rpb.NewServerReflectionClient(s.conn)
+	// NOTE(fdymylja): we use grpcreflect because it solves imports too
+	// so that we can always assert that given a reflection server it is
+	// possible to fully query all the methods, without having any context
+	// on the proto registry
+	rc := grpcreflect.NewClient(ctx, stub)
+
+	services, err := rc.ListServices()
 	s.Require().NoError(err)
-	s.Require().NoError(stream.Send(&rpb.ServerReflectionRequest{
-		MessageRequest: &rpb.ServerReflectionRequest_ListServices{},
-	}))
-	res, err := stream.Recv()
-	s.Require().NoError(err)
-	services := res.GetListServicesResponse().Service
-	servicesMap := make(map[string]bool)
-	for _, s := range services {
-		servicesMap[s.Name] = true
+	s.Require().Greater(len(services), 0)
+
+	for _, svc := range services {
+		file, err := rc.FileContainingSymbol(svc)
+		s.Require().NoError(err)
+		sd := file.FindSymbol(svc)
+		s.Require().NotNil(sd)
 	}
-	// Make sure the following services are present
-	s.Require().True(servicesMap["cosmos.bank.v1beta1.Query"])
 }
 
 func (s *IntegrationTestSuite) TestGRPCServer_GetTxsEvent() {