Add protoset support. Fixes #5. (#10)

Add protoset support. Fixes #5. (#10)

README.md

@@ -24,8 +24,10 @@ Download a prebuilt executable binary from the [releases page](https://github.co
 Usage: grpcannon [options...] <host>
 Options:
   -proto	The protocol buffer file.
+  -protoset The compiled protoset file. Alternative to proto. -proto takes precedence.
   -call		A fully-qualified method name in 'service/method' or 'service.method' format.
   -cert		The file containing the CA root cert file.
+  -cname	an override of the expect Server Cname presented by the server.
 
   -c  Number of requests to run concurrently. Total number of requests cannot
 	  be smaller than the concurrency level. Default is 50.
@@ -81,6 +83,22 @@ grpcannon -proto ./greeter.proto -call helloworld.Greeter.SayHelloCS -d '[{"name
 
 If a single object is given for data it is sent as every message.
 
+We can also use `.protoset` files which can bundle multiple protoco buffer files into one binary file.
+
+Create a protoset
+
+```
+protoc --proto_path=. --descriptor_set_out=bundle.protoset *.proto
+```
+
+And then use it as input to `grpcannon` with `-protoset` option:
+
+```
+./grpcannon -protoset ./bundle.protoset -call helloworld.Greeter.SayHello -d '{"name":"Bob"}' -n 1000 -c 10 0.0.0.0:50051
+```
+
+Note that only one of `-proto` or `-protoset` options will be used. `-proto` takes precedence.
+
 Example `grpcannon.json`
 
 ```json

cmd/grpcannon/main.go

@@ -13,16 +13,18 @@ import (
 	"github.com/bojand/grpcannon/config"
 	"github.com/bojand/grpcannon/printer"
 	"github.com/bojand/grpcannon/protodesc"
+	"github.com/jhump/protoreflect/desc"
 )
 
 var (
 	// set by goreleaser with -ldflags="-X main.version=..."
 	version = "dev"
 
-	proto = flag.String("proto", "", `The .proto file.`)
-	call  = flag.String("call", "", `A fully-qualified symbol name.`)
-	cert  = flag.String("cert", "", "Client certificate file. If Omitted insecure is used.")
-	cname = flag.String("cname", "", "Server Cert CName Override - useful for self signed certs")
+	proto    = flag.String("proto", "", `The .proto file.`)
+	protoset = flag.String("protoset", "", `The .protoset file.`)
+	call     = flag.String("call", "", `A fully-qualified symbol name.`)
+	cert     = flag.String("cert", "", "Client certificate file. If Omitted insecure is used.")
+	cname    = flag.String("cname", "", "Server Cert CName Override - useful for self signed certs")
 
 	c = flag.Int("c", 50, "Number of requests to run concurrently.")
 	n = flag.Int("n", 200, "Number of requests to run. Default is 200.")
@@ -53,6 +55,7 @@ var (
 var usage = `Usage: grpcannon [options...] <host>
 Options:
   -proto	The protocol buffer file.
+  -protoset The compiled protoset file. Alternative to proto. -proto takes precedence.
   -call		A fully-qualified method name in 'service/method' or 'service.method' format.
   -cert		The file containing the CA root cert file.
   -cname	an override of the expect Server Cname presented by the server.
@@ -119,7 +122,7 @@ func main() {
 			iPaths = strings.Split(pathsTrimmed, ",")
 		}
 
-		cfg, err = config.New(*proto, *call, *cert, *cname, *n, *c, *q, *z, *t,
+		cfg, err = config.New(*proto, *protoset, *call, *cert, *cname, *n, *c, *q, *z, *t,
 			*data, *dataPath, *md, *mdPath, *output, *format, host, *ct, *kt, *cpus, iPaths)
 		if err != nil {
 			errAndExit(err.Error())
@@ -157,15 +160,15 @@ func usageAndExit(msg string) {
 }
 
 func runTest(config *config.Config) (*grpcannon.Report, error) {
-	mtd, err := protodesc.GetMethodDesc(config.Call, config.Proto, config.ImportPaths)
+	mtd, err := getMethodDesc(config)
 	if err != nil {
 		return nil, err
 	}
 
 	opts := &grpcannon.Options{
 		Host:          config.Host,
 		Cert:          config.Cert,
-		CName:		   config.CName,
+		CName:         config.CName,
 		N:             config.N,
 		C:             config.C,
 		QPS:           config.QPS,
@@ -199,3 +202,11 @@ func runTest(config *config.Config) (*grpcannon.Report, error) {
 
 	return reqr.Run()
 }
+
+func getMethodDesc(config *config.Config) (*desc.MethodDescriptor, error) {
+	if config.Proto != "" {
+		return protodesc.GetMethodDescFromProto(config.Call, config.Proto, config.ImportPaths)
+	} else {
+		return protodesc.GetMethodDescFromProtoSet(config.Call, config.Protoset)
+	}
+}

config/config.go

@@ -15,9 +15,10 @@ import (
 // Config for the run.
 type Config struct {
 	Proto         string             `json:"proto"`
+	Protoset      string             `json:"protoset"`
 	Call          string             `json:"call"`
 	Cert          string             `json:"cert"`
-	CName		  string			 `json:"cName"`
+	CName         string             `json:"cName"`
 	N             int                `json:"n"`
 	C             int                `json:"c"`
 	QPS           int                `json:"q"`
@@ -36,16 +37,17 @@ type Config struct {
 	ImportPaths   []string           `json:"i,omitempty"`
 }
 
-// NewConfig creates a new config
-func New(proto, call, cert, cName string, n, c, qps int, z time.Duration, timeout int,
+// New creates a new config
+func New(proto, protoset, call, cert, cName string, n, c, qps int, z time.Duration, timeout int,
 	data, dataPath, metadata, mdPath, output, format, host string,
 	dialTimout, keepaliveTime, cpus int, importPaths []string) (*Config, error) {
 
 	cfg := &Config{
 		Proto:         proto,
+		Protoset:      protoset,
 		Call:          call,
 		Cert:          cert,
-		CName:		   cName,
+		CName:         cName,
 		N:             n,
 		C:             c,
 		QPS:           qps,
@@ -110,12 +112,18 @@ func (c *Config) Default() {
 
 // Validate the config
 func (c *Config) Validate() error {
-	if err := requiredString(c.Proto); err != nil {
-		return errors.Wrap(err, "proto")
+	if strings.TrimSpace(c.Proto) == "" && strings.TrimSpace(c.Protoset) == "" {
+		return errors.New("Proto or Protoset required")
 	}
 
-	if filepath.Ext(c.Proto) != ".proto" {
-		return errors.Errorf(fmt.Sprintf("proto: must have .proto extension"))
+	if strings.TrimSpace(c.Proto) != "" {
+		if filepath.Ext(c.Proto) != ".proto" {
+			return errors.Errorf(fmt.Sprintf("proto: must have .proto extension"))
+		}
+	} else {
+		if filepath.Ext(c.Protoset) != ".protoset" {
+			return errors.Errorf(fmt.Sprintf("protoset: must have .protoset extension"))
+		}
 	}
 
 	if err := requiredString(c.Call); err != nil {

config/config_test.go

@@ -11,7 +11,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-const expected = `{"proto":"asdf","call":"","cert":"","cName":"","n":0,"c":0,"q":0,"t":0,"D":"","M":"","o":"","O":"oval","host":"","T":0,"L":0,"cpus":0,"z":"4h30m0s"}`
+const expected = `{"proto":"asdf","protoset":"","call":"","cert":"","cName":"","n":0,"c":0,"q":0,"t":0,"D":"","M":"","o":"","O":"oval","host":"","T":0,"L":0,"cpus":0,"z":"4h30m0s"}`
 
 func TestConfig_MarshalJSON(t *testing.T) {
 	z, _ := time.ParseDuration("4h30m")
@@ -260,7 +260,7 @@ func TestConfig_ReadConfig(t *testing.T) {
 				Call:          "mycall",
 				Data:          data,
 				Cert:          "mycert",
-				CName: 		   "localhost",
+				CName:         "localhost",
 				N:             200,
 				C:             50,
 				QPS:           0,
@@ -299,7 +299,7 @@ func TestConfig_ReadConfig(t *testing.T) {
 				Call:          "mycall",
 				Data:          data,
 				Cert:          "mycert",
-				CName: 		   "localhost",
+				CName:         "localhost",
 				N:             200,
 				C:             50,
 				QPS:           0,
@@ -331,7 +331,7 @@ func TestConfig_Validate(t *testing.T) {
 	t.Run("missing proto", func(t *testing.T) {
 		c := &Config{}
 		err := c.Validate()
-		assert.Equal(t, "proto: is required", err.Error())
+		assert.Equal(t, "Proto or Protoset required", err.Error())
 	})
 
 	t.Run("invalid proto", func(t *testing.T) {

data_test.go

@@ -159,15 +159,15 @@ func TestData_isMapData(t *testing.T) {
 }
 
 func TestData_createPayloads(t *testing.T) {
-	mtdUnary, err := protodesc.GetMethodDesc(
+	mtdUnary, err := protodesc.GetMethodDescFromProto(
 		"helloworld.Greeter.SayHello",
 		"./testdata/greeter.proto",
 		nil)
 
 	assert.NoError(t, err)
 	assert.NotNil(t, mtdUnary)
 
-	mtdClientStreaming, err := protodesc.GetMethodDesc(
+	mtdClientStreaming, err := protodesc.GetMethodDescFromProto(
 		"helloworld.Greeter.SayHelloCS",
 		"./testdata/greeter.proto",
 		nil)

protodesc/protodesc.go

@@ -2,16 +2,19 @@ package protodesc
 
 import (
 	"fmt"
+	"io/ioutil"
 	"path/filepath"
 	"strings"
 
+	"github.com/golang/protobuf/proto"
+	"github.com/golang/protobuf/protoc-gen-go/descriptor"
 	"github.com/jhump/protoreflect/desc"
 	"github.com/jhump/protoreflect/desc/protoparse"
 )
 
-// GetMethodDesc gets method descitor for the given call symbol in the file given my protoPath
+// GetMethodDescFromProto gets method descritor for the given call symbol from proto file given my path proto
 // imports is used for import paths in parsing the proto file
-func GetMethodDesc(call, proto string, imports []string) (*desc.MethodDescriptor, error) {
+func GetMethodDescFromProto(call, proto string, imports []string) (*desc.MethodDescriptor, error) {
 	p := &protoparse.Parser{ImportPaths: imports}
 
 	filename := proto
@@ -26,19 +29,57 @@ func GetMethodDesc(call, proto string, imports []string) (*desc.MethodDescriptor
 
 	fileDesc := fds[0]
 
+	files := map[string]*desc.FileDescriptor{}
+	files[fileDesc.GetName()] = fileDesc
+
+	return getMethodDesc(call, files)
+}
+
+// GetMethodDescFromProtoSet gets method descritor for the given call symbol from protoset file given my path protoset
+func GetMethodDescFromProtoSet(call, protoset string) (*desc.MethodDescriptor, error) {
+	b, err := ioutil.ReadFile(protoset)
+	if err != nil {
+		return nil, fmt.Errorf("could not load protoset file %q: %v", protoset, err)
+	}
+
+	var fds descriptor.FileDescriptorSet
+	err = proto.Unmarshal(b, &fds)
+	if err != nil {
+		return nil, fmt.Errorf("could not parse contents of protoset file %q: %v", protoset, err)
+	}
+
+	unresolved := map[string]*descriptor.FileDescriptorProto{}
+	for _, fd := range fds.File {
+		unresolved[fd.GetName()] = fd
+	}
+	resolved := map[string]*desc.FileDescriptor{}
+	for _, fd := range fds.File {
+		_, err := resolveFileDescriptor(unresolved, resolved, fd.GetName())
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return getMethodDesc(call, resolved)
+}
+
+func getMethodDesc(call string, files map[string]*desc.FileDescriptor) (*desc.MethodDescriptor, error) {
 	svc, mth := parseSymbol(call)
 	if svc == "" || mth == "" {
 		return nil, fmt.Errorf("given method name %q is not in expected format: 'service/method' or 'service.method'", call)
 	}
 
-	dsc := fileDesc.FindSymbol(svc)
+	dsc, err := findServiceSymbol(files, svc)
+	if err != nil {
+		return nil, err
+	}
 	if dsc == nil {
-		return nil, fmt.Errorf("target server does not expose service %q", svc)
+		return nil, fmt.Errorf("cannot find service %q", svc)
 	}
 
 	sd, ok := dsc.(*desc.ServiceDescriptor)
 	if !ok {
-		return nil, fmt.Errorf("target server does not expose service %q", svc)
+		return nil, fmt.Errorf("cannot find service %q", svc)
 	}
 
 	mtd := sd.FindMethodByName(mth)
@@ -49,6 +90,39 @@ func GetMethodDesc(call, proto string, imports []string) (*desc.MethodDescriptor
 	return mtd, nil
 }
 
+func resolveFileDescriptor(unresolved map[string]*descriptor.FileDescriptorProto, resolved map[string]*desc.FileDescriptor, filename string) (*desc.FileDescriptor, error) {
+	if r, ok := resolved[filename]; ok {
+		return r, nil
+	}
+	fd, ok := unresolved[filename]
+	if !ok {
+		return nil, fmt.Errorf("no descriptor found for %q", filename)
+	}
+	deps := make([]*desc.FileDescriptor, 0, len(fd.GetDependency()))
+	for _, dep := range fd.GetDependency() {
+		depFd, err := resolveFileDescriptor(unresolved, resolved, dep)
+		if err != nil {
+			return nil, err
+		}
+		deps = append(deps, depFd)
+	}
+	result, err := desc.CreateFileDescriptor(fd, deps...)
+	if err != nil {
+		return nil, err
+	}
+	resolved[filename] = result
+	return result, nil
+}
+
+func findServiceSymbol(resolved map[string]*desc.FileDescriptor, fullyQualifiedName string) (desc.Descriptor, error) {
+	for _, fd := range resolved {
+		if dsc := fd.FindSymbol(fullyQualifiedName); dsc != nil {
+			return dsc, nil
+		}
+	}
+	return nil, fmt.Errorf("cannot find service %q", fullyQualifiedName)
+}
+
 func parseSymbol(svcAndMethod string) (string, string) {
 	pos := strings.LastIndex(svcAndMethod, "/")
 	if pos < 0 {