add WriteProtoset method, for creating protoset during symbol resolution

cmd/grpcurl/grpcurl.go

@@ -59,13 +59,14 @@ var (
 	rpcHeaders    multiString
 	reflHeaders   multiString
 	expandHeaders = flags.Bool("expand-headers", false, prettify(`
-		If set, headers may use '${NAME}' syntax to reference environment variables.
-		These will be expanded to the actual environment variable value before
-		sending to the server. For example, if there is an environment variable
-		defined like FOO=bar, then a header of 'key: ${FOO}' would expand to 'key: bar'.
-		This applies to -H, -rpc-header, and -reflect-header options. No other
-		expansion/escaping is performed. This can be used to supply
-		credentials/secrets without having to put them in command-line arguments.`))
+		If set, headers may use '${NAME}' syntax to reference environment
+		variables. These will be expanded to the actual environment variable
+		value before sending to the server. For example, if there is an
+		environment variable defined like FOO=bar, then a header of
+		'key: ${FOO}' would expand to 'key: bar'. This applies to -H,
+		-rpc-header, and -reflect-header options. No other expansion/escaping is
+		performed. This can be used to supply credentials/secrets without having
+		to put them in command-line arguments.`))
 	authority = flags.String("authority", "", prettify(`
 		Value of :authority pseudo-header to be use with underlying HTTP/2
 		requests. It defaults to the given address.`))
@@ -101,6 +102,13 @@ var (
 		will accept. If not specified, defaults to 4,194,304 (4 megabytes).`))
 	emitDefaults = flags.Bool("emit-defaults", false, prettify(`
 		Emit default values for JSON-encoded responses.`))
+	protosetOut = flags.String("protoset-out", "", prettify(`
+		The name of a file to be written that will contain a FileDescriptorSet
+		proto. With the list and describe verbs, the listed or described
+		elements, and their transitive dependencies, will be written to the
+		named file if this option is given. When invoking an RPC and this option
+		is given, the method being invoked and its transitive dependencies will
+		be included in the output file.`))
 	msgTemplate = flags.Bool("msg-template", false, prettify(`
 		When describing messages, show a template of input data.`))
 	verbose = flags.Bool("v", false, prettify(`
@@ -391,6 +399,9 @@ func main() {
 					fmt.Printf("%s\n", svc)
 				}
 			}
+			if err := writeProtoset(descSource, svcs...); err != nil {
+				fail(err, "Failed to write protoset to %s", *protosetOut)
+			}
 		} else {
 			methods, err := grpcurl.ListMethods(descSource, symbol)
 			if err != nil {
@@ -403,6 +414,9 @@ func main() {
 					fmt.Printf("%s\n", m)
 				}
 			}
+			if err := writeProtoset(descSource, symbol); err != nil {
+				fail(err, "Failed to write protoset to %s", *protosetOut)
+			}
 		}
 
 	} else if describe {
@@ -503,6 +517,9 @@ func main() {
 				fmt.Println(str)
 			}
 		}
+		if err := writeProtoset(descSource, symbols...); err != nil {
+			fail(err, "Failed to write protoset to %s", *protosetOut)
+		}
 
 	} else {
 		// Invoke an RPC
@@ -619,3 +636,15 @@ func fail(err error, msg string, args ...interface{}) {
 		exit(2)
 	}
 }
+
+func writeProtoset(descSource grpcurl.DescriptorSource, symbols ...string) error {
+	if *protosetOut == "" {
+		return nil
+	}
+	f, err := os.Create(*protosetOut)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+	return grpcurl.WriteProtoset(f, descSource, symbols...)
+}

desc_source.go

@@ -3,6 +3,7 @@ package grpcurl
 import (
 	"errors"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"sync"
 
@@ -251,3 +252,53 @@ func reflectionSupport(err error) error {
 	}
 	return err
 }
+
+// WriteProtoset will use the given descriptor source to resolve all of the given
+// symbols and write a proto file descriptor set with their definitions to the
+// given output. The output will include descriptors for all files in which the
+// symbols are defined as well as their transitive dependencies.
+func WriteProtoset(out io.Writer, descSource DescriptorSource, symbols ...string) error {
+	// compute set of file descriptors
+	filenames := make([]string, 0, len(symbols))
+	fds := make(map[string]*desc.FileDescriptor, len(symbols))
+	for _, sym := range symbols {
+		d, err := descSource.FindSymbol(sym)
+		if err != nil {
+			return fmt.Errorf("failed to find descriptor for %q: %v", sym, err)
+		}
+		fd := d.GetFile()
+		if _, ok := fds[fd.GetName()]; !ok {
+			fds[fd.GetName()] = fd
+			filenames = append(filenames, fd.GetName())
+		}
+	}
+	// now expand that to include transitive dependencies in topologically sorted
+	// order (such that file always appears after its dependencies)
+	expandedFiles := make(map[string]struct{}, len(fds))
+	allFilesSlice := make([]*descpb.FileDescriptorProto, 0, len(fds))
+	for _, filename := range filenames {
+		addFilesToSet(expandedFiles, &allFilesSlice, fds[filename])
+	}
+	// now we can serialize to file
+	b, err := proto.Marshal(&descpb.FileDescriptorSet{File: allFilesSlice})
+	if err != nil {
+		return fmt.Errorf("failed to serialize file descriptor set: %v", err)
+	}
+	if _, err := out.Write(b); err != nil {
+		return fmt.Errorf("failed to write file descriptor set: %v", err)
+	}
+	return nil
+}
+
+func addFilesToSet(seen map[string]struct{}, fds *[]*descpb.FileDescriptorProto, fd *desc.FileDescriptor) {
+	if _, ok := seen[fd.GetName()]; ok {
+		// already seen this one
+		return
+	}
+	seen[fd.GetName()] = struct{}{}
+	// add all dependencies first
+	for _, dep := range fd.GetDependencies() {
+		addFilesToSet(seen, fds, dep)
+	}
+	*fds = append(*fds, fd.AsFileDescriptorProto())
+}