From 77b94acc21ec201af3fdae340f06b364679f46d6 Mon Sep 17 00:00:00 2001
From: Josh Humphries <jh@fullstory.com>
Date: Fri, 23 Apr 2021 17:01:07 -0400
Subject: [PATCH] use alternate interface; change signature of Register

---
 reflection/serverreflection.go | 34 +++++++++++++++-------------------
 1 file changed, 15 insertions(+), 19 deletions(-)

diff --git a/reflection/serverreflection.go b/reflection/serverreflection.go
index 73446673a11..82a5ba7f244 100644
--- a/reflection/serverreflection.go
+++ b/reflection/serverreflection.go
@@ -54,9 +54,19 @@ import (
 	"google.golang.org/grpc/status"
 )
 
+// GRPCServer is the interface provided by a gRPC server. It is implemented by
+// *grpc.Server, but could also be implemented by other concrete types. It acts
+// as a registry, for accumulating the services exposed by the server.
+type GRPCServer interface {
+	grpc.ServiceRegistrar
+	GetServiceInfo() map[string]grpc.ServiceInfo
+}
+
+var _ GRPCServer = (*grpc.Server)(nil)
+
 type serverReflectionServer struct {
 	rpb.UnimplementedServerReflectionServer
-	queryServices func() map[string]grpc.ServiceInfo
+	s GRPCServer
 
 	initSymbols  sync.Once
 	serviceNames []string
@@ -64,23 +74,9 @@ type serverReflectionServer struct {
 }
 
 // Register registers the server reflection service on the given gRPC server.
-func Register(s *grpc.Server) {
-	register(s, s.GetServiceInfo)
-}
-
-// RegisterWithRegistrar registers the server reflection service with the given registrar.
-// The given map describes the services that will be exposed via reflection.
-func RegisterWithRegistrar(r grpc.ServiceRegistrar, svcs map[string]grpc.ServiceInfo) {
-	rpb.RegisterServerReflectionServer(r, &serverReflectionServer{
-		queryServices: func() map[string]grpc.ServiceInfo {
-			return svcs
-		},
-	})
-}
-
-func register(r grpc.ServiceRegistrar, queryServices func() map[string]grpc.ServiceInfo) {
-	rpb.RegisterServerReflectionServer(r, &serverReflectionServer{
-		queryServices: queryServices,
+func Register(s GRPCServer) {
+	rpb.RegisterServerReflectionServer(s, &serverReflectionServer{
+		s: s,
 	})
 }
 
@@ -94,7 +90,7 @@ type protoMessage interface {
 
 func (s *serverReflectionServer) getSymbols() (svcNames []string, symbolIndex map[string]*dpb.FileDescriptorProto) {
 	s.initSymbols.Do(func() {
-		serviceInfo := s.queryServices()
+		serviceInfo := s.s.GetServiceInfo()
 
 		s.symbols = map[string]*dpb.FileDescriptorProto{}
 		s.serviceNames = make([]string, 0, len(serviceInfo))