Remove ban on recursive query parameters (#2022)

* implements a max recursive depth check

* bazel is just at root and update build.bazel

* review comments

Author: drewwells

.gitignore

@@ -7,7 +7,7 @@ bazel-genfiles
 bazel-grpc-gateway
 bazel-out
 bazel-testlogs
-.bazelrc
+/.bazelrc
 
 # Go vendor directory
 vendor

internal/descriptor/registry.go

@@ -109,6 +109,9 @@ type Registry struct {
 
 	// omitPackageDoc, if false, causes a package comment to be included in the generated code.
 	omitPackageDoc bool
+
+	// recursiveDepth sets the maximum depth of a field parameter
+	recursiveDepth int
 }
 
 type repeatedFieldSeparator struct {
@@ -134,6 +137,7 @@ func NewRegistry() *Registry {
 		messageOptions: make(map[string]*options.Schema),
 		serviceOptions: make(map[string]*options.Tag),
 		fieldOptions:   make(map[string]*options.JSONSchema),
+		recursiveDepth: 1000,
 	}
 }
 
@@ -356,6 +360,16 @@ func (r *Registry) SetStandalone(standalone bool) {
 	r.standalone = standalone
 }
 
+// SetRecursiveDepth records the max recursion count
+func (r *Registry) SetRecursiveDepth(count int) {
+	r.recursiveDepth = count
+}
+
+// GetRecursiveDepth returns the max recursion count
+func (r *Registry) GetRecursiveDepth() int {
+	return r.recursiveDepth
+}
+
 // ReserveGoPackageAlias reserves the unique alias of go package.
 // If succeeded, the alias will be never used for other packages in generated go files.
 // If failed, the alias is already taken by another package, so you need to use another

protoc-gen-openapiv2/internal/genopenapi/BUILD.bazel

@@ -34,7 +34,10 @@ go_library(
 go_test(
     name = "go_default_test",
     size = "small",
-    srcs = ["template_test.go"],
+    srcs = [
+        "cycle_test.go",
+        "template_test.go",
+    ],
     embed = [":go_default_library"],
     deps = [
         "//internal/descriptor:go_default_library",

protoc-gen-openapiv2/internal/genopenapi/cycle_test.go

@@ -0,0 +1,41 @@
+package genopenapi
+
+import (
+	"testing"
+)
+
+func TestCycle(t *testing.T) {
+	for _, tt := range []struct {
+		max     int
+		attempt int
+		e       bool
+	}{
+		{
+			max:     3,
+			attempt: 3,
+			e:       true,
+		},
+		{
+			max:     5,
+			attempt: 6,
+		},
+		{
+			max:     1000,
+			attempt: 1001,
+		},
+	} {
+
+		c := newCycleChecker(tt.max)
+		var final bool
+		for i := 0; i < tt.attempt; i++ {
+			final = c.Check("a")
+			if !final {
+				break
+			}
+		}
+
+		if final != tt.e {
+			t.Errorf("got: %t wanted: %t", final, tt.e)
+		}
+	}
+}

protoc-gen-openapiv2/internal/genopenapi/generator.go

@@ -36,6 +36,11 @@ type wrapper struct {
 	swagger  *openapiSwaggerObject
 }
 
+type GeneratorOptions struct {
+	Registry       *descriptor.Registry
+	RecursiveDepth int
+}
+
 // New returns a new generator which generates grpc gateway files.
 func New(reg *descriptor.Registry) gen.Generator {
 	return &generator{reg: reg}

protoc-gen-openapiv2/internal/genopenapi/template.go

@@ -120,7 +120,7 @@ func getEnumDefault(enum *descriptor.Enum) string {
 // messageToQueryParameters converts a message to a list of OpenAPI query parameters.
 func messageToQueryParameters(message *descriptor.Message, reg *descriptor.Registry, pathParams []descriptor.Parameter, body *descriptor.Body) (params []openapiParameterObject, err error) {
 	for _, field := range message.Fields {
-		p, err := queryParams(message, field, "", reg, pathParams, body)
+		p, err := queryParams(message, field, "", reg, pathParams, body, reg.GetRecursiveDepth())
 		if err != nil {
 			return nil, err
 		}
@@ -130,17 +130,64 @@ func messageToQueryParameters(message *descriptor.Message, reg *descriptor.Regis
 }
 
 // queryParams converts a field to a list of OpenAPI query parameters recursively through the use of nestedQueryParams.
-func queryParams(message *descriptor.Message, field *descriptor.Field, prefix string, reg *descriptor.Registry, pathParams []descriptor.Parameter, body *descriptor.Body) (params []openapiParameterObject, err error) {
-	return nestedQueryParams(message, field, prefix, reg, pathParams, body, map[string]bool{})
+func queryParams(message *descriptor.Message, field *descriptor.Field, prefix string, reg *descriptor.Registry, pathParams []descriptor.Parameter, body *descriptor.Body, recursiveCount int) (params []openapiParameterObject, err error) {
+	return nestedQueryParams(message, field, prefix, reg, pathParams, body, newCycleChecker(recursiveCount))
+}
+
+type cycleChecker struct {
+	m     map[string]int
+	count int
+}
+
+func newCycleChecker(recursive int) *cycleChecker {
+	return &cycleChecker{
+		m:     make(map[string]int),
+		count: recursive,
+	}
+}
+
+// Check returns whether name is still within recursion
+// toleration
+func (c *cycleChecker) Check(name string) bool {
+	count, ok := c.m[name]
+	count = count + 1
+	isCycle := count > c.count
+
+	if isCycle {
+		return false
+	}
+
+	// provision map entry if not available
+	if !ok {
+		c.m[name] = 1
+		return true
+	}
+
+	c.m[name] = count
+
+	return true
+}
+
+func (c *cycleChecker) Branch() *cycleChecker {
+	copy := &cycleChecker{
+		count: c.count,
+		m:     map[string]int{},
+	}
+
+	for k, v := range c.m {
+		copy.m[k] = v
+	}
+
+	return copy
 }
 
 // nestedQueryParams converts a field to a list of OpenAPI query parameters recursively.
 // This function is a helper function for queryParams, that keeps track of cyclical message references
 //  through the use of
-//      touched map[string]bool
-// If a cycle is discovered, an error is returned, as cyclical data structures aren't allowed
+//      touched map[string]int
+// If a cycle is discovered, an error is returned, as cyclical data structures are dangerous
 //  in query parameters.
-func nestedQueryParams(message *descriptor.Message, field *descriptor.Field, prefix string, reg *descriptor.Registry, pathParams []descriptor.Parameter, body *descriptor.Body, touchedIn map[string]bool) (params []openapiParameterObject, err error) {
+func nestedQueryParams(message *descriptor.Message, field *descriptor.Field, prefix string, reg *descriptor.Registry, pathParams []descriptor.Parameter, body *descriptor.Body, cycle *cycleChecker) (params []openapiParameterObject, err error) {
 	// make sure the parameter is not already listed as a path parameter
 	for _, pathParam := range pathParams {
 		if pathParam.Target == field {
@@ -248,19 +295,15 @@ func nestedQueryParams(message *descriptor.Message, field *descriptor.Field, pre
 	}
 
 	// Check for cyclical message reference:
-	isCycle := touchedIn[*msg.Name]
-	if isCycle {
-		return nil, fmt.Errorf("recursive types are not allowed for query parameters, cycle found on %q", fieldType)
+	isOK := cycle.Check(*msg.Name)
+	if !isOK {
+		return nil, fmt.Errorf("exceeded recursive count (%d) for query parameter %q", cycle.count, fieldType)
 	}
 
 	// Construct a new map with the message name so a cycle further down the recursive path can be detected.
 	// Do not keep anything in the original touched reference and do not pass that reference along.  This will
 	// prevent clobbering adjacent records while recursing.
-	touchedOut := make(map[string]bool)
-	for k, v := range touchedIn {
-		touchedOut[k] = v
-	}
-	touchedOut[*msg.Name] = true
+	touchedOut := cycle.Branch()
 
 	for _, nestedField := range msg.Fields {
 		var fieldName string

protoc-gen-openapiv2/main.go

@@ -33,6 +33,7 @@ var (
 	simpleOperationIDs         = flag.Bool("simple_operation_ids", false, "whether to remove the service prefix in the operationID generation. Can introduce duplicate operationIDs, use with caution.")
 	openAPIConfiguration       = flag.String("openapi_configuration", "", "path to file which describes the OpenAPI Configuration in YAML format")
 	generateUnboundMethods     = flag.Bool("generate_unbound_methods", false, "generate swagger metadata even for RPC methods that have no HttpRule annotation")
+	recursiveDepth             = flag.Int("recursive-depth", 1000, "maximum recursion count allowed for a field type")
 )
 
 // Variables set by goreleaser at build time
@@ -89,6 +90,7 @@ func main() {
 	reg.SetDisableDefaultErrors(*disableDefaultErrors)
 	reg.SetSimpleOperationIDs(*simpleOperationIDs)
 	reg.SetGenerateUnboundMethods(*generateUnboundMethods)
+	reg.SetRecursiveDepth(*recursiveDepth)
 	if err := reg.SetRepeatedPathParamSeparator(*repeatedPathParamSeparator); err != nil {
 		emitError(err)
 		return