Repository: ogen-go/ogen
Pull request #768 (2023-02-07 02:42:50)

feat(gen): support `apiKey` cookie security


Patch:
diff --git a/gen/_template/security.tmpl b/gen/_template/security.tmpl
index 687196b0..f9d121df 100644
--- a/gen/_template/security.tmpl
+++ b/gen/_template/security.tmpl
@@ -45,6 +45,16 @@ func (s *Server) security{{ $s.Type.Name }}(ctx context.Context, operationName s
 			return ctx, false, nil
 		}
 		value := q.Get(parameterName)
+        {{- else if $s.Kind.IsCookie }}
+	    var value string
+		switch cookie, err := req.Cookie(parameterName); err {
+		case nil:
+			value = cookie.Value
+		case http.ErrNoCookie:
+			return ctx, false, nil
+		default:
+			return nil, false, err
+		}
 		{{- else }}
             {{ errorf "unexpected security %q:%q" $s.Kind $s.Format }}
 		{{- end }}
@@ -103,6 +113,11 @@ func (s *Client) security{{ $s.Type.Name }}(ctx context.Context, operationName s
 		q := req.URL.Query()
 		q.Set({{ quote $s.ParameterName }}, t.APIKey)
 		req.URL.RawQuery = q.Encode()
+        {{- else if $s.Kind.IsCookie }}
+		req.AddCookie(&http.Cookie{
+			Name: {{ quote $s.ParameterName }},
+			Value: t.APIKey,
+		})
 		{{- else }}
             {{ errorf "unexpected security %q:%q" $s.Kind $s.Format }}
 		{{- end }}
diff --git a/gen/gen_security.go b/gen/gen_security.go
index c342fc45..d57da84b 100644
--- a/gen/gen_security.go
+++ b/gen/gen_security.go
@@ -29,7 +29,7 @@ func (g *Generator) generateSecurityAPIKey(s *ir.Security, spec openapi.Security
 		s.Kind = ir.HeaderSecurity
 		vetHeaderParameterName(g.log, s.ParameterName, spec.Security)
 	case "cookie":
-		return nil, &ErrNotImplemented{Name: "cookie security"}
+		s.Kind = ir.CookieSecurity
 	default:
 		return nil, errors.Errorf(`unknown "in" value %q`, in)
 	}
diff --git a/gen/ir/security.go b/gen/ir/security.go
index a011d643..fa51a620 100644
--- a/gen/ir/security.go
+++ b/gen/ir/security.go
@@ -12,6 +12,8 @@ const (
 	QuerySecurity SecurityKind = "query"
 	// HeaderSecurity is HTTP header security parameter. Matches some "http" schemes and "apiKey" with "in" = "header".
 	HeaderSecurity SecurityKind = "header"
+	// CookieSecurity is HTTP cookie security parameter. Matches some "http" schemes and "apiKey" with "in" = "cookie".
+	CookieSecurity SecurityKind = "cookie"
 	// OAuth2Security is special type for OAuth2-based authentication. Matches "oauth2" and "openIdConnect".
 	OAuth2Security SecurityKind = "oauth2"
 )
@@ -26,6 +28,11 @@ func (s SecurityKind) IsHeader() bool {
 	return s == HeaderSecurity
 }
 
+// IsCookie whether s is CookieSecurity.
+func (s SecurityKind) IsCookie() bool {
+	return s == CookieSecurity
+}
+
 // IsOAuth2 whether s is OAuth2Security.
 func (s SecurityKind) IsOAuth2() bool {
 	return s == OAuth2Security


Test patch:
diff --git a/_testdata/positive/security.json b/_testdata/positive/security.json
index cc1bbf38..1a032ab0 100644
--- a/_testdata/positive/security.json
+++ b/_testdata/positive/security.json
@@ -30,6 +30,7 @@
             "queryKey": []
           },
           {
+            "cookieKey": [],
             "headerKey": []
           }
         ],
@@ -73,6 +74,11 @@
         "in": "header",
         "name": "X-Api-Key"
       },
+      "cookieKey": {
+        "type": "apiKey",
+        "in": "cookie",
+        "name": "api_key"
+      },
       "basicAuth": {
         "type": "http",
         "scheme": "basic"
diff --git a/internal/integration/security_test.go b/internal/integration/security_test.go
index c61aef5e..b8300481 100644
--- a/internal/integration/security_test.go
+++ b/internal/integration/security_test.go
@@ -19,6 +19,7 @@ type testSecurity struct {
 	bearerToken api.BearerToken
 	headerKey   api.HeaderKey
 	queryKey    api.QueryKey
+	cookieKey   api.CookieKey
 }
 
 func (t *testSecurity) OptionalSecurity(ctx context.Context) error {
@@ -63,11 +64,19 @@ func (t *testSecurity) HandleQueryKey(ctx context.Context, operationName string,
 	return context.WithValue(ctx, tokenKey("QueryKey"), v), nil
 }
 
+func (t *testSecurity) HandleCookieKey(ctx context.Context, operationName string, v api.CookieKey) (context.Context, error) {
+	if v != t.cookieKey {
+		return nil, errors.Errorf("invalid api key: %q", v.APIKey)
+	}
+	return context.WithValue(ctx, tokenKey("CookieKey"), v), nil
+}
+
 type testSecuritySource struct {
 	basicAuth   *api.BasicAuth
 	bearerToken *api.BearerToken
 	headerKey   *api.HeaderKey
 	queryKey    *api.QueryKey
+	cookieKey   *api.CookieKey
 }
 
 func (t *testSecuritySource) BasicAuth(ctx context.Context, operationName string) (r api.BasicAuth, _ error) {
@@ -98,12 +107,20 @@ func (t *testSecuritySource) QueryKey(ctx context.Context, operationName string)
 	return r, ogenerrors.ErrSkipClientSecurity
 }
 
+func (t *testSecuritySource) CookieKey(ctx context.Context, operationName string) (r api.CookieKey, _ error) {
+	if v := t.cookieKey; v != nil {
+		return *v, nil
+	}
+	return r, ogenerrors.ErrSkipClientSecurity
+}
+
 func TestSecurity(t *testing.T) {
 	h := &testSecurity{
 		basicAuth:   api.BasicAuth{Username: "username", Password: "password"},
 		bearerToken: api.BearerToken{Token: "BearerToken"},
 		headerKey:   api.HeaderKey{APIKey: "HeaderKey"},
 		queryKey:    api.QueryKey{APIKey: "QueryKey"},
+		cookieKey:   api.CookieKey{APIKey: "CookieKey"},
 	}
 	srv, err := api.NewServer(h, h)
 	require.NoError(t, err)
@@ -171,6 +188,7 @@ func TestSecurity(t *testing.T) {
 
 		resp = sendReq(t, "/disjointSecurity", func(r *http.Request) {
 			r.Header.Set("X-Api-Key", h.headerKey.APIKey)
+			r.AddCookie(&http.Cookie{Name: "api_key", Value: h.cookieKey.APIKey})
 		})
 		require.Equal(t, http.StatusOK, resp.StatusCode)
 
@@ -207,6 +225,7 @@ func TestSecurityClientCheck(t *testing.T) {
 		bearerToken: api.BearerToken{Token: "BearerToken"},
 		headerKey:   api.HeaderKey{APIKey: "HeaderKey"},
 		queryKey:    api.QueryKey{APIKey: "QueryKey"},
+		cookieKey:   api.CookieKey{APIKey: "CookieKey"},
 	}
 	srv, err := api.NewServer(h, h)
 	require.NoError(t, err)
@@ -247,9 +266,10 @@ func TestSecurityClientCheck(t *testing.T) {
 	t.Run("DisjointSecurity", test((*api.Client).DisjointSecurity, []testCase{
 		{wantErr: true},
 		{source: testSecuritySource{queryKey: &h.queryKey}, wantErr: true},
+		{source: testSecuritySource{headerKey: &h.headerKey}, wantErr: true},
 
 		{source: testSecuritySource{queryKey: &h.queryKey, basicAuth: &h.basicAuth}, wantErr: false},
-		{source: testSecuritySource{headerKey: &h.headerKey}, wantErr: false},
+		{source: testSecuritySource{headerKey: &h.headerKey, cookieKey: &h.cookieKey}, wantErr: false},
 	}))
 	t.Run("IntersectSecurity", test((*api.Client).IntersectSecurity, []testCase{
 		{wantErr: true},
diff --git a/internal/integration/test_security/oas_client_gen.go b/internal/integration/test_security/oas_client_gen.go
index 246a34f2..bfb9d57b 100644
--- a/internal/integration/test_security/oas_client_gen.go
+++ b/internal/integration/test_security/oas_client_gen.go
@@ -136,11 +136,22 @@ func (c *Client) sendDisjointSecurity(ctx context.Context) (res *DisjointSecurit
 				return res, errors.Wrap(err, "security \"QueryKey\"")
 			}
 		}
+		{
+			stage = "Security:CookieKey"
+			switch err := c.securityCookieKey(ctx, "DisjointSecurity", r); err {
+			case nil:
+				satisfied[0] |= 1 << 2
+			case ogenerrors.ErrSkipClientSecurity:
+				// Skip this security.
+			default:
+				return res, errors.Wrap(err, "security \"CookieKey\"")
+			}
+		}
 		{
 			stage = "Security:HeaderKey"
 			switch err := c.securityHeaderKey(ctx, "DisjointSecurity", r); err {
 			case nil:
-				satisfied[0] |= 1 << 2
+				satisfied[0] |= 1 << 3
 			case ogenerrors.ErrSkipClientSecurity:
 				// Skip this security.
 			default:
@@ -152,7 +163,7 @@ func (c *Client) sendDisjointSecurity(ctx context.Context) (res *DisjointSecurit
 		nextRequirement:
 			for _, requirement := range []bitset{
 				{0b00000011},
-				{0b00000100},
+				{0b00001100},
 			} {
 				for i, mask := range requirement {
 					if satisfied[i]&mask != mask {
diff --git a/internal/integration/test_security/oas_handlers_gen.go b/internal/integration/test_security/oas_handlers_gen.go
index 360b4766..a628efa8 100644
--- a/internal/integration/test_security/oas_handlers_gen.go
+++ b/internal/integration/test_security/oas_handlers_gen.go
@@ -90,6 +90,23 @@ func (s *Server) handleDisjointSecurityRequest(args [0]string, w http.ResponseWr
 				ctx = sctx
 			}
 		}
+		{
+			sctx, ok, err := s.securityCookieKey(ctx, "DisjointSecurity", r)
+			if err != nil {
+				err = &ogenerrors.SecurityError{
+					OperationContext: opErrContext,
+					Security:         "CookieKey",
+					Err:              err,
+				}
+				recordError("Security:CookieKey", err)
+				s.cfg.ErrorHandler(ctx, w, r, err)
+				return
+			}
+			if ok {
+				satisfied[0] |= 1 << 2
+				ctx = sctx
+			}
+		}
 		{
 			sctx, ok, err := s.securityHeaderKey(ctx, "DisjointSecurity", r)
 			if err != nil {
@@ -103,7 +120,7 @@ func (s *Server) handleDisjointSecurityRequest(args [0]string, w http.ResponseWr
 				return
 			}
 			if ok {
-				satisfied[0] |= 1 << 2
+				satisfied[0] |= 1 << 3
 				ctx = sctx
 			}
 		}
@@ -112,7 +129,7 @@ func (s *Server) handleDisjointSecurityRequest(args [0]string, w http.ResponseWr
 		nextRequirement:
 			for _, requirement := range []bitset{
 				{0b00000011},
-				{0b00000100},
+				{0b00001100},
 			} {
 				for i, mask := range requirement {
 					if satisfied[i]&mask != mask {
diff --git a/internal/integration/test_security/oas_schemas_gen.go b/internal/integration/test_security/oas_schemas_gen.go
index 6fce3a9f..808a1f52 100644
--- a/internal/integration/test_security/oas_schemas_gen.go
+++ b/internal/integration/test_security/oas_schemas_gen.go
@@ -41,6 +41,20 @@ func (s *BearerToken) SetToken(val string) {
 	s.Token = val
 }
 
+type CookieKey struct {
+	APIKey string
+}
+
+// GetAPIKey returns the value of APIKey.
+func (s *CookieKey) GetAPIKey() string {
+	return s.APIKey
+}
+
+// SetAPIKey sets the value of APIKey.
+func (s *CookieKey) SetAPIKey(val string) {
+	s.APIKey = val
+}
+
 // DisjointSecurityOK is response for DisjointSecurity operation.
 type DisjointSecurityOK struct{}
 
diff --git a/internal/integration/test_security/oas_security_gen.go b/internal/integration/test_security/oas_security_gen.go
index dc187035..f06b9a41 100644
--- a/internal/integration/test_security/oas_security_gen.go
+++ b/internal/integration/test_security/oas_security_gen.go
@@ -18,6 +18,8 @@ type SecurityHandler interface {
 	HandleBasicAuth(ctx context.Context, operationName string, t BasicAuth) (context.Context, error)
 	// HandleBearerToken handles bearerToken security.
 	HandleBearerToken(ctx context.Context, operationName string, t BearerToken) (context.Context, error)
+	// HandleCookieKey handles cookieKey security.
+	HandleCookieKey(ctx context.Context, operationName string, t CookieKey) (context.Context, error)
 	// HandleHeaderKey handles headerKey security.
 	HandleHeaderKey(ctx context.Context, operationName string, t HeaderKey) (context.Context, error)
 	// HandleQueryKey handles queryKey security.
@@ -69,6 +71,25 @@ func (s *Server) securityBearerToken(ctx context.Context, operationName string,
 	}
 	return rctx, true, err
 }
+func (s *Server) securityCookieKey(ctx context.Context, operationName string, req *http.Request) (context.Context, bool, error) {
+	var t CookieKey
+	const parameterName = "api_key"
+	var value string
+	switch cookie, err := req.Cookie(parameterName); err {
+	case nil:
+		value = cookie.Value
+	case http.ErrNoCookie:
+		return ctx, false, nil
+	default:
+		return nil, false, err
+	}
+	t.APIKey = value
+	rctx, err := s.sec.HandleCookieKey(ctx, operationName, t)
+	if err != nil {
+		return nil, false, err
+	}
+	return rctx, true, err
+}
 func (s *Server) securityHeaderKey(ctx context.Context, operationName string, req *http.Request) (context.Context, bool, error) {
 	var t HeaderKey
 	const parameterName = "X-Api-Key"
@@ -105,6 +126,8 @@ type SecuritySource interface {
 	BasicAuth(ctx context.Context, operationName string) (BasicAuth, error)
 	// BearerToken provides bearerToken security value.
 	BearerToken(ctx context.Context, operationName string) (BearerToken, error)
+	// CookieKey provides cookieKey security value.
+	CookieKey(ctx context.Context, operationName string) (CookieKey, error)
 	// HeaderKey provides headerKey security value.
 	HeaderKey(ctx context.Context, operationName string) (HeaderKey, error)
 	// QueryKey provides queryKey security value.
@@ -133,6 +156,20 @@ func (s *Client) securityBearerToken(ctx context.Context, operationName string,
 	req.Header.Set("Authorization", "Bearer "+t.Token)
 	return nil
 }
+func (s *Client) securityCookieKey(ctx context.Context, operationName string, req *http.Request) error {
+	t, err := s.sec.CookieKey(ctx, operationName)
+	if err != nil {
+		if errors.Is(err, ogenerrors.ErrSkipClientSecurity) {
+			return ogenerrors.ErrSkipClientSecurity
+		}
+		return errors.Wrap(err, "security source \"CookieKey\"")
+	}
+	req.AddCookie(&http.Cookie{
+		Name:  "api_key",
+		Value: t.APIKey,
+	})
+	return nil
+}
 func (s *Client) securityHeaderKey(ctx context.Context, operationName string, req *http.Request) error {
 	t, err := s.sec.HeaderKey(ctx, operationName)
 	if err != nil {