add: extensions reflection

server/grpc/gogoreflection/fix_registration.go

@@ -98,3 +98,46 @@ func getMessageType(name string) reflect.Type {
 	// 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/serverreflection.go

@@ -47,7 +47,6 @@ import (
 	"sort"
 	"sync"
 
-	gogoproto "github.com/gogo/protobuf/proto"
 	// nolint: staticcheck
 	"github.com/golang/protobuf/proto"
 	dpb "github.com/golang/protobuf/protoc-gen-go/descriptor"
@@ -243,13 +242,7 @@ func fileDescContainingExtension(st reflect.Type, ext int32) (*dpb.FileDescripto
 		return nil, fmt.Errorf("failed to create message from type: %v", st)
 	}
 
-	var extDesc *gogoproto.ExtensionDesc
-	for id, desc := range gogoproto.RegisteredExtensions(m) {
-		if id == ext {
-			extDesc = desc
-			break
-		}
-	}
+	extDesc := getExtension(ext, m)
 
 	if extDesc == nil {
 		return nil, fmt.Errorf("failed to find registered extension for extension number %v", ext)
@@ -264,11 +257,7 @@ func (s *serverReflectionServer) allExtensionNumbersForType(st reflect.Type) ([]
 		return nil, fmt.Errorf("failed to create message from type: %v", st)
 	}
 
-	exts := gogoproto.RegisteredExtensions(m)
-	out := make([]int32, 0, len(exts))
-	for id := range exts {
-		out = append(out, id)
-	}
+	out := getExtensionsNumbers(m)
 	return out, nil
 }