fix to register file descriptors and message descriptors (#544)

.github/workflows/main.yml

@@ -31,7 +31,8 @@ jobs:
       run: go vet ./...
 
     - name: Test
-      run: go test -p 1 -coverpkg ./... -covermode atomic -coverprofile coverage.txt ./...
+      # TODO: Remove conflictPolicy flag.
+      run: go test -race -coverpkg ./... -covermode atomic -coverprofile coverage.txt -ldflags "-X google.golang.org/protobuf/reflect/protoregistry.conflictPolicy=warn" ./...
 
     - name: Upload coverage to Codecov
       uses: codecov/[email protected]

Makefile

@@ -32,7 +32,8 @@ credits:
 
 .PHONY: gotest
 gotest: lint
-	go test -race ./...
+	# TODO: Remove conflictPolicy flag.
+	go test -race -ldflags "-X google.golang.org/protobuf/reflect/protoregistry.conflictPolicy=warn" ./...
 
 .PHONY: lint
 lint:

idl/proto/proto.go

@@ -14,6 +14,9 @@ import (
 	"github.com/ktr0731/evans/grpc/grpcreflection"
 	"github.com/ktr0731/evans/idl"
 	"github.com/pkg/errors"
+	"google.golang.org/protobuf/reflect/protodesc"
+	"google.golang.org/protobuf/reflect/protoregistry"
+	"google.golang.org/protobuf/types/dynamicpb"
 )
 
 type spec struct {
@@ -135,12 +138,7 @@ func LoadFiles(importPaths []string, fnames []string) (idl.Spec, error) {
 		return nil, errors.Wrap(err, "proto: failed to parse passed proto files")
 	}
 
-	// Collect dependency file descriptors
-	for _, d := range fileDescs {
-		fileDescs = append(fileDescs, d.GetDependencies()...)
-	}
-
-	return newSpec(fileDescs), nil
+	return newSpec(fileDescs)
 }
 
 // LoadByReflection receives a gRPC reflection client, then tries to instantiate a new idl.Spec by using gRPC reflection.
@@ -149,10 +147,10 @@ func LoadByReflection(client grpcreflection.Client) (idl.Spec, error) {
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to list packages by gRPC reflection")
 	}
-	return newSpec(fileDescs), nil
+	return newSpec(fileDescs)
 }
 
-func newSpec(fds []*desc.FileDescriptor) idl.Spec {
+func newSpec(fds []*desc.FileDescriptor) (idl.Spec, error) {
 	var (
 		encounteredPkgs = make(map[string]interface{})
 		encounteredSvcs = make(map[string]interface{})
@@ -162,6 +160,10 @@ func newSpec(fds []*desc.FileDescriptor) idl.Spec {
 		msgDescs        = make(map[string]*desc.MessageDescriptor)
 	)
 	for _, f := range fds {
+		if err := registerFileAndType(f); err != nil {
+			return nil, errors.Wrap(err, "failed to register file descriptor")
+		}
+
 		if _, encountered := encounteredPkgs[f.GetPackage()]; !encountered {
 			pkgNames = append(pkgNames, f.GetPackage())
 			encounteredPkgs[f.GetPackage()] = nil
@@ -190,7 +192,38 @@ func newSpec(fds []*desc.FileDescriptor) idl.Spec {
 		svcDescs:  svcDescs,
 		rpcDescs:  rpcDescs,
 		msgDescs:  msgDescs,
+	}, nil
+}
+
+func registerFileAndType(fd *desc.FileDescriptor) error {
+	deps := fd.GetDependencies()
+	for _, d := range deps {
+		if err := registerFileAndType(d); err != nil {
+			return err
+		}
+	}
+
+	prfd, err := protodesc.NewFile(fd.AsFileDescriptorProto(), protoregistry.GlobalFiles)
+	if err != nil {
+		return errors.Wrap(err, "failed to new file descriptor")
 	}
+
+	if _, err := protoregistry.GlobalFiles.FindFileByPath(prfd.Path()); errors.Is(err, protoregistry.NotFound) {
+		if err := protoregistry.GlobalFiles.RegisterFile(prfd); err != nil {
+			return err
+		}
+	}
+
+	for i := 0; i < prfd.Messages().Len(); i++ {
+		md := prfd.Messages().Get(i)
+		if _, err := protoregistry.GlobalTypes.FindMessageByName(md.FullName()); errors.Is(err, protoregistry.NotFound) {
+			if err := protoregistry.GlobalTypes.RegisterMessage(dynamicpb.NewMessageType(md)); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
 }
 
 // FullyQualifiedServiceName returns the fully-qualified service name.