Support for automatically hadling both Bearer and Cookie auth

sillen102 · Mar 11, 2024 · 858b238 · 858b238
1 parent 2283ec8
commit 858b238
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 52 deletions.
diff --git a/chiJwk/jwk_auth.go b/chiJwk/jwk_auth.go
@@ -17,13 +17,12 @@ const JwtTokenKey = "jwt-token"
 
 // JwkAuthOptions is the struct for the jwk auth middleware.
 type JwkAuthOptions struct {
-    AuthenticationType AuthenticationType
-    CookieOptions      CookieOptions
-    JwkSet             jwk.Set
-    Issuer             string
-    IssuerJwkUrl       string
-    Filter             Filter
-    CreateToken        func(claims map[string]interface{}) (Token, error)
+    CookieOptions CookieOptions
+    JwkSet        jwk.Set
+    Issuer        string
+    IssuerJwkUrl  string
+    Filter        Filter
+    CreateToken   func(claims map[string]interface{}) (Token, error)
 }
 
 type AuthenticationType int
@@ -32,11 +31,6 @@ type CookieOptions struct {
     Name string
 }
 
-const (
-    Cookie AuthenticationType = iota
-    Bearer
-)
-
 type Token interface {
     Roles() []string
     Scopes() []string
@@ -68,23 +62,15 @@ func NewJwkOptions(issuer string, jwksUrl string) (*JwkAuthOptions, error) {
     }
 
     return &JwkAuthOptions{
-        AuthenticationType: Cookie,
-        CookieOptions:      CookieOptions{Name: "access-token"},
-        JwkSet:             jwksSet,
-        Issuer:             issuer,
-        IssuerJwkUrl:       jwksUrl,
-        Filter:             DefaultFilter{FilterRoles: make([]string, 0), FilterScopes: make([]string, 0)},
-        CreateToken:        CreateTokenFromClaims[Token],
+        CookieOptions: CookieOptions{Name: "access-token"},
+        JwkSet:        jwksSet,
+        Issuer:        issuer,
+        IssuerJwkUrl:  jwksUrl,
+        Filter:        DefaultFilter{FilterRoles: make([]string, 0), FilterScopes: make([]string, 0)},
+        CreateToken:   CreateTokenFromClaims[Token],
     }, nil
 }
 
-// WithAuthenticationType sets the authentication type option that determines how the token
-// is extracted from the request.
-func (options *JwkAuthOptions) WithAuthenticationType(authenticationType AuthenticationType) *JwkAuthOptions {
-    options.AuthenticationType = authenticationType
-    return options
-}
-
 // WithCookieOptions sets the cookie options that determines how the cookie is extracted from the request.
 func (options *JwkAuthOptions) WithCookieOptions(cookieOptions CookieOptions) *JwkAuthOptions {
     options.CookieOptions = cookieOptions
@@ -127,33 +113,21 @@ func (options *JwkAuthOptions) AuthMiddleware(filter ...Filter) func(next http.H
         options.JwkSet = jwksSet
     }
 
-
     return func(next http.Handler) http.Handler {
 
         fn := func(w http.ResponseWriter, r *http.Request) {
             ctx := r.Context()
             log := FromCtx(ctx)
 
             // Get the token from the request
-            var tokenStr string
-            switch options.AuthenticationType {
-            case Bearer:
-                token, err := request.AuthorizationHeaderExtractor.ExtractToken(r)
-                if err != nil {
-                    log.Debug(fmt.Sprintf("invalid authorization header %v", err))
-                    w.WriteHeader(http.StatusUnauthorized)
-                    return
-                }
-                tokenStr = token
-            case Cookie:
-                // Get the access token from the cookie
-                accessCookie, err := r.Cookie(options.CookieOptions.Name)
-                if err != nil {
-                    log.Debug(fmt.Sprintf("access token cookie required %v", err))
+            tokenStr := options.extractTokenFromCookie(r)
+            if tokenStr == "" {
+                tokenStr = options.extractTokenFromHeader(r)
+                if tokenStr == "" {
+                    log.Debug("no token found in request")
                     w.WriteHeader(http.StatusUnauthorized)
                     return
                 }
-                tokenStr = accessCookie.Value
             }
 
             // Parse and verify the token
@@ -217,6 +191,22 @@ func (options *JwkAuthOptions) AuthMiddleware(filter ...Filter) func(next http.H
     }
 }
 
+func (options *JwkAuthOptions) extractTokenFromHeader(r *http.Request) string {
+    token, err := request.AuthorizationHeaderExtractor.ExtractToken(r)
+    if err != nil {
+        return ""
+    }
+    return token
+}
+
+func (options *JwkAuthOptions) extractTokenFromCookie(r *http.Request) string {
+    accessCookie, err := r.Cookie(options.CookieOptions.Name)
+    if err != nil {
+        return ""
+    }
+    return accessCookie.Value
+}
+
 // GetClaims extracts the token claims from the context into the provided object.
 func GetClaims(ctx context.Context, token Token) error {
     claims, ok := ctx.Value(JwtTokenKey).(map[string]interface{})

diff --git a/chiJwk/jwk_auth_test.go b/chiJwk/jwk_auth_test.go
@@ -46,11 +46,10 @@ func TestAuthMiddleware_Cookie(t *testing.T) {
 
     // Create a mock JwkAuthOptions
     jwkAuthOptions := &chiJwk.JwkAuthOptions{
-        AuthenticationType: chiJwk.Cookie,
-        CookieOptions:      chiJwk.CookieOptions{Name: "access-token"},
-        JwkSet:             jwkSet,
-        Issuer:             issuerValue,
-        Filter:             chiJwk.DefaultFilter{FilterRoles: make([]string, 0), FilterScopes: make([]string, 0)},
+        CookieOptions: chiJwk.CookieOptions{Name: "access-token"},
+        JwkSet:        jwkSet,
+        Issuer:        issuerValue,
+        Filter:        chiJwk.DefaultFilter{FilterRoles: make([]string, 0), FilterScopes: make([]string, 0)},
         CreateToken: func(claims map[string]interface{}) (chiJwk.Token, error) {
             var token keycloak.JwtToken
             err := mapstructure.Decode(claims, &token)
@@ -180,10 +179,9 @@ func TestAuthMiddleware_BearerToken(t *testing.T) {
 
     // Create a mock JwkAuthOptions
     jwkAuthOptions := &chiJwk.JwkAuthOptions{
-        AuthenticationType: chiJwk.Bearer,
-        JwkSet:             jwkSet,
-        Issuer:             issuerValue,
-        Filter:             chiJwk.DefaultFilter{FilterRoles: make([]string, 0), FilterScopes: make([]string, 0)},
+        JwkSet: jwkSet,
+        Issuer: issuerValue,
+        Filter: chiJwk.DefaultFilter{FilterRoles: make([]string, 0), FilterScopes: make([]string, 0)},
         CreateToken: func(claims map[string]interface{}) (chiJwk.Token, error) {
             var token keycloak.JwtToken
             err := mapstructure.Decode(claims, &token)