Attach instance of AnyResolver to jsonpb.Marshaler for gRPC (uber#422)

internal/cmd/BUILD.bazel

@@ -34,6 +34,8 @@ go_test(
         "@com_github_stretchr_testify//require:go_default_library",
         "@io_bazel_rules_go//proto/wkt:descriptor_go_proto",
         "@org_golang_google_grpc//:go_default_library",
+        "@org_golang_google_grpc//codes:go_default_library",
         "@org_golang_google_grpc//credentials:go_default_library",
+        "@org_golang_google_grpc//status:go_default_library",
     ],
 )

internal/cmd/cmd_test.go

@@ -46,7 +46,9 @@ import (
 	"github.com/uber/prototool/internal/settings"
 	"github.com/uber/prototool/internal/vars"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/status"
 )
 
 func TestCompile(t *testing.T) {
@@ -1444,10 +1446,7 @@ func TestGRPC(t *testing.T) {
 		ssClientCrt,
 		"--server-name", "client.local",
 	)
-}
 
-func TestGRPCError(t *testing.T) {
-	t.Parallel()
 	assertGRPCExclamationError(t,
 		errors.New("test"),
 		1,
@@ -1459,6 +1458,20 @@ func TestGRPCError(t *testing.T) {
 		`{"value":"hello"}`,
 		`--details`,
 	)
+
+	st, err := status.New(codes.InvalidArgument, "test").WithDetails(&grpcpb.Foo{Bar: "baz"})
+	assert.NoError(t, err)
+	assertGRPCExclamationError(t,
+		status.ErrorProto(st.Proto()),
+		1,
+		`{"status":{"code":3,"message":"test","details":[{"@type":"type.googleapis.com/grpc.Foo","bar":"baz"}]}}
+		{"trailers":{"content-type":["application/grpc"]}}
+		rpc error: code = InvalidArgument desc = test`,
+		"testdata/grpc/grpc.proto",
+		"grpc.ExcitedService/Exclamation",
+		`{"value":"hello"}`,
+		`--details`,
+	)
 }
 
 func TestVersion(t *testing.T) {

internal/cmd/testdata/grpc/grpc.proto

@@ -7,6 +7,10 @@ option java_multiple_files = true;
 option java_outer_classname = "GrpcProto";
 option java_package = "com.grpc";
 
+message Foo {
+  string bar = 1;
+}
+
 message ExclamationRequest {
   string value = 1;
 }

internal/grpc/BUILD.bazel

@@ -15,6 +15,7 @@ go_library(
         "@com_github_golang_protobuf//jsonpb:go_default_library_gen",
         "@com_github_golang_protobuf//proto:go_default_library",
         "@com_github_jhump_protoreflect//desc:go_default_library",
+        "@com_github_jhump_protoreflect//dynamic:go_default_library",
         "@io_bazel_rules_go//proto/wkt:descriptor_go_proto",
         "@org_golang_google_grpc//:go_default_library",
         "@org_golang_google_grpc//credentials:go_default_library",

internal/grpc/handler.go

@@ -33,6 +33,8 @@ import (
 	"github.com/golang/protobuf/jsonpb"
 	"github.com/golang/protobuf/proto"
 	"github.com/golang/protobuf/protoc-gen-go/descriptor"
+	protoreflectdesc "github.com/jhump/protoreflect/desc"
+	"github.com/jhump/protoreflect/dynamic"
 	"github.com/uber/prototool/internal/desc"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
@@ -72,7 +74,11 @@ func newHandler(options ...HandlerOption) *handler {
 }
 
 func (h *handler) Invoke(fileDescriptorSets []*descriptor.FileDescriptorSet, address string, method string, inputReader io.Reader, outputWriter io.Writer) error {
-	descriptorSource, err := h.getDescriptorSourceForMethod(fileDescriptorSets, method)
+	descriptorSource, err := getDescriptorSourceForMethod(fileDescriptorSets, method)
+	if err != nil {
+		return err
+	}
+	anyResolver, err := getAnyResolver(fileDescriptorSets)
 	if err != nil {
 		return err
 	}
@@ -81,7 +87,7 @@ func (h *handler) Invoke(fileDescriptorSets []*descriptor.FileDescriptorSet, add
 		return err
 	}
 	defer func() { _ = clientConn.Close() }()
-	invocationEventHandler := newInvocationEventHandler(outputWriter, h.logger, h.details)
+	invocationEventHandler := newInvocationEventHandler(anyResolver, outputWriter, h.logger, h.details)
 	ctx, cancel := context.WithTimeout(context.Background(), h.callTimeout)
 	defer cancel()
 	if err := grpcurl.InvokeRPC(
@@ -157,7 +163,28 @@ func (h *handler) getDialOptions() []grpc.DialOption {
 	return dialOptions
 }
 
-func (h *handler) getDescriptorSourceForMethod(fileDescriptorSets []*descriptor.FileDescriptorSet, method string) (grpcurl.DescriptorSource, error) {
+func getAnyResolver(fileDescriptorSets []*descriptor.FileDescriptorSet) (jsonpb.AnyResolver, error) {
+	var fileDescriptors []*protoreflectdesc.FileDescriptor
+	for _, fileDescriptorSet := range fileDescriptorSets {
+		for _, fileDescriptorProto := range fileDescriptorSet.File {
+			iFileDescriptorSet, err := desc.SortFileDescriptorSet(fileDescriptorSet, fileDescriptorProto)
+			if err != nil {
+				return nil, err
+			}
+			fileDescriptor, err := protoreflectdesc.CreateFileDescriptorFromSet(iFileDescriptorSet)
+			if err != nil {
+				return nil, err
+			}
+			fileDescriptors = append(fileDescriptors, fileDescriptor)
+		}
+	}
+	return dynamic.AnyResolver(
+		dynamic.NewMessageFactoryWithDefaults(),
+		fileDescriptors...,
+	), nil
+}
+
+func getDescriptorSourceForMethod(fileDescriptorSets []*descriptor.FileDescriptorSet, method string) (grpcurl.DescriptorSource, error) {
 	servicePath, err := getServiceForMethod(method)
 	if err != nil {
 		return nil, err

internal/grpc/invocation_event_handler.go

@@ -33,19 +33,21 @@ import (
 	"google.golang.org/grpc/status"
 )
 
-var jsonpbMarshaler = &jsonpb.Marshaler{}
-
 var _ grpcurl.InvocationEventHandler = &invocationEventHandler{}
 
 type invocationEventHandler struct {
-	output  io.Writer
-	logger  *zap.Logger
-	details bool
-	err     error
+	jsonpbMarshaler *jsonpb.Marshaler
+	output          io.Writer
+	logger          *zap.Logger
+	details         bool
+	err             error
 }
 
-func newInvocationEventHandler(output io.Writer, logger *zap.Logger, details bool) *invocationEventHandler {
+func newInvocationEventHandler(anyResolver jsonpb.AnyResolver, output io.Writer, logger *zap.Logger, details bool) *invocationEventHandler {
 	return &invocationEventHandler{
+		jsonpbMarshaler: &jsonpb.Marshaler{
+			AnyResolver: anyResolver,
+		},
 		output:  output,
 		logger:  logger,
 		details: details,
@@ -90,7 +92,7 @@ func (i *invocationEventHandler) printProtoMessage(input proto.Message, detailsK
 	if input == nil {
 		return
 	}
-	s, err := jsonpbMarshaler.MarshalToString(input)
+	s, err := i.jsonpbMarshaler.MarshalToString(input)
 	if err != nil {
 		i.logger.Error("marshal error", zap.Error(err))
 		return