protoc-gen-swagger/genswagger: add custom Extensions to SecurityScheme

Signed-off-by: Stephan Renatus <[email protected]>

examples/proto/examplepb/a_bit_of_everything.proto

@@ -51,6 +51,31 @@ option (grpc.gateway.protoc_gen_swagger.options.openapiv2_swagger) = {
 				type: TYPE_API_KEY;
 				in: IN_HEADER;
 				name: "X-API-Key";
+				extensions: {
+					key: "x-amazon-apigateway-authtype";
+					value {
+						string_value: "oauth2";
+					}
+				}
+				extensions: {
+					key: "x-amazon-apigateway-authorizer";
+					value {
+						struct_value {
+							fields {
+								key: "type";
+								value {
+									string_value: "token";
+								}
+							}
+							fields {
+								key: "authorizerResultTtlInSeconds";
+								value {
+									number_value: 60;
+								}
+							}
+						}
+					}
+				}
 			}
 		}
 		security: {
@@ -129,13 +154,13 @@ option (grpc.gateway.protoc_gen_swagger.options.openapiv2_swagger) = {
 		}
 	}
 	extensions: {
-		key: "x-grpc-gateway-foo"
+		key: "x-grpc-gateway-foo";
 		value {
-			string_value: "bar"
+			string_value: "bar";
 		}
 	}
 	extensions: {
-		key: "x-grpc-gateway-baz-list"
+		key: "x-grpc-gateway-baz-list";
 		value {
 			list_value: {
 				values: {

examples/proto/examplepb/a_bit_of_everything.swagger.json

@@ -2283,7 +2283,12 @@
     "ApiKeyAuth": {
       "type": "apiKey",
       "name": "X-API-Key",
-      "in": "header"
+      "in": "header",
+      "x-amazon-apigateway-authorizer": {
+        "authorizerResultTtlInSeconds": 60,
+        "type": "token"
+      },
+      "x-amazon-apigateway-authtype": "oauth2"
     },
     "BasicAuth": {
       "type": "basic"

protoc-gen-swagger/genswagger/generator.go

@@ -71,8 +71,17 @@ func fieldName(k string) string {
 	return strings.ReplaceAll(strings.Title(k), "-", "")
 }
 
-// encodeSwagger converts swagger file obj to plugin.CodeGeneratorResponse_File
-func encodeSwagger(file *wrapper) (*plugin.CodeGeneratorResponse_File, error) {
+func (so swaggerObject) MarshalJSON() ([]byte, error) {
+	type alias swaggerObject
+	return extensionMarshalJSON(alias(so), so.extensions)
+}
+
+func (so swaggerSecuritySchemeObject) MarshalJSON() ([]byte, error) {
+	type alias swaggerSecuritySchemeObject
+	return extensionMarshalJSON(alias(so), so.extensions)
+}
+
+func extensionMarshalJSON(so interface{}, extensions []extension) ([]byte, error) {
 	// To append arbitrary keys to the struct we'll render into json,
 	// we're creating another struct that embeds the original one, and
 	// its extra fields:
@@ -88,11 +97,11 @@ func encodeSwagger(file *wrapper) (*plugin.CodeGeneratorResponse_File, error) {
 	fields := []reflect.StructField{
 		reflect.StructField{ // embedded
 			Name:      "Embedded",
-			Type:      reflect.TypeOf(file.swagger),
+			Type:      reflect.TypeOf(so),
 			Anonymous: true,
 		},
 	}
-	for _, ext := range file.swagger.extensions {
+	for _, ext := range extensions {
 		fields = append(fields, reflect.StructField{
 			Name: fieldName(ext.key),
 			Type: reflect.TypeOf(ext.value),
@@ -102,15 +111,19 @@ func encodeSwagger(file *wrapper) (*plugin.CodeGeneratorResponse_File, error) {
 
 	t := reflect.StructOf(fields)
 	s := reflect.New(t).Elem()
-	s.Field(0).Set(reflect.ValueOf(file.swagger))
-	for _, ext := range file.swagger.extensions {
+	s.Field(0).Set(reflect.ValueOf(so))
+	for _, ext := range extensions {
 		s.FieldByName(fieldName(ext.key)).Set(reflect.ValueOf(ext.value))
 	}
+	return json.Marshal(s.Interface())
+}
 
+// encodeSwagger converts swagger file obj to plugin.CodeGeneratorResponse_File
+func encodeSwagger(file *wrapper) (*plugin.CodeGeneratorResponse_File, error) {
 	var formatted bytes.Buffer
 	enc := json.NewEncoder(&formatted)
 	enc.SetIndent("", "  ")
-	if err := enc.Encode(s.Interface()); err != nil {
+	if err := enc.Encode(*file.swagger); err != nil {
 		return nil, err
 	}
 	name := file.fileName

protoc-gen-swagger/genswagger/template.go

@@ -14,6 +14,7 @@ import (
 	"github.com/golang/glog"
 	"github.com/golang/protobuf/jsonpb"
 	"github.com/golang/protobuf/proto"
+	structpb "github.com/golang/protobuf/ptypes/struct"
 	pbdescriptor "github.com/golang/protobuf/protoc-gen-go/descriptor"
 	gogen "github.com/golang/protobuf/protoc-gen-go/generator"
 	"github.com/grpc-ecosystem/grpc-gateway/protoc-gen-grpc-gateway/descriptor"
@@ -1164,6 +1165,13 @@ func applyTemplate(p param) (*swaggerObject, error) {
 						newSecDefValue.Scopes[scopeKey] = scopeDesc
 					}
 				}
+				if secDefValue.Extensions != nil {
+					exts, err := processExtensions(secDefValue.Extensions)
+					if err != nil {
+						return nil, err
+					}
+					newSecDefValue.extensions = exts
+				}
 				s.SecurityDefinitions[secDefKey] = newSecDefValue
 			}
 		}
@@ -1223,18 +1231,10 @@ func applyTemplate(p param) (*swaggerObject, error) {
 		}
 
 		if spb.Extensions != nil {
-			exts := []extension{}
-			for k, v := range spb.Extensions {
-				if !strings.HasPrefix(k, "x-") {
-					return nil, fmt.Errorf("Extension keys need to start with \"x-\": %q", k)
-				}
-				ext, err := (&jsonpb.Marshaler{Indent: "  "}).MarshalToString(v)
-				if err != nil {
-					return nil, err
-				}
-				exts = append(exts, extension{key: k, value: json.RawMessage(ext)})
+			exts, err := processExtensions(spb.Extensions)
+			if err != nil {
+				return nil, err
 			}
-			sort.Slice(exts, func(i, j int) bool { return exts[i].key < exts[j].key })
 			s.extensions = exts
 		}
 
@@ -1249,6 +1249,22 @@ func applyTemplate(p param) (*swaggerObject, error) {
 	return &s, nil
 }
 
+func processExtensions(inputExts map[string]*structpb.Value) ([]extension, error) {
+	exts := []extension{}
+	for k, v := range inputExts {
+		if !strings.HasPrefix(k, "x-") {
+			return nil, fmt.Errorf("Extension keys need to start with \"x-\": %q", k)
+		}
+		ext, err := (&jsonpb.Marshaler{Indent: "  "}).MarshalToString(v)
+		if err != nil {
+			return nil, err
+		}
+		exts = append(exts, extension{key: k, value: json.RawMessage(ext)})
+	}
+	sort.Slice(exts, func(i, j int) bool { return exts[i].key < exts[j].key })
+	return exts, nil
+}
+
 // updateSwaggerDataFromComments updates a Swagger object based on a comment
 // from the proto file.
 //

protoc-gen-swagger/genswagger/template_test.go

@@ -372,7 +372,7 @@ func TestApplyTemplateSimple(t *testing.T) {
 	}
 }
 
-func TestApplyTemplateTopLevelOptions(t *testing.T) {
+func TestApplyTemplateExtensions(t *testing.T) {
 	file := descriptor.File{
 		FileDescriptorProto: &protodescriptor.FileDescriptorProto{
 			SourceCodeInfo: &protodescriptor.SourceCodeInfo{},
@@ -393,7 +393,18 @@ func TestApplyTemplateTopLevelOptions(t *testing.T) {
 	swagger := swagger_options.Swagger{
 		Extensions: map[string]*structpb.Value{
 			"x-foo": &structpb.Value{Kind: &structpb.Value_StringValue{StringValue: "bar"}},
-			"x-bar": &structpb.Value{Kind: &structpb.Value_ListValue{ListValue: &structpb.ListValue{Values: []*structpb.Value{{Kind: &structpb.Value_StringValue{StringValue: "baz"}}}}}},
+			"x-bar": &structpb.Value{Kind: &structpb.Value_ListValue{ListValue: &structpb.ListValue{
+				Values: []*structpb.Value{{Kind: &structpb.Value_StringValue{StringValue: "baz"}}},
+			}}},
+		},
+		SecurityDefinitions: &swagger_options.SecurityDefinitions{
+			Security: map[string]*swagger_options.SecurityScheme{
+				"somescheme": &swagger_options.SecurityScheme{
+					Extensions: map[string]*structpb.Value{
+						"x-security-baz": &structpb.Value{Kind: &structpb.Value_BoolValue{BoolValue: true}},
+					},
+				},
+			},
 		},
 	}
 	err := proto.SetExtension(proto.Message(file.FileDescriptorProto.Options), swagger_options.E_Openapiv2Swagger, &swagger)
@@ -414,6 +425,16 @@ func TestApplyTemplateTopLevelOptions(t *testing.T) {
 	}, result.extensions, "Extensions"; !reflect.DeepEqual(is, want) {
 		t.Errorf("applyTemplate(%#v).%s = %s want to be %s", file, name, is, want)
 	}
+
+	var scheme swaggerSecuritySchemeObject
+	for _, v := range result.SecurityDefinitions {
+		scheme = v
+	}
+	if want, is, name := []extension{
+		{key: "x-security-baz", value: json.RawMessage("true")},
+	}, scheme.extensions, "SecurityScheme.Extensions"; !reflect.DeepEqual(is, want) {
+		t.Errorf("applyTemplate(%#v).%s = %s want to be %s", file, name, is, want)
+	}
 }
 
 func TestApplyTemplateRequestWithoutClientStreaming(t *testing.T) {

protoc-gen-swagger/genswagger/types.go

@@ -83,6 +83,8 @@ type swaggerSecuritySchemeObject struct {
 	AuthorizationURL string              `json:"authorizationUrl,omitempty"`
 	TokenURL         string              `json:"tokenUrl,omitempty"`
 	Scopes           swaggerScopesObject `json:"scopes,omitempty"`
+
+	extensions []extension
 }
 
 // http://swagger.io/specification/#scopesObject