diff --git a/cmd/grpcurl/grpcurl.go b/cmd/grpcurl/grpcurl.go index b05c34c..aaa3bbf 100644 --- a/cmd/grpcurl/grpcurl.go +++ b/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...) +} diff --git a/desc_source.go b/desc_source.go index c23ae3d..635ddef 100644 --- a/desc_source.go +++ b/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 { + allFilesSlice = addFilesToSet(allFilesSlice, expandedFiles, 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(allFiles []*descpb.FileDescriptorProto, expanded map[string]struct{}, fd *desc.FileDescriptor) []*descpb.FileDescriptorProto { + if _, ok := expanded[fd.GetName()]; ok { + // already seen this one + return allFiles + } + expanded[fd.GetName()] = struct{}{} + // add all dependencies first + for _, dep := range fd.GetDependencies() { + allFiles = addFilesToSet(allFiles, expanded, dep) + } + return append(allFiles, fd.AsFileDescriptorProto()) +} diff --git a/desc_source_test.go b/desc_source_test.go new file mode 100644 index 0000000..cd58357 --- /dev/null +++ b/desc_source_test.go @@ -0,0 +1,62 @@ +package grpcurl + +import ( + "bytes" + "io/ioutil" + "testing" + + "github.com/golang/protobuf/proto" + "github.com/golang/protobuf/protoc-gen-go/descriptor" +) + +func TestWriteProtoset(t *testing.T) { + exampleProtoset, err := loadProtoset("./testing/example.protoset") + if err != nil { + t.Fatalf("failed to load example.protoset: %v", err) + } + testProtoset, err := loadProtoset("./testing/test.protoset") + if err != nil { + t.Fatalf("failed to load test.protoset: %v", err) + } + + mergedProtoset := &descriptor.FileDescriptorSet{ + File: append(exampleProtoset.File, testProtoset.File...), + } + + descSrc, err := DescriptorSourceFromFileDescriptorSet(mergedProtoset) + if err != nil { + t.Fatalf("failed to create descriptor source: %v", err) + } + + checkWriteProtoset(t, descSrc, exampleProtoset, "TestService") + checkWriteProtoset(t, descSrc, testProtoset, "grpc.testing.TestService") + checkWriteProtoset(t, descSrc, mergedProtoset, "TestService", "grpc.testing.TestService") +} + +func loadProtoset(path string) (*descriptor.FileDescriptorSet, error) { + b, err := ioutil.ReadFile(path) + if err != nil { + return nil, err + } + var protoset descriptor.FileDescriptorSet + if err := proto.Unmarshal(b, &protoset); err != nil { + return nil, err + } + return &protoset, nil +} + +func checkWriteProtoset(t *testing.T, descSrc DescriptorSource, protoset *descriptor.FileDescriptorSet, symbols ...string) { + var buf bytes.Buffer + if err := WriteProtoset(&buf, descSrc, symbols...); err != nil { + t.Fatalf("failed to write protoset: %v", err) + } + + var result descriptor.FileDescriptorSet + if err := proto.Unmarshal(buf.Bytes(), &result); err != nil { + t.Fatalf("failed to unmarshal written protoset: %v", err) + } + + if !proto.Equal(protoset, &result) { + t.Fatalf("written protoset not equal to input:\nExpecting: %s\nActual: %s", protoset, &result) + } +}