fixup! Add support for the anchor redirects

.traefik.yml

@@ -11,7 +11,6 @@ summary: |
   role-based access control, token caching, and more.
 
   The middleware has been tested with Auth0, Logto, Google, and other standard OIDC providers.
-  It includes special handling for Google's OAuth implementation to ensure compatibility.
   It supports various authentication scenarios including:
 
   - Basic authentication with customizable callback and logout URLs
@@ -45,10 +44,6 @@ testData:
     - company.com
     - subsidiary.com
 
-  allowedUsers: # Restricts access to specific email addresses regardless of domain
-    - [email protected]
-    - [email protected]
-
   allowedRolesAndGroups: # Restricts access to users with specific roles or groups (if not provided, no role/group restrictions)
     - guest-endpoints
     - admin
@@ -63,16 +58,6 @@ testData:
     - /public
     - /health
     - /metrics
-    
-  headers: # Custom headers to set with templated values from claims and tokens
-    - name: "X-User-Email"
-      value: "{{.Claims.email}}"
-    - name: "X-User-ID"
-      value: "{{.Claims.sub}}"
-    - name: "Authorization"
-      value: "Bearer {{.AccessToken}}"
-    - name: "X-User-Roles"
-      value: "{{range $i, $e := .Claims.roles}}{{if $i}},{{end}}{{$e}}{{end}}"
 
   # Advanced parameters (usually discovered automatically from provider metadata)
   revocationURL: https://accounts.google.com/revoke # Endpoint for revoking tokens
@@ -157,9 +142,6 @@ configuration:
       Default: ["openid", "profile", "email"]
 
       Include "roles" or similar scope if you need role/group information.
-      Note: For Google OAuth, the middleware automatically handles the 
-      proper authentication parameters and does NOT require the "offline_access"
-      scope (which Google rejects as invalid). See documentation for details.
     required: false
     items:
       type: string
@@ -219,21 +201,6 @@ configuration:
     items:
       type: string
 
-  allowedUsers:
-    type: array
-    description: |
-      Restricts access to specific email addresses.
-      If provided, only users with these exact email addresses will be allowed access,
-      in addition to any domain-level restrictions set by allowedUserDomains.
-      
-      This provides fine-grained control over individual access and can be used
-      together with allowedUserDomains for flexible access control strategies.
-
-      Examples: ["[email protected]", "[email protected]"]
-    required: false
-    items:
-      type: string
-
   allowedRolesAndGroups:
     type: array
     description: |
@@ -276,31 +243,3 @@ configuration:
 
       Default: false
     required: false
-
-  headers:
-    type: array
-    description: |
-      Custom HTTP headers to set with templated values derived from OIDC claims and tokens.
-      Each header has a name and a value template that can access:
-      - {{.Claims.field}} - Access ID token claims (e.g., email, sub, name)
-      - {{.AccessToken}} - The raw access token string
-      - {{.IdToken}} - The raw ID token string
-      - {{.RefreshToken}} - The raw refresh token string
-
-      Templates support Go template syntax including conditionals and iteration.
-      Variable names are case-sensitive - use .Claims not .claims.
-
-      Examples:
-      - name: "X-User-Email", value: "{{.Claims.email}}"
-      - name: "Authorization", value: "Bearer {{.AccessToken}}"
-      - name: "X-User-Roles", value: "{{range $i, $e := .Claims.roles}}{{if $i}},{{end}}{{$e}}{{end}}"
-    required: false
-    items:
-      type: object
-      properties:
-        name:
-          type: string
-          description: The HTTP header name to set
-        value:
-          type: string
-          description: Template string for the header value

LICENSE

@@ -1,21 +0,0 @@
-MIT License
-
-Copyright (c) 2025 Lukasz Raczylo
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.

README.md

@@ -13,7 +13,7 @@ The Traefik OIDC middleware provides a complete OIDC authentication solution wit
 - Rate limiting
 - Excluded paths (public URLs)
 
-The middleware has been tested with Auth0, Logto, Google and other standard OIDC providers. It includes special handling for Google's OAuth implementation.
+The middleware has been tested with Auth0 and Logto, but should work with any standard OIDC provider.
 
 ## Traefik Version Compatibility
 
@@ -69,17 +69,14 @@ The middleware supports the following configuration options:
 | `postLogoutRedirectURI` | The URL to redirect to after logout | `/` | `/logged-out-page` |
 | `scopes` | The OAuth 2.0 scopes to request | `["openid", "profile", "email"]` | `["openid", "email", "profile", "roles"]` |
 | `logLevel` | Sets the logging verbosity | `info` | `debug`, `info`, `error` |
-| `forceHTTPS` | Forces the use of HTTPS for all URLs | `true` | `true`, `false` |
+| | `forceHTTPS` | Forces the use of HTTPS for all URLs | `true` | `true`, `false` |
-| `rateLimit` | Sets the maximum number of requests per second | `100` | `500` |
+| | `rateLimit` | Sets the maximum number of requests per second | `100` | `500` |
-| `excludedURLs` | Lists paths that bypass authentication | none | `["/health", "/metrics", "/public"]` |
+| | `excludedURLs` | Lists paths that bypass authentication | `["/favicon"]` | `["/health", "/metrics", "/public"]` |
-| `allowedUserDomains` | Restricts access to specific email domains | none | `["company.com", "subsidiary.com"]` |
+| | `allowedUserDomains` | Restricts access to specific email domains | none | `["company.com", "subsidiary.com"]` |
-| `allowedUsers` | A list of specific email addresses that are allowed access | none | `["[email protected]", "[email protected]"]` |
+| | `allowedRolesAndGroups` | Restricts access to users with specific roles or groups | none | `["admin", "developer"]` |
-| `allowedRolesAndGroups` | Restricts access to users with specific roles or groups | none | `["admin", "developer"]` |
+| | `revocationURL` | The endpoint for revoking tokens | auto-discovered | `https://accounts.google.com/revoke` |
-| `revocationURL` | The endpoint for revoking tokens | auto-discovered | `https://accounts.google.com/revoke` |
+| | `oidcEndSessionURL` | The provider's end session endpoint | auto-discovered | `https://accounts.google.com/logout` |
-| `oidcEndSessionURL` | The provider's end session endpoint | auto-discovered | `https://accounts.google.com/logout` |
+| | `enablePKCE` | Enables PKCE (Proof Key for Code Exchange) for authorization code flow | `false` | `true`, `false` |
-| `enablePKCE` | Enables PKCE (Proof Key for Code Exchange) for authorization code flow | `false` | `true`, `false` |
-| `refreshGracePeriodSeconds` | Seconds before token expiry to attempt proactive refresh | `60` | `120` |
-| `headers` | Custom HTTP headers with templates that can access OIDC claims and tokens | none | See "Templated Headers" section |
 
 ## Usage Examples
 
@@ -160,67 +157,6 @@ spec:
         - subsidiary.com
 ```
 
-### With Specific User Access
-
-```yaml
-apiVersion: traefik.io/v1alpha1
-kind: Middleware
-metadata:
-  name: oidc-specific-users
-  namespace: traefik
-spec:
-  plugin:
-    traefikoidc:
-      providerURL: https://accounts.google.com
-      clientID: 1234567890.apps.googleusercontent.com
-      clientSecret: your-client-secret
-      sessionEncryptionKey: potato-secret-is-at-least-32-bytes-long
-      callbackURL: /oauth2/callback
-      logoutURL: /oauth2/logout
-      scopes:
-        - openid
-        - email
-        - profile
-      allowedUsers:
-        - [email protected]
-        - [email protected]
-```
-
-### With Both Domain and Specific User Access
-
-```yaml
-apiVersion: traefik.io/v1alpha1
-kind: Middleware
-metadata:
-  name: oidc-domain-and-users
-  namespace: traefik
-spec:
-  plugin:
-    traefikoidc:
-      providerURL: https://accounts.google.com
-      clientID: 1234567890.apps.googleusercontent.com
-      clientSecret: your-client-secret
-      sessionEncryptionKey: potato-secret-is-at-least-32-bytes-long
-      callbackURL: /oauth2/callback
-      logoutURL: /oauth2/logout
-      scopes:
-        - openid
-        - email
-        - profile
-      allowedUserDomains:
-        - company.com
-      allowedUsers:
-        - [email protected]
-        - [email protected]
-```
-
-When configuring access control:
-- If only `allowedUsers` is set, only the specified email addresses will be granted access
-- If only `allowedUserDomains` is set, only users with email addresses from those domains will be granted access
-- If both are set, access is granted if the user's email is in `allowedUsers` OR their email's domain is in `allowedUserDomains`
-- If neither is set, any authenticated user will be granted access
-- Email matching is case-insensitive
-
 ### With Role-Based Access Control
 
 ```yaml
@@ -298,41 +234,6 @@ spec:
         - profile
 ```
 
-### With Templated Headers
-
-```yaml
-apiVersion: traefik.io/v1alpha1
-kind: Middleware
-metadata:
-  name: oidc-with-headers
-  namespace: traefik
-spec:
-  plugin:
-    traefikoidc:
-      providerURL: https://accounts.google.com
-      clientID: 1234567890.apps.googleusercontent.com
-      clientSecret: your-client-secret
-      sessionEncryptionKey: potato-secret-is-at-least-32-bytes-long
-      callbackURL: /oauth2/callback
-      logoutURL: /oauth2/logout
-      scopes:
-        - openid
-        - email
-        - profile
-        - roles
-      headers:
-        - name: "X-User-Email"
-          value: "{{.Claims.email}}"
-        - name: "X-User-ID"
-          value: "{{.Claims.sub}}"
-        - name: "Authorization"
-          value: "Bearer {{.AccessToken}}"
-        - name: "X-User-Roles"
-          value: "{{range $i, $e := .Claims.roles}}{{if $i}},{{end}}{{$e}}{{end}}"
-        - name: "X-Is-Admin"
-          value: "{{if eq .Claims.role \"admin\"}}true{{else}}false{{end}}"
-```
-
 ### With PKCE Enabled
 
 ```yaml
@@ -357,37 +258,6 @@ spec:
         - profile
 ```
 
-### Google OIDC Configuration Example
-
-This example shows a configuration specifically tailored for Google OIDC:
-
-```yaml
-apiVersion: traefik.io/v1alpha1
-kind: Middleware
-metadata:
-  name: oidc-google
-  namespace: traefik
-spec:
-  plugin:
-    traefikoidc:
-      providerURL: https://accounts.google.com
-      clientID: your-google-client-id.apps.googleusercontent.com # Replace with your Client ID
-      clientSecret: your-google-client-secret                     # Replace with your Client Secret
-      sessionEncryptionKey: your-secure-encryption-key-min-32-chars # Replace with your key
-      callbackURL: /oauth2/callback                             # Adjust if needed
-      logoutURL: /oauth2/logout                                 # Optional: Adjust if needed
-      scopes:
-        - openid
-        - email
-        - profile
-        # Note: DO NOT manually add offline_access scope for Google
-        # The middleware automatically handles Google-specific requirements
-      refreshGracePeriodSeconds: 300  # Optional: Start refresh 5 min before expiry (default 60)
-      # Other optional parameters like allowedUserDomains, etc. can be added here
-```
-
-The middleware automatically detects Google as the provider and applies the necessary adjustments to ensure proper authentication and token refresh. See the [Google OAuth Fix](#google-oauth-compatibility-fix) section for details.
-
 ### Keeping Secrets Secret in Kubernetes
 
 For Kubernetes environments, you can reference secrets instead of hardcoding sensitive values:
@@ -514,9 +384,6 @@ http:
             - profile
           allowedUserDomains:
             - company.com
-          allowedUsers:
-            - [email protected]
-            - [email protected]
           allowedRolesAndGroups:
             - admin
             - developer
@@ -528,15 +395,6 @@ http:
             - /public
             - /health
             - /metrics
-          headers:
-            - name: "X-User-Email"
-              value: "{{.Claims.email}}"
-            - name: "X-User-ID"
-              value: "{{.Claims.sub}}"
-            - name: "Authorization"
-              value: "Bearer {{.AccessToken}}"
-            - name: "X-User-Roles"
-              value: "{{range $i, $e := .Claims.roles}}{{if $i}},{{end}}{{$e}}{{end}}"
 ```
 
 ## Advanced Configuration
@@ -556,87 +414,11 @@ PKCE is recommended when:
 
 Note that not all OIDC providers support PKCE, so check your provider's documentation before enabling this feature.
 
-### Session Duration and Token Refresh
-
-This middleware aims to provide long-lived user sessions, typically up to 24 hours, by utilizing OIDC refresh tokens.
-
-**How it works:**
-- When a user authenticates, the middleware requests an access token and, if available, a refresh token from the OIDC provider.
-- The access token usually has a short lifespan (e.g., 1 hour).
-- Before the access token expires (controlled by `refreshGracePeriodSeconds`), the middleware uses the refresh token to obtain a new access token from the provider without requiring the user to log in again.
-- This process repeats, allowing the session to remain valid for as long as the refresh token is valid (often 24 hours or more, depending on the provider).
-
-**Provider-Specific Considerations (e.g., Google):**
-- Some providers, like Google, issue short-lived access tokens (e.g., 1 hour) and require specific configurations for long-term sessions.
-- To enable session extension beyond the initial token expiry with Google and similar providers, the middleware automatically includes the `offline_access` scope in the authentication request. This scope is necessary to obtain a refresh token.
-- For Google specifically, the middleware also adds the `prompt=consent` parameter to the initial authorization request. This ensures Google issues a refresh token, which is crucial for extending the session.
-- If a refresh attempt fails (e.g., the refresh token is revoked or expired), the user will be required to re-authenticate. The middleware includes enhanced error handling and logging for these scenarios.
-- Ensure your OIDC provider is configured to issue refresh tokens and allows their use for extending sessions. Check your provider's documentation for details on refresh token validity periods.
-
-### Google OAuth Compatibility Fix
-
-The middleware includes a specific fix for Google's OAuth implementation, which differs from the standard OIDC specification in how it handles refresh tokens:
-
-- **Issue**: Google does not support the standard `offline_access` scope for requesting refresh tokens and instead requires special parameters.
-  
-- **Automatic Solution**: The middleware detects Google as the provider based on the issuer URL and:
-  - Uses `access_type=offline` query parameter instead of the `offline_access` scope
-  - Adds `prompt=consent` to ensure refresh tokens are consistently issued
-  - Properly handles token refresh with Google's implementation
-
-You do not need any special configuration to use Google OAuth - just set `providerURL` to `https://accounts.google.com` and the middleware will automatically apply the proper parameters.
-
-For detailed information on the Google OAuth fix, see the [dedicated documentation](docs/google-oauth-fix.md).
-
 ### Token Caching and Blacklisting
 
 The middleware automatically caches validated tokens to improve performance and maintains a blacklist of revoked tokens.
-### Templated Headers
-
-The middleware supports setting custom HTTP headers with values templated from OIDC claims and tokens. This allows you to pass authentication information to downstream services in a flexible, customized format.
-
-Templates can access the following variables:
-- `{{.Claims.field}}` - Access individual claims from the ID token (e.g., `{{.Claims.email}}`, `{{.Claims.sub}}`)
-- `{{.AccessToken}}` - The raw access token string
-- `{{.IdToken}}` - The raw ID token string (same as AccessToken in most configurations)
-- `{{.RefreshToken}}` - The raw refresh token string
-
-**Example configuration:**
-```yaml
-headers:
-  - name: "X-User-Email"
-    value: "{{.Claims.email}}"
-  - name: "X-User-ID"
-    value: "{{.Claims.sub}}"
-  - name: "Authorization"
-    value: "Bearer {{.AccessToken}}"
-  - name: "X-User-Name"
-    value: "{{.Claims.given_name}} {{.Claims.family_name}}"
-```
-
-**Advanced template examples:**
-
-Conditional logic:
-```yaml
-headers:
-  - name: "X-Is-Admin"
-    value: "{{if eq .Claims.role \"admin\"}}true{{else}}false{{end}}"
-```
-
-Array handling:
-```yaml
-headers:
-  - name: "X-User-Roles"
-    value: "{{range $i, $e := .Claims.roles}}{{if $i}},{{end}}{{$e}}{{end}}"
-```
-
-**Notes:**
-- Variable names are case-sensitive (use `.Claims`, not `.claims`)
-- Missing claims will result in `<no value>` in the header value
-- The middleware validates templates during startup and logs errors for invalid templates
-
-### Default Headers Set for Downstream Services
 
+### Headers Set for Downstream Services
 
 When a user is authenticated, the middleware sets the following headers for downstream services:
 
@@ -673,12 +455,6 @@ logLevel: debug
 3. **No matching public key found**: The JWKS endpoint might be unavailable or the token's key ID (kid) doesn't match any key in the JWKS.
 4. **Access denied: Your email domain is not allowed**: The user's email domain is not in the `allowedUserDomains` list.
 5. **Access denied: You do not have any of the allowed roles or groups**: The user doesn't have any of the roles or groups specified in `allowedRolesAndGroups`.
-6. **Google sessions expire after ~1 hour**: If using Google as the OIDC provider and sessions expire prematurely (around 1 hour instead of longer), ensure:
-   - Do NOT manually add the `offline_access` scope. Google rejects this scope as invalid.
-   - The middleware automatically applies the required Google parameters (`access_type=offline` and `prompt=consent`).
-   - Your Google Cloud OAuth consent screen is set to "External" and "Production" mode. "Testing" mode often limits refresh token validity.
-   - Verify you're using a version of the middleware that includes the Google OAuth compatibility fix.
-   - For more details, see the [Google OAuth Compatibility Fix](#google-oauth-compatibility-fix) section or the [detailed documentation](docs/google-oauth-fix.md).
 
 ## Contributing
 

autocleanup.go

@@ -2,16 +2,8 @@ package traefikoidc
 
 import "time"
 
-// autoCleanupRoutine periodically calls the provided cleanup function.
+// autoCleanupRoutine runs a ticker that calls the provided cleanup function at the specified interval.
-// It starts a ticker with the given interval and executes the cleanup function
+// It stops when a value is received on the stop channel.
-// on each tick. The routine stops gracefully when a signal is received on the
-// stop channel. This is typically used for background cleanup tasks like
-// expiring cache entries.
-//
-// Parameters:
-//   - interval: The time duration between cleanup calls.
-//   - stop: A channel used to signal the routine to stop. Receiving any value will terminate the loop.
-//   - cleanup: The function to call periodically for cleanup tasks.
 func autoCleanupRoutine(interval time.Duration, stop <-chan struct{}, cleanup func()) {
 	ticker := time.NewTicker(interval)
 	defer ticker.Stop()

blacklist.go

@@ -0,0 +1,110 @@
+package traefikoidc
+
+import (
+	"sync"
+	"time"
+)
+
+// TokenBlacklist manages a thread-safe list of revoked tokens with expiration.
+type TokenBlacklist struct {
+	tokens map[string]time.Time
+	mutex  sync.RWMutex
+}
+
+// NewTokenBlacklist creates a new token blacklist instance.
+func NewTokenBlacklist() *TokenBlacklist {
+	return &TokenBlacklist{
+		tokens: make(map[string]time.Time),
+	}
+}
+
+// Add adds a token to the blacklist with an expiration time.
+func (b *TokenBlacklist) Add(token string, expiry time.Time) {
+	b.mutex.Lock()
+	defer b.mutex.Unlock()
+
+	// Clean up expired tokens if we're at capacity
+	if len(b.tokens) >= 1000 {
+		now := time.Now()
+		futureThreshold := now.Add(time.Minute)
+		for t, exp := range b.tokens {
+			if now.After(exp) || futureThreshold.After(exp) {
+				delete(b.tokens, t)
+			}
+		}
+
+		// If still at capacity, remove oldest token
+		if len(b.tokens) >= 1000 {
+			var oldestToken string
+			var oldestTime time.Time
+			first := true
+			for t, exp := range b.tokens {
+				if first || exp.Before(oldestTime) {
+					oldestToken = t
+					oldestTime = exp
+					first = false
+				}
+			}
+			if oldestToken != "" {
+				delete(b.tokens, oldestToken)
+			}
+		}
+	}
+
+	b.tokens[token] = expiry
+}
+
+// IsBlacklisted checks if a token is in the blacklist and not expired.
+func (b *TokenBlacklist) IsBlacklisted(token string) bool {
+	b.mutex.RLock()
+	defer b.mutex.RUnlock()
+
+	expiry, exists := b.tokens[token]
+	if !exists {
+		return false
+	}
+
+	// If token is expired, remove it and return false
+	if time.Now().After(expiry) {
+		// Switch to write lock to remove expired token
+		b.mutex.RUnlock()
+		b.mutex.Lock()
+		delete(b.tokens, token)
+		b.mutex.Unlock()
+		b.mutex.RLock()
+		return false
+	}
+
+	return true
+}
+
+// Cleanup removes expired tokens from the blacklist.
+// Also removes tokens that will expire within the next minute to prevent edge cases.
+func (b *TokenBlacklist) Cleanup() {
+	b.mutex.Lock()
+	defer b.mutex.Unlock()
+
+	now := time.Now()
+	futureThreshold := now.Add(time.Minute)
+
+	for token, expiry := range b.tokens {
+		// Remove tokens that are expired or will expire soon
+		if now.After(expiry) || futureThreshold.After(expiry) {
+			delete(b.tokens, token)
+		}
+	}
+}
+
+// Remove removes a token from the blacklist regardless of its expiration.
+func (b *TokenBlacklist) Remove(token string) {
+	b.mutex.Lock()
+	defer b.mutex.Unlock()
+	delete(b.tokens, token)
+}
+
+// Count returns the current number of tokens in the blacklist.
+func (b *TokenBlacklist) Count() int {
+	b.mutex.RLock()
+	defer b.mutex.RUnlock()
+	return len(b.tokens)
+}

blacklist_test.go

@@ -0,0 +1,74 @@
+package traefikoidc
+
+import (
+	"testing"
+	"time"
+)
+
+func TestTokenBlacklist_Add(t *testing.T) {
+	blacklist := NewTokenBlacklist()
+	token := "testToken"
+	expiry := time.Now().Add(time.Hour)
+
+	blacklist.Add(token, expiry)
+
+	if !blacklist.IsBlacklisted(token) {
+		t.Errorf("Expected token to be blacklisted, but it was not")
+	}
+}
+
+func TestTokenBlacklist_IsBlacklisted(t *testing.T) {
+	blacklist := NewTokenBlacklist()
+	token := "testToken"
+	expiry := time.Now().Add(time.Hour)
+
+	blacklist.Add(token, expiry)
+
+	if !blacklist.IsBlacklisted(token) {
+		t.Errorf("Expected token to be blacklisted, but it was not")
+	}
+
+	if blacklist.IsBlacklisted("nonExistentToken") {
+		t.Errorf("Expected non-existent token to not be blacklisted, but it was")
+	}
+}
+
+func TestTokenBlacklist_Cleanup(t *testing.T) {
+	blacklist := NewTokenBlacklist()
+	token := "testToken"
+	expiry := time.Now().Add(-time.Hour) // Expired token
+
+	blacklist.Add(token, expiry)
+	blacklist.Cleanup()
+
+	if blacklist.IsBlacklisted(token) {
+		t.Errorf("Expected expired token to be removed after cleanup, but it was not")
+	}
+}
+
+func TestTokenBlacklist_Remove(t *testing.T) {
+	blacklist := NewTokenBlacklist()
+	token := "testToken"
+	expiry := time.Now().Add(time.Hour)
+
+	blacklist.Add(token, expiry)
+	blacklist.Remove(token)
+
+	if blacklist.IsBlacklisted(token) {
+		t.Errorf("Expected token to be removed, but it was not")
+	}
+}
+
+func TestTokenBlacklist_Count(t *testing.T) {
+	blacklist := NewTokenBlacklist()
+	token1 := "token1"
+	token2 := "token2"
+	expiry := time.Now().Add(time.Hour)
+
+	blacklist.Add(token1, expiry)
+	blacklist.Add(token2, expiry)
+
+	if blacklist.Count() != 2 {
+		t.Errorf("Expected blacklist count to be 2, but got %d", blacklist.Count())
+	}
+}

cache.go

@@ -46,9 +46,7 @@ type Cache struct {
 // DefaultMaxSize is the default maximum number of items in the cache.
 const DefaultMaxSize = 500
 
-// NewCache creates a new empty cache instance with default settings.
+// NewCache creates a new empty cache instance that is ready for use.
-// It initializes the internal maps and list, sets the default maximum size,
-// and starts the automatic cleanup goroutine.
 func NewCache() *Cache {
 	c := &Cache{
 		items:               make(map[string]CacheItem, DefaultMaxSize),
@@ -62,12 +60,8 @@ func NewCache() *Cache {
 	return c
 }
 
-// Set adds or updates an item in the cache with the specified key, value, and expiration duration.
+// Set adds or updates an item in the cache with the specified expiration duration.
-// If the key already exists, its value and expiration time are updated, and it's moved
+// It moves the item to the most recently used position.
-// to the most recently used position in the LRU list.
-// If the key does not exist and the cache is full, the least recently used item is evicted
-// before adding the new item.
-// The expiration duration is relative to the time Set is called.
 func (c *Cache) Set(key string, value interface{}, expiration time.Duration) {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
@@ -101,11 +95,8 @@ func (c *Cache) Set(key string, value interface{}, expiration time.Duration) {
 	c.elems[key] = elem
 }
 
-// Get retrieves an item from the cache by its key.
+// Get retrieves an item from the cache if it exists and hasn't expired.
-// If the item exists and has not expired, its value and true are returned.
+// Moving the accessed item to the most recently used position.
-// Accessing an item moves it to the most recently used position in the LRU list.
-// If the item does not exist or has expired, nil and false are returned, and the
-// expired item is removed from the cache.
 func (c *Cache) Get(key string) (interface{}, bool) {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
@@ -129,9 +120,7 @@ func (c *Cache) Get(key string) (interface{}, bool) {
 	return item.Value, true
 }
 
-// Delete removes an item from the cache by its key.
+// Delete removes an item from the cache.
-// If the key exists, the corresponding item is removed from the cache storage
-// and the LRU list.
 func (c *Cache) Delete(key string) {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
@@ -139,28 +128,22 @@ func (c *Cache) Delete(key string) {
 	c.removeItem(key)
 }
 
-// Cleanup iterates through the cache and removes all items that have expired.
+// Cleanup removes all expired items from the cache. This should be called periodically
-// An item is considered expired if the current time is after its ExpiresAt timestamp.
+// to prevent memory bloat from expired entries.
-// This method is called automatically by the auto-cleanup goroutine, but can also
-// be called manually.
 func (c *Cache) Cleanup() {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
 
 	now := time.Now()
 	for key, item := range c.items {
-		// Remove items that are expired
+		// Remove items that are expired or within 10% of expiration
-		if now.After(item.ExpiresAt) {
+		if now.After(item.ExpiresAt) || now.Add(time.Duration(float64(item.ExpiresAt.Sub(now))*0.1)).After(item.ExpiresAt) {
 			c.removeItem(key)
 		}
 	}
 }
 
-// evictOldest removes the least recently used (oldest) item from the cache.
+// evictOldest removes the least recently used item from the cache.
-// It first attempts to find and remove an expired item from the front of the LRU list.
-// If no expired items are found at the front, it removes the absolute oldest item (front of the list).
-// This method is called internally by Set when the cache reaches its maximum size.
-// Note: This function assumes the write lock is already held.
 func (c *Cache) evictOldest() {
 	now := time.Now()
 	elem := c.order.Front()
@@ -184,28 +167,7 @@ func (c *Cache) evictOldest() {
 	}
 }
 
-// SetMaxSize changes the maximum number of items the cache can hold.
+// removeItem removes an item from both the cache and the LRU tracking structures.
-// If the new size is smaller than the current number of items in the cache,
-// oldest items will be evicted until the cache size is within the new limit.
-func (c *Cache) SetMaxSize(size int) {
-	if size <= 0 {
-		return // Invalid size, ignore
-	}
-
-	c.mutex.Lock()
-	defer c.mutex.Unlock()
-
-	c.maxSize = size
-
-	// If cache exceeds the new max size, evict oldest items
-	for len(c.items) > c.maxSize {
-		c.evictOldest()
-	}
-}
-
-// removeItem removes an item specified by the key from the cache's internal storage (items map)
-// and its corresponding entry from the LRU list (order list and elems map).
-// Note: This function assumes the write lock is already held.
 func (c *Cache) removeItem(key string) {
 	delete(c.items, key)
 	if elem, ok := c.elems[key]; ok {
@@ -214,15 +176,12 @@ func (c *Cache) removeItem(key string) {
 	}
 }
 
-// startAutoCleanup starts the background goroutine that automatically calls the Cleanup method
+// startAutoCleanup initiates a goroutine that periodically cleans up expired cache items.
-// at the interval specified by c.autoCleanupInterval.
-// It uses the autoCleanupRoutine helper function.
 func (c *Cache) startAutoCleanup() {
 	autoCleanupRoutine(c.autoCleanupInterval, c.stopCleanup, c.Cleanup)
 }
 
-// Close stops the automatic cleanup goroutine associated with this cache instance.
+// Close terminates the auto cleanup goroutine.
-// It should be called when the cache is no longer needed to prevent resource leaks.
 func (c *Cache) Close() {
 	close(c.stopCleanup)
 }

cache_test.go

@@ -1,99 +1,306 @@
 package traefikoidc
 
 import (
+	"reflect"
 	"testing"
 	"time"
 )
 
-func TestCache_Cleanup(t *testing.T) {
+func TestCache(t *testing.T) {
-	c := NewCache()
+	t.Run("Basic Set and Get", func(t *testing.T) {
+		cache := NewCache()
+		key := "test-key"
+		value := "test-value"
+		expiration := 1 * time.Second
 
-	// Add some items with different expiration times
+		// Test Set
-	now := time.Now()
+		cache.Set(key, value, expiration)
-	pastTime := now.Add(-1 * time.Hour)  // Already expired
-	futureTime := now.Add(1 * time.Hour) // Not expired
 
-	// Create test items
+		// Test Get
-	c.items["expired"] = CacheItem{
+		got, found := cache.Get(key)
-		Value:     "expired-value",
+		if !found {
-		ExpiresAt: pastTime,
+			t.Error("Expected to find key in cache")
-	}
+		}
+		if got != value {
+			t.Errorf("Expected value %v, got %v", value, got)
+		}
+	})
 
-	c.items["valid"] = CacheItem{
+	t.Run("Expiration", func(t *testing.T) {
-		Value:     "valid-value",
+		cache := NewCache()
-		ExpiresAt: futureTime,
+		key := "test-key"
-	}
+		value := "test-value"
+		expiration := 10 * time.Millisecond
 
-	// Store original elements in the order list to match items
+		// Set with short expiration
-	c.elems["expired"] = c.order.PushBack(lruEntry{key: "expired"})
+		cache.Set(key, value, expiration)
-	c.elems["valid"] = c.order.PushBack(lruEntry{key: "valid"})
 
-	// Call cleanup, which should only remove expired items
+		// Wait for expiration
-	c.Cleanup()
+		time.Sleep(20 * time.Millisecond)
 
-	// Check that only the expired item was removed
+		// Should not find expired key
-	if _, exists := c.items["expired"]; exists {
+		_, found := cache.Get(key)
-		t.Error("Expired item was not removed by Cleanup()")
+		if found {
-	}
+			t.Error("Expected key to be expired")
+		}
+	})
 
-	if _, exists := c.items["valid"]; !exists {
+	t.Run("Delete", func(t *testing.T) {
-		t.Error("Valid item was incorrectly removed by Cleanup()")
+		cache := NewCache()
-	}
+		key := "test-key"
+		value := "test-value"
+		expiration := 1 * time.Second
+
+		// Set and then delete
+		cache.Set(key, value, expiration)
+		cache.Delete(key)
+
+		// Should not find deleted key
+		_, found := cache.Get(key)
+		if found {
+			t.Error("Expected key to be deleted")
+		}
+	})
+
+	t.Run("Cleanup", func(t *testing.T) {
+		cache := NewCache()
+		// Add multiple items with different expirations
+		cache.Set("expired1", "value1", 10*time.Millisecond)
+		cache.Set("expired2", "value2", 10*time.Millisecond)
+		cache.Set("valid", "value3", 1*time.Second)
+
+		// Wait for some items to expire
+		time.Sleep(20 * time.Millisecond)
+
+		// Run cleanup
+		cache.Cleanup()
+
+		// Check expired items are removed
+		_, found1 := cache.Get("expired1")
+		_, found2 := cache.Get("expired2")
+		_, found3 := cache.Get("valid")
+
+		if found1 {
+			t.Error("Expected expired1 to be cleaned up")
+		}
+		if found2 {
+			t.Error("Expected expired2 to be cleaned up")
+		}
+		if !found3 {
+			t.Error("Expected valid item to remain in cache")
+		}
+	})
+
+	t.Run("Concurrent Access", func(t *testing.T) {
+		cache := NewCache()
+		done := make(chan bool)
+
+		// Start multiple goroutines to access cache concurrently
+		for i := 0; i < 10; i++ {
+			go func(id int) {
+				key := "key"
+				value := "value"
+				expiration := 1 * time.Second
+
+				// Perform multiple operations
+				cache.Set(key, value, expiration)
+				cache.Get(key)
+				cache.Delete(key)
+				cache.Cleanup()
+
+				done <- true
+			}(i)
+		}
+
+		// Wait for all goroutines to complete
+		for i := 0; i < 10; i++ {
+			<-done
+		}
+	})
+
+	t.Run("Zero Expiration", func(t *testing.T) {
+		cache := NewCache()
+		key := "test-key"
+		value := "test-value"
+
+		// Set with zero expiration
+		cache.Set(key, value, 0)
+
+		// Should not find the key
+		_, found := cache.Get(key)
+		if found {
+			t.Error("Expected key with zero expiration to be immediately expired")
+		}
+	})
+
+	t.Run("Negative Expiration", func(t *testing.T) {
+		cache := NewCache()
+		key := "test-key"
+		value := "test-value"
+
+		// Set with negative expiration
+		cache.Set(key, value, -1*time.Second)
+
+		// Should not find the key
+		_, found := cache.Get(key)
+		if found {
+			t.Error("Expected key with negative expiration to be immediately expired")
+		}
+	})
+
+	t.Run("Update Existing Key", func(t *testing.T) {
+		cache := NewCache()
+		key := "test-key"
+		value1 := "value1"
+		value2 := "value2"
+		expiration := 1 * time.Second
+
+		// Set initial value
+		cache.Set(key, value1, expiration)
+
+		// Update value
+		cache.Set(key, value2, expiration)
+
+		// Check updated value
+		got, found := cache.Get(key)
+		if !found {
+			t.Error("Expected to find key in cache")
+		}
+		if got != value2 {
+			t.Errorf("Expected updated value %v, got %v", value2, got)
+		}
+	})
+
+	t.Run("Different Value Types", func(t *testing.T) {
+		cache := NewCache()
+		expiration := 1 * time.Second
+
+		// Test with different value types
+		testCases := []struct {
+			key   string
+			value interface{}
+		}{
+			{"string", "test"},
+			{"int", 42},
+			{"float", 3.14},
+			{"bool", true},
+			{"slice", []string{"a", "b", "c"}},
+			{"map", map[string]int{"a": 1, "b": 2}},
+			{"struct", struct{ Name string }{"test"}},
+		}
+
+		for _, tc := range testCases {
+			t.Run(tc.key, func(t *testing.T) {
+				cache.Set(tc.key, tc.value, expiration)
+				got, found := cache.Get(tc.key)
+				if !found {
+					t.Error("Expected to find key in cache")
+				}
+				// Use reflect.DeepEqual for comparing complex types like slices and maps
+				if !reflect.DeepEqual(got, tc.value) {
+					t.Errorf("Expected value %v, got %v", tc.value, got)
+				}
+			})
+		}
+	})
 }
 
-func TestCache_SetMaxSize(t *testing.T) {
+func TestTokenCache(t *testing.T) {
-	c := NewCache()
+	t.Run("Basic Operations", func(t *testing.T) {
+		tc := NewTokenCache()
+		token := "test-token"
+		claims := map[string]interface{}{
+			"sub":   "1234567890",
+			"name":  "John Doe",
+			"admin": true,
+		}
+		expiration := 1 * time.Second
 
-	// Set a lower max size
+		// Test Set and Get
-	originalMaxSize := c.maxSize
+		tc.Set(token, claims, expiration)
-	newMaxSize := 3
+		gotClaims, found := tc.Get(token)
+		if !found {
+			t.Error("Expected to find token in cache")
+		}
+		if len(gotClaims) != len(claims) {
+			t.Errorf("Expected %d claims, got %d", len(claims), len(gotClaims))
+		}
+		for k, v := range claims {
+			if gotClaims[k] != v {
+				t.Errorf("Expected claim %s to be %v, got %v", k, v, gotClaims[k])
+			}
+		}
 
-	// Add more items than the new max size
+		// Test Delete
-	for i := 0; i < originalMaxSize; i++ {
+		tc.Delete(token)
-		key := "key" + string(rune('A'+i))
+		_, found = tc.Get(token)
-		c.Set(key, i, 1*time.Hour)
+		if found {
-	}
+			t.Error("Expected token to be deleted")
+		}
+	})
 
-	// Verify items were added
+	t.Run("Expiration", func(t *testing.T) {
-	if len(c.items) != originalMaxSize {
+		tc := NewTokenCache()
-		t.Errorf("Expected %d items before SetMaxSize, got %d", originalMaxSize, len(c.items))
+		token := "test-token"
-	}
+		claims := map[string]interface{}{"sub": "1234567890"}
+		expiration := 10 * time.Millisecond
 
-	// Change the max size to a smaller value
+		// Set with short expiration
-	c.SetMaxSize(newMaxSize)
+		tc.Set(token, claims, expiration)
 
-	// Check that the cache was reduced to the new max size
+		// Wait for expiration
-	if len(c.items) > newMaxSize {
+		time.Sleep(20 * time.Millisecond)
-		t.Errorf("Cache size %d exceeds new max size %d after SetMaxSize", len(c.items), newMaxSize)
-	}
 
-	if c.maxSize != newMaxSize {
+		// Should not find expired token
-		t.Errorf("Cache maxSize not updated, expected %d, got %d", newMaxSize, c.maxSize)
+		_, found := tc.Get(token)
-	}
+		if found {
+			t.Error("Expected token to be expired")
+		}
+	})
 
-	// Check that the oldest items were evicted (should keep "keyC", "keyD", "keyE", etc.)
+	t.Run("Cleanup", func(t *testing.T) {
-	if _, exists := c.items["keyA"]; exists {
+		tc := NewTokenCache()
-		t.Error("Expected oldest item 'keyA' to be evicted, but it still exists")
+
-	}
+		// Add multiple tokens with different expirations
-}
+		tc.Set("expired1", map[string]interface{}{"sub": "1"}, 10*time.Millisecond)
-
+		tc.Set("expired2", map[string]interface{}{"sub": "2"}, 10*time.Millisecond)
-func TestJWKCache_WithInternalCache(t *testing.T) {
+		tc.Set("valid", map[string]interface{}{"sub": "3"}, 1*time.Second)
-	cache := NewJWKCache()
+
-
+		// Wait for some tokens to expire
-	// Check that the internal cache is properly initialized
+		time.Sleep(20 * time.Millisecond)
-	if cache.internalCache == nil {
+
-		t.Error("internalCache field was not initialized")
+		// Run cleanup
-	}
+		tc.Cleanup()
-
+
-	// Test max size configuration
+		// Check expired tokens are removed
-	testSize := 50
+		_, found1 := tc.Get("expired1")
-	cache.SetMaxSize(testSize)
+		_, found2 := tc.Get("expired2")
-
+		_, found3 := tc.Get("valid")
-	if cache.maxSize != testSize {
+
-		t.Errorf("JWKCache maxSize not updated, expected %d, got %d", testSize, cache.maxSize)
+		if found1 {
-	}
+			t.Error("Expected expired1 to be cleaned up")
-
+		}
-	if cache.internalCache.maxSize != testSize {
+		if found2 {
-		t.Errorf("internalCache maxSize not updated, expected %d, got %d", testSize, cache.internalCache.maxSize)
+			t.Error("Expected expired2 to be cleaned up")
-	}
+		}
+		if !found3 {
+			t.Error("Expected valid token to remain in cache")
+		}
+	})
+
+	t.Run("Token Prefix", func(t *testing.T) {
+		tc := NewTokenCache()
+		token := "test-token"
+		claims := map[string]interface{}{"sub": "1234567890"}
+		expiration := 1 * time.Second
+
+		// Set token
+		tc.Set(token, claims, expiration)
+
+		// Verify internal storage uses prefix
+		_, found := tc.cache.Get("t-" + token)
+		if !found {
+			t.Error("Expected to find prefixed token in underlying cache")
+		}
+	})
 }

docs/google-oauth-fix.md

@@ -1,163 +0,0 @@
-# Google OAuth Integration Fix
-
-## Problem Overview
-
-The Traefik OIDC plugin encountered an authentication issue when using Google as an OAuth provider. Authentication would fail with the following error:
-
-```
-Some requested scopes were invalid. {valid=[openid, https://www.googleapis.com/auth/userinfo.email, https://www.googleapis.com/auth/userinfo.profile], invalid=[offline_access]}
-```
-
-This occurred because Google's OAuth implementation differs from the standard OIDC specification in how it handles refresh tokens and offline access.
-
-## Technical Details of the Issue
-
-### Standard OIDC Provider Behavior
-
-Most OpenID Connect (OIDC) providers follow the standard specification, where:
-- To obtain a refresh token, clients include the `offline_access` scope in their authorization request
-- This allows authenticated sessions to persist beyond the initial access token expiration
-
-### Google's Non-Standard Approach
-
-Google's OAuth implementation deviates from the standard by:
-1. Not supporting the `offline_access` scope, instead rejecting it as an invalid scope
-2. Requiring the `access_type=offline` query parameter for requesting refresh tokens
-3. Needing the `prompt=consent` parameter to consistently issue refresh tokens (especially for repeat authentications)
-
-This difference caused the plugin to fail when configured for Google OAuth, as it was using a standard approach that didn't work with Google's implementation.
-
-## Solution Implementation
-
-The fix involved modifying the authentication flow to specifically handle Google providers:
-
-1. **Google Provider Detection**: Added code to detect if the OIDC provider is Google based on the issuer URL:
-
-```go
-// Check if we're dealing with a Google OIDC provider
-isGoogleProvider := strings.Contains(t.issuerURL, "google") || 
-                   strings.Contains(t.issuerURL, "accounts.google.com")
-```
-
-2. **Provider-Specific Auth URL Building**: Modified the `buildAuthURL` function to handle Google and non-Google providers differently:
-
-```go
-// Handle offline access differently for Google vs other providers
-if isGoogleProvider {
-    // For Google, use access_type=offline parameter instead of offline_access scope
-    params.Set("access_type", "offline")
-    t.logger.Debug("Google OIDC provider detected, added access_type=offline for refresh tokens")
-
-    // Add prompt=consent for Google to ensure refresh token is issued
-    params.Set("prompt", "consent")
-    t.logger.Debug("Google OIDC provider detected, added prompt=consent to ensure refresh tokens")
-} else {
-    // For non-Google providers, use the offline_access scope
-    hasOfflineAccess := false
-    for _, scope := range scopes {
-        if scope == "offline_access" {
-            hasOfflineAccess = true
-            break
-        }
-    }
-
-    if !hasOfflineAccess {
-        scopes = append(scopes, "offline_access")
-    }
-}
-```
-
-3. **Token Refresh Enhancement**: Improved the token refresh logic to better handle Google's behavior, particularly when refresh tokens aren't returned in refresh responses (as Google often uses the same refresh token for multiple requests).
-
-## Why This Approach Works
-
-This solution aligns with Google's OAuth 2.0 documentation which specifies:
-
-1. **Access Type Parameter**: Google's [OAuth 2.0 documentation](https://developers.google.com/identity/protocols/oauth2/web-server#offline) states that to request a refresh token, applications must include `access_type=offline` in the authorization request.
-
-2. **Prompt Parameter**: The [`prompt=consent`](https://developers.google.com/identity/protocols/oauth2/web-server#forceapprovalprompt) parameter forces the consent screen to appear, ensuring a refresh token is issued even if the user has previously granted access.
-
-3. **Scope Validation**: Google strictly validates scopes and rejects non-standard ones like `offline_access`, instead relying on the `access_type` parameter to indicate whether a refresh token should be issued.
-
-By adapting to these Google-specific requirements, the OIDC plugin can now seamlessly work with both standard OIDC providers and Google's OAuth implementation.
-
-## Testing and Verification
-
-Comprehensive tests were implemented to verify the solution:
-
-1. **Provider Detection Test**: Ensures the code correctly identifies Google providers and applies the appropriate parameters.
-
-2. **Auth URL Parameter Tests**: Verifies that:
-   - For Google providers: `access_type=offline` and `prompt=consent` are included; `offline_access` scope is NOT included
-   - For non-Google providers: `offline_access` scope IS included; `access_type` parameter is NOT added
-
-3. **Token Refresh Tests**: Validates that Google's token refresh process works correctly, including the preservation of refresh tokens when Google doesn't return a new one.
-
-4. **Integration Test**: Tests the complete authentication flow with a mocked Google provider to ensure all components work together seamlessly.
-
-Sample test case (simplified):
-
-```go
-t.Run("Google provider detection adds required parameters", func(t *testing.T) {
-    // Test buildAuthURL to ensure it adds access_type=offline and prompt=consent for Google
-    authURL := tOidc.buildAuthURL("https://example.com/callback", "state123", "nonce123", "")
-
-    // Check that access_type=offline was added (not offline_access scope for Google)
-    if !strings.Contains(authURL, "access_type=offline") {
-        t.Errorf("access_type=offline not added to Google auth URL: %s", authURL)
-    }
-
-    // Verify offline_access scope is NOT included for Google providers
-    if strings.Contains(authURL, "offline_access") {
-        t.Errorf("offline_access scope incorrectly added to Google auth URL: %s", authURL)
-    }
-
-    // Check that prompt=consent was added
-    if !strings.Contains(authURL, "prompt=consent") {
-        t.Errorf("prompt=consent not added to Google auth URL: %s", authURL)
-    }
-})
-```
-
-## Usage Guidance for Developers
-
-When configuring the Traefik OIDC middleware for Google:
-
-1. **Provider URL**: Use `https://accounts.google.com` as the `providerURL` value
-
-2. **Client Configuration**: Create OAuth 2.0 credentials in the Google Cloud Console:
-   - Configure the authorized redirect URI to match your `callbackURL` setting
-   - Ensure your OAuth consent screen is properly configured (especially if you want long-lived refresh tokens)
-
-3. **Configuration Example**:
-
-```yaml
-apiVersion: traefik.io/v1alpha1
-kind: Middleware
-metadata:
-  name: oidc-google
-  namespace: traefik
-spec:
-  plugin:
-    traefikoidc:
-      providerURL: https://accounts.google.com
-      clientID: your-google-client-id.apps.googleusercontent.com
-      clientSecret: your-google-client-secret
-      sessionEncryptionKey: your-secure-encryption-key-min-32-chars
-      callbackURL: /oauth2/callback
-      scopes:
-        - openid
-        - email
-        - profile
-        # Note: DO NOT manually add offline_access scope for Google
-        # The middleware handles this automatically and correctly
-```
-
-4. **Troubleshooting**: If sessions still expire prematurely with Google (typically after 1 hour):
-   - Ensure your Google Cloud OAuth consent screen is set to "External" and "Production" mode (not "Testing" mode, which limits refresh token validity)
-   - Review your application logs with `logLevel: debug` to check for refresh token errors
-   - Verify you're using a version of the middleware that includes this fix
-
-## Conclusion
-
-This fix ensures that the Traefik OIDC plugin works seamlessly with Google's OAuth implementation without requiring users to make provider-specific configuration changes. The middleware now intelligently adapts to the provider's requirements, making it more robust and user-friendly while maintaining compatibility with the standard OIDC specification for other providers.

google_session_test.go

@@ -1,592 +0,0 @@
-package traefikoidc
-
-import (
-	"context"
-	"crypto/rand"
-	"crypto/rsa"
-	"encoding/base64"
-	"fmt"
-	"math/big"
-	"net/http/httptest"
-	"net/url"
-	"strings"
-	"testing"
-	"time"
-
-	"golang.org/x/time/rate"
-)
-
-// MockJWTVerifier implements the JWTVerifier interface for testing
-type MockJWTVerifier struct {
-	VerifyJWTFunc func(jwt *JWT, token string) error
-}
-
-func (m *MockJWTVerifier) VerifyJWTSignatureAndClaims(jwt *JWT, token string) error {
-	if m.VerifyJWTFunc != nil {
-		return m.VerifyJWTFunc(jwt, token)
-	}
-	return nil
-}
-
-func TestGoogleOIDCRefreshTokenHandling(t *testing.T) {
-	// Create a mocked TraefikOidc instance that simulates Google provider behavior
-	mockLogger := NewLogger("debug")
-
-	// Create a test instance with a Google-like issuer URL
-	tOidc := &TraefikOidc{
-		issuerURL:          "https://accounts.google.com",
-		clientID:           "test-client-id",
-		clientSecret:       "test-client-secret",
-		logger:             mockLogger,
-		scopes:             []string{"openid", "profile", "email"},
-		refreshGracePeriod: 60,
-	}
-
-	// Create a session manager
-	sessionManager, _ := NewSessionManager("0123456789abcdef0123456789abcdef", true, mockLogger)
-	tOidc.sessionManager = sessionManager
-
-	t.Run("Google provider detection adds required parameters", func(t *testing.T) {
-		// Test buildAuthURL to ensure it adds access_type=offline and prompt=consent for Google
-		authURL := tOidc.buildAuthURL("https://example.com/callback", "state123", "nonce123", "")
-
-		// Check that access_type=offline was added (not offline_access scope for Google)
-		if !strings.Contains(authURL, "access_type=offline") {
-			t.Errorf("access_type=offline not added to Google auth URL: %s", authURL)
-		}
-
-		// Verify offline_access scope is NOT included for Google providers
-		if strings.Contains(authURL, "offline_access") {
-			t.Errorf("offline_access scope incorrectly added to Google auth URL: %s", authURL)
-		}
-
-		// Check that prompt=consent was added
-		if !strings.Contains(authURL, "prompt=consent") {
-			t.Errorf("prompt=consent not added to Google auth URL: %s", authURL)
-		}
-	})
-
-	t.Run("Non-Google provider doesn't add Google-specific params", func(t *testing.T) {
-		// Create a test instance with a non-Google issuer URL
-		nonGoogleOidc := &TraefikOidc{
-			issuerURL:    "https://auth.example.com",
-			clientID:     "test-client-id",
-			clientSecret: "test-client-secret",
-			logger:       mockLogger,
-			scopes:       []string{"openid", "profile", "email"},
-		}
-
-		// Test buildAuthURL without Google-specific parameters
-		authURL := nonGoogleOidc.buildAuthURL("https://example.com/callback", "state123", "nonce123", "")
-
-		// Check that prompt=consent is not automatically added
-		if strings.Contains(authURL, "prompt=consent") {
-			t.Errorf("prompt=consent added to non-Google auth URL: %s", authURL)
-		}
-	})
-
-	t.Run("Session refresh with Google provider", func(t *testing.T) {
-		// Create a request and response recorder
-		req := httptest.NewRequest("GET", "/test", nil)
-		rw := httptest.NewRecorder()
-
-		// Create a session and set a refresh token
-		session, _ := sessionManager.GetSession(req)
-		session.SetAuthenticated(true)
-		session.SetEmail("[email protected]")
-		session.SetAccessToken("old-access-token")
-		session.SetRefreshToken("valid-refresh-token")
-
-		// Create a mock token exchanger that simulates Google's behavior
-		mockTokenExchanger := &MockTokenExchanger{
-			RefreshTokenFunc: func(refreshToken string) (*TokenResponse, error) {
-				// Check that the refresh token is passed correctly
-				if refreshToken != "valid-refresh-token" {
-					t.Errorf("Incorrect refresh token passed: %s", refreshToken)
-					return nil, fmt.Errorf("invalid token")
-				}
-
-				// Return a simulated Google token response with a new access token
-				// but without a new refresh token (Google doesn't always return a new refresh token)
-				return &TokenResponse{
-					IDToken:      "new-id-token-from-google",
-					AccessToken:  "new-access-token-from-google",
-					RefreshToken: "", // Google often doesn't return a new refresh token
-					ExpiresIn:    3600,
-				}, nil
-			},
-		}
-
-		// Set the mock token exchanger
-		tOidc.tokenExchanger = mockTokenExchanger
-
-		// Create a struct that implements the TokenVerifier interface
-		tOidc.tokenVerifier = &MockTokenVerifier{
-			VerifyFunc: func(token string) error {
-				return nil
-			},
-		}
-
-		tOidc.extractClaimsFunc = func(token string) (map[string]interface{}, error) {
-			// Return mock claims
-			return map[string]interface{}{
-				"email": "[email protected]",
-				"exp":   float64(time.Now().Add(1 * time.Hour).Unix()),
-			}, nil
-		}
-
-		// Attempt to refresh the token
-		refreshed := tOidc.refreshToken(rw, req, session)
-
-		// Verify the refresh was successful
-		if !refreshed {
-			t.Error("Token refresh failed for Google provider")
-		}
-
-		// Check that we kept the original refresh token since Google didn't provide a new one
-		if session.GetRefreshToken() != "valid-refresh-token" {
-			t.Errorf("Original refresh token not preserved: got %s, expected 'valid-refresh-token'",
-				session.GetRefreshToken())
-		}
-
-		// Check that the tokens were updated correctly
-		if session.GetIDToken() != "new-id-token-from-google" {
-			t.Errorf("ID token not updated: got %s, expected 'new-id-token-from-google'",
-				session.GetIDToken())
-		}
-
-		if session.GetAccessToken() != "new-access-token-from-google" {
-			t.Errorf("Access token not updated: got %s, expected 'new-access-token-from-google'",
-				session.GetAccessToken())
-		}
-	})
-	// Test that our fix specifically addresses the reported Google error
-	t.Run("Google provider handles offline access correctly", func(t *testing.T) {
-		// Build the auth URL with Google provider detection
-		authURL := tOidc.buildAuthURL("https://example.com/callback", "state123", "nonce123", "")
-
-		// Parse the URL to examine its parameters
-		parsedURL, err := url.Parse(authURL)
-		if err != nil {
-			t.Fatalf("Failed to parse auth URL: %v", err)
-		}
-
-		params := parsedURL.Query()
-
-		// Verify that access_type=offline is set (Google's way of requesting refresh tokens)
-		if params.Get("access_type") != "offline" {
-			t.Errorf("access_type=offline not set in Google auth URL")
-		}
-
-		// Verify that the scope parameter doesn't contain offline_access
-		// (which Google reports as invalid: {invalid=[offline_access]})
-		scope := params.Get("scope")
-		if strings.Contains(scope, "offline_access") {
-			t.Errorf("offline_access incorrectly included in scope for Google provider: %s", scope)
-		}
-
-		// Verify that the necessary scopes are still included
-		for _, requiredScope := range []string{"openid", "profile", "email"} {
-			if !strings.Contains(scope, requiredScope) {
-				t.Errorf("Required scope '%s' missing from auth URL", requiredScope)
-			}
-		}
-	})
-
-	// Enhanced test for verifying non-Google provider includes offline_access scope
-	t.Run("Non-Google provider includes offline_access scope", func(t *testing.T) {
-		// Create a test instance with a non-Google issuer URL
-		nonGoogleOidc := &TraefikOidc{
-			issuerURL:    "https://auth.example.com",
-			clientID:     "test-client-id",
-			clientSecret: "test-client-secret",
-			logger:       mockLogger,
-			scopes:       []string{"openid", "profile", "email"},
-		}
-
-		// Test buildAuthURL for a non-Google provider
-		authURL := nonGoogleOidc.buildAuthURL("https://example.com/callback", "state123", "nonce123", "")
-
-		// Parse the URL to examine its parameters
-		parsedURL, err := url.Parse(authURL)
-		if err != nil {
-			t.Fatalf("Failed to parse auth URL: %v", err)
-		}
-
-		params := parsedURL.Query()
-
-		// Verify that access_type=offline is NOT set for non-Google providers
-		if params.Get("access_type") == "offline" {
-			t.Errorf("access_type=offline incorrectly added to non-Google auth URL")
-		}
-
-		// Verify that offline_access scope IS included for non-Google providers
-		scope := params.Get("scope")
-		if !strings.Contains(scope, "offline_access") {
-			t.Errorf("offline_access scope missing from non-Google auth URL scope: %s", scope)
-		}
-
-		// Verify that the necessary scopes are still included
-		for _, requiredScope := range []string{"openid", "profile", "email"} {
-			if !strings.Contains(scope, requiredScope) {
-				t.Errorf("Required scope '%s' missing from non-Google auth URL", requiredScope)
-			}
-		}
-	})
-
-	// Additional test for complete URL construction for Google provider
-	t.Run("Complete Google auth URL construction", func(t *testing.T) {
-		// Build the auth URL with additional parameters
-		redirectURL := "https://example.com/callback"
-		state := "state123"
-		nonce := "nonce123"
-		codeChallenge := "code_challenge_value" // For PKCE
-
-		// Enable PKCE for this test
-		tOidc.enablePKCE = true
-
-		// Build auth URL
-		authURL := tOidc.buildAuthURL(redirectURL, state, nonce, codeChallenge)
-
-		// Parse the URL to examine its structure and parameters
-		parsedURL, err := url.Parse(authURL)
-		if err != nil {
-			t.Fatalf("Failed to parse auth URL: %v", err)
-		}
-
-		// Verify the base URL
-		expectedBaseURL := "https://accounts.google.com/o/oauth2/v2/auth"
-		if !strings.HasPrefix(authURL, expectedBaseURL) && !strings.Contains(authURL, "accounts.google.com") {
-			t.Errorf("Auth URL doesn't start with expected Google OAuth endpoint: %s", authURL)
-		}
-
-		// Check all required parameters
-		params := parsedURL.Query()
-		expectedParams := map[string]string{
-			"client_id":     "test-client-id",
-			"response_type": "code",
-			"redirect_uri":  redirectURL,
-			"state":         state,
-			"nonce":         nonce,
-			"access_type":   "offline",
-			"prompt":        "consent",
-		}
-
-		// Also check PKCE parameters if enabled
-		if tOidc.enablePKCE {
-			expectedParams["code_challenge"] = codeChallenge
-			expectedParams["code_challenge_method"] = "S256"
-		}
-
-		for key, expectedValue := range expectedParams {
-			if value := params.Get(key); value != expectedValue {
-				t.Errorf("Parameter %s has incorrect value. Expected: %s, Got: %s",
-					key, expectedValue, value)
-			}
-		}
-
-		// Verify scope parameter separately due to it being space-separated values
-		scope := params.Get("scope")
-		if scope == "" {
-			t.Error("Scope parameter missing from Google auth URL")
-		}
-
-		// Check that all required scopes are present
-		scopeList := strings.Split(scope, " ")
-		expectedScopes := []string{"openid", "profile", "email"}
-		for _, expectedScope := range expectedScopes {
-			found := false
-			for _, actualScope := range scopeList {
-				if actualScope == expectedScope {
-					found = true
-					break
-				}
-			}
-			if !found {
-				t.Errorf("Expected scope '%s' not found in scope parameter: %s", expectedScope, scope)
-			}
-		}
-
-		// Verify offline_access is NOT in the scope list
-		for _, actualScope := range scopeList {
-			if actualScope == "offline_access" {
-				t.Errorf("offline_access scope incorrectly included in Google auth URL: %s", scope)
-			}
-		}
-	})
-
-	// Integration test with mocked Google provider
-	t.Run("Integration test with mocked Google provider", func(t *testing.T) {
-		// Generate an RSA key for signing the test JWTs
-		rsaPrivateKey, err := rsa.GenerateKey(rand.Reader, 2048)
-		if err != nil {
-			t.Fatalf("Failed to generate RSA key: %v", err)
-		}
-
-		// Create JWK for the RSA public key
-		jwk := JWK{
-			Kty: "RSA",
-			Kid: "test-key-id",
-			Alg: "RS256",
-			N:   base64.RawURLEncoding.EncodeToString(rsaPrivateKey.PublicKey.N.Bytes()),
-			E:   base64.RawURLEncoding.EncodeToString(bigIntToBytes(big.NewInt(int64(rsaPrivateKey.PublicKey.E)))),
-		}
-		jwks := &JWKSet{
-			Keys: []JWK{jwk},
-		}
-
-		// Create a mock JWK cache
-		mockJWKCache := &MockJWKCache{
-			JWKS: jwks,
-			Err:  nil,
-		}
-
-		// Create a complete test instance with all required fields
-		mockLogger := NewLogger("debug")
-		googleTOidc := &TraefikOidc{
-			issuerURL:          "https://accounts.google.com",
-			clientID:           "test-client-id",
-			clientSecret:       "test-client-secret",
-			logger:             mockLogger,
-			scopes:             []string{"openid", "profile", "email"},
-			refreshGracePeriod: 60,
-			tokenCache:         NewTokenCache(), // Initialize tokenCache
-			tokenBlacklist:     NewCache(),      // Initialize tokenBlacklist
-			enablePKCE:         false,
-			limiter:            rate.NewLimiter(rate.Inf, 0), // No rate limiting for tests
-			jwkCache:           mockJWKCache,
-			jwksURL:            "https://accounts.google.com/jwks",
-		}
-
-		// Create a session manager
-		sessionManager, _ := NewSessionManager("0123456789abcdef0123456789abcdef", true, mockLogger)
-		googleTOidc.sessionManager = sessionManager
-
-		// Create a mock token verifier
-		mockTokenVerifier := &MockTokenVerifier{
-			VerifyFunc: func(token string) error {
-				return nil // Always verify successfully for this test
-			},
-		}
-		googleTOidc.tokenVerifier = mockTokenVerifier
-
-		// Create JWT tokens for the test
-		now := time.Now()
-		exp := now.Add(1 * time.Hour).Unix()
-		iat := now.Unix()
-		nbf := now.Unix()
-
-		// Create initial ID token
-		initialIDToken, err := createTestJWT(rsaPrivateKey, "RS256", "test-key-id", map[string]interface{}{
-			"iss":   "https://accounts.google.com",
-			"aud":   "test-client-id",
-			"exp":   exp,
-			"iat":   iat,
-			"nbf":   nbf,
-			"sub":   "test-subject",
-			"email": "[email protected]",
-			"nonce": "nonce123", // For initial authentication verification
-			"jti":   generateRandomString(16),
-		})
-		if err != nil {
-			t.Fatalf("Failed to create test ID token: %v", err)
-		}
-
-		// Create refresh ID token
-		refreshedIDToken, err := createTestJWT(rsaPrivateKey, "RS256", "test-key-id", map[string]interface{}{
-			"iss":   "https://accounts.google.com",
-			"aud":   "test-client-id",
-			"exp":   exp,
-			"iat":   iat,
-			"nbf":   nbf,
-			"sub":   "test-subject",
-			"email": "[email protected]",
-			"jti":   generateRandomString(16),
-		})
-		if err != nil {
-			t.Fatalf("Failed to create refreshed test ID token: %v", err)
-		}
-
-		// Set up token verifier with mock
-		googleTOidc.tokenVerifier = &MockTokenVerifier{
-			VerifyFunc: func(token string) error {
-				return nil // Always verify successfully for this test
-			},
-		}
-
-		// Set up JWT verifier with mock
-		googleTOidc.jwtVerifier = &MockJWTVerifier{
-			VerifyJWTFunc: func(jwt *JWT, token string) error {
-				return nil // Always verify successfully for this test
-			},
-		}
-
-		// Create a mock token exchanger that simulates Google's OAuth behavior
-		mockTokenExchanger := &MockTokenExchanger{
-			ExchangeCodeFunc: func(ctx context.Context, grantType, codeOrToken, redirectURL, codeVerifier string) (*TokenResponse, error) {
-				// Verify the correct parameters are passed
-				if grantType != "authorization_code" {
-					t.Errorf("Expected grant_type=authorization_code, got %s", grantType)
-				}
-				if codeOrToken != "test_auth_code" {
-					t.Errorf("Expected code=test_auth_code, got %s", codeOrToken)
-				}
-				if redirectURL != "https://example.com/callback" {
-					t.Errorf("Expected redirect_uri=https://example.com/callback, got %s", redirectURL)
-				}
-
-				// Return a successful token response with a proper JWT
-				return &TokenResponse{
-					IDToken:      initialIDToken,
-					AccessToken:  initialIDToken, // Use a valid JWT as the access token too
-					RefreshToken: "google_refresh_token",
-					ExpiresIn:    3600,
-				}, nil
-			},
-			RefreshTokenFunc: func(refreshToken string) (*TokenResponse, error) {
-				// Verify the correct refresh token is passed
-				if refreshToken != "google_refresh_token" {
-					t.Errorf("Expected refresh_token=google_refresh_token, got %s", refreshToken)
-				}
-
-				// Return a successful refresh response with a proper JWT
-				return &TokenResponse{
-					IDToken:      refreshedIDToken,
-					AccessToken:  refreshedIDToken, // Use a valid JWT as the access token
-					RefreshToken: "",               // Google doesn't always return a new refresh token
-					ExpiresIn:    3600,
-				}, nil
-			},
-		}
-
-		googleTOidc.tokenExchanger = mockTokenExchanger
-
-		// Use the real extractClaimsFunc to parse the proper JWT tokens
-		googleTOidc.extractClaimsFunc = extractClaims
-
-		// 1. Test building the authorization URL
-		authURL := googleTOidc.buildAuthURL("https://example.com/callback", "state123", "nonce123", "")
-
-		// Verify Google-specific parameters
-		if !strings.Contains(authURL, "access_type=offline") {
-			t.Errorf("Google auth URL missing access_type=offline: %s", authURL)
-		}
-		if !strings.Contains(authURL, "prompt=consent") {
-			t.Errorf("Google auth URL missing prompt=consent: %s", authURL)
-		}
-		if strings.Contains(authURL, "offline_access") {
-			t.Errorf("Google auth URL incorrectly includes offline_access scope: %s", authURL)
-		}
-
-		// 2. Test handling the callback and token exchange
-		// Create a request and response recorder for the callback
-		req := httptest.NewRequest("GET", "/callback?code=test_auth_code&state=state123", nil)
-		rw := httptest.NewRecorder()
-
-		// Create a session and set the necessary values
-		session, _ := googleTOidc.sessionManager.GetSession(req)
-		session.SetCSRF("state123") // Must match the state parameter
-		session.SetNonce("nonce123")
-
-		// Save the session to the request
-		if err := session.Save(req, rw); err != nil {
-			t.Fatalf("Failed to save session: %v", err)
-		}
-
-		// Get cookies from the response and add them to a new request
-		cookies := rw.Result().Cookies()
-		callbackReq := httptest.NewRequest("GET", "/callback?code=test_auth_code&state=state123", nil)
-		for _, cookie := range cookies {
-			callbackReq.AddCookie(cookie)
-		}
-		callbackRw := httptest.NewRecorder()
-
-		// Handle the callback
-		googleTOidc.handleCallback(callbackRw, callbackReq, "https://example.com/callback")
-
-		// Verify the response is a redirect (302 Found)
-		if callbackRw.Code != 302 {
-			t.Errorf("Expected 302 redirect, got %d", callbackRw.Code)
-		}
-
-		// Create a new request to get the updated session
-		newReq := httptest.NewRequest("GET", "/", nil)
-		for _, cookie := range callbackRw.Result().Cookies() {
-			newReq.AddCookie(cookie)
-		}
-
-		// Get the updated session
-		newSession, err := googleTOidc.sessionManager.GetSession(newReq)
-		if err != nil {
-			t.Fatalf("Failed to get session after callback: %v", err)
-		}
-
-		// Verify the session contains the expected values
-		if !newSession.GetAuthenticated() {
-			t.Error("Session not marked as authenticated after callback")
-		}
-		if newSession.GetEmail() != "[email protected]" {
-			t.Errorf("Session email incorrect: got %s, expected [email protected]",
-				newSession.GetEmail())
-		}
-
-		// Check for non-empty access token that can be parsed as JWT
-		accessToken := newSession.GetAccessToken()
-		if accessToken == "" {
-			t.Error("Session access token is empty")
-		} else {
-			claims, err := extractClaims(accessToken)
-			if err != nil {
-				t.Errorf("Failed to parse access token as JWT: %v", err)
-			} else if email, ok := claims["email"].(string); !ok || email != "[email protected]" {
-				t.Errorf("Access token JWT doesn't contain expected email claim")
-			}
-		}
-
-		// Check refresh token
-		if newSession.GetRefreshToken() != "google_refresh_token" {
-			t.Errorf("Session refresh token incorrect: got %s, expected google_refresh_token",
-				newSession.GetRefreshToken())
-		}
-
-		// 3. Test token refresh
-		refreshReq := httptest.NewRequest("GET", "/", nil)
-		for _, cookie := range callbackRw.Result().Cookies() {
-			refreshReq.AddCookie(cookie)
-		}
-		refreshRw := httptest.NewRecorder()
-
-		// Get the session for refresh
-		refreshSession, _ := googleTOidc.sessionManager.GetSession(refreshReq)
-
-		// Refresh the token
-		refreshed := googleTOidc.refreshToken(refreshRw, refreshReq, refreshSession)
-
-		// Verify refresh was successful
-		if !refreshed {
-			t.Error("Token refresh failed")
-		}
-
-		// Verify the session data after refresh
-		// Check for non-empty refreshed access token that can be parsed as JWT
-		refreshedAccessToken := refreshSession.GetAccessToken()
-		if refreshedAccessToken == "" {
-			t.Error("Session access token is empty after refresh")
-		} else {
-			claims, err := extractClaims(refreshedAccessToken)
-			if err != nil {
-				t.Errorf("Failed to parse refreshed access token as JWT: %v", err)
-			} else if email, ok := claims["email"].(string); !ok || email != "[email protected]" {
-				t.Errorf("Refreshed access token JWT doesn't contain expected email claim")
-			}
-		}
-
-		// Since Google didn't return a new refresh token, the original should be preserved
-		if refreshSession.GetRefreshToken() != "google_refresh_token" {
-			t.Errorf("Original refresh token not preserved: got %s, expected google_refresh_token",
-				refreshSession.GetRefreshToken())
-		}
-	})
-}
-
-// No need to redefine MockTokenExchanger - it's already defined in main_test.go

helpers.go

@@ -15,14 +15,10 @@ import (
 	"time"
 )
 
-// generateNonce creates a cryptographically secure random string suitable for use as an OIDC nonce.
+// generateNonce creates a cryptographically secure random nonce
-// The nonce is used during the authentication flow to mitigate replay attacks by associating
+// for use in the OIDC authentication flow. The nonce is used to
-// the ID token with the specific authentication request.
+// prevent replay attacks by ensuring the token received matches
-// It generates 32 random bytes and encodes them using base64 URL encoding.
+// the authentication request.
-//
-// Returns:
-//   - A base64 URL encoded random string (nonce).
-//   - An error if the random byte generation fails.
 func generateNonce() (string, error) {
 	nonceBytes := make([]byte, 32)
 	_, err := rand.Read(nonceBytes)
@@ -32,13 +28,9 @@ func generateNonce() (string, error) {
 	return base64.URLEncoding.EncodeToString(nonceBytes), nil
 }
 
-// generateCodeVerifier creates a cryptographically secure random string suitable for use as a PKCE code verifier.
+// generateCodeVerifier creates a cryptographically secure random string
-// According to RFC 7636, the verifier should be a high-entropy string between 43 and 128 characters long.
+// for use as a PKCE code verifier. The code verifier must be between 43 and 128
-// This function generates 32 random bytes, resulting in a 43-character base64 URL encoded string.
+// characters long, per the PKCE spec (RFC 7636).
-//
-// Returns:
-//   - A base64 URL encoded random string (code verifier).
-//   - An error if the random byte generation fails.
 func generateCodeVerifier() (string, error) {
 	// Using 32 bytes (256 bits) will produce a 43 character base64url string
 	verifierBytes := make([]byte, 32)
@@ -49,15 +41,8 @@ func generateCodeVerifier() (string, error) {
 	return base64.RawURLEncoding.EncodeToString(verifierBytes), nil
 }
 
-// deriveCodeChallenge computes the PKCE code challenge from a given code verifier.
+// deriveCodeChallenge creates a code challenge from a code verifier
-// It uses the S256 challenge method (SHA-256 hash followed by base64 URL encoding)
+// using the SHA-256 method as specified in the PKCE standard (RFC 7636).
-// as defined in RFC 7636.
-//
-// Parameters:
-//   - codeVerifier: The high-entropy string generated by generateCodeVerifier.
-//
-// Returns:
-//   - The base64 URL encoded SHA-256 hash of the code verifier (code challenge).
 func deriveCodeChallenge(codeVerifier string) string {
 	// Calculate SHA-256 hash of the code verifier
 	hasher := sha256.New()
@@ -88,23 +73,15 @@ type TokenResponse struct {
 	TokenType string `json:"token_type"`
 }
 
-// exchangeTokens performs the OAuth 2.0 token exchange with the OIDC provider's token endpoint.
+// exchangeTokens performs the OAuth 2.0 token exchange with the OIDC provider.
-// It handles both the "authorization_code" grant type (exchanging an authorization code for tokens)
+// It supports both authorization code and refresh token grant types.
-// and the "refresh_token" grant type (using a refresh token to obtain new tokens).
-// It includes necessary parameters like client credentials and handles PKCE verification if applicable.
-// The function follows redirects and handles potential errors during the exchange.
-//
 // Parameters:
-//   - ctx: The context for the outgoing HTTP request.
+//   - ctx: Context for the HTTP request
-//   - grantType: The OAuth 2.0 grant type ("authorization_code" or "refresh_token").
+//   - grantType: The OAuth 2.0 grant type ("authorization_code" or "refresh_token")
-//   - codeOrToken: The authorization code (for "authorization_code" grant) or the refresh token (for "refresh_token" grant).
+//   - codeOrToken: Either the authorization code or refresh token
-//   - redirectURL: The redirect URI that was used in the initial authorization request (required for "authorization_code" grant).
+//   - redirectURL: The callback URL for authorization code grant
-//   - codeVerifier: The PKCE code verifier (required for "authorization_code" grant if PKCE was used).
+//   - codeVerifier: Optional PKCE code verifier for authorization code grant
-//
+func (t *TraefikOidc) exchangeTokens(ctx context.Context, grantType, codeOrToken, redirectURL string, codeVerifier string) (*TokenResponse, error) {
-// Returns:
-//   - A TokenResponse containing the obtained tokens (ID, access, refresh).
-//   - An error if the token exchange fails (e.g., network error, provider error, invalid grant).
-func (t *TraefikOidc) exchangeTokens(ctx context.Context, grantType string, codeOrToken string, redirectURL string, codeVerifier string) (*TokenResponse, error) {
 	data := url.Values{
 		"grant_type":    {grantType},
 		"client_id":     {t.clientID},
@@ -163,16 +140,8 @@ func (t *TraefikOidc) exchangeTokens(ctx context.Context, grantType string, code
 	return &tokenResponse, nil
 }
 
-// getNewTokenWithRefreshToken uses a refresh token to obtain a new set of tokens (ID, access, refresh)
+// getNewTokenWithRefreshToken obtains new tokens using a refresh token.
-// from the OIDC provider's token endpoint. It wraps the exchangeTokens function with the
+// This is used to refresh access tokens before they expire.
-// "refresh_token" grant type.
-//
-// Parameters:
-//   - refreshToken: The refresh token previously obtained during authentication or a prior refresh.
-//
-// Returns:
-//   - A TokenResponse containing the newly obtained tokens.
-//   - An error if the refresh operation fails.
 func (t *TraefikOidc) getNewTokenWithRefreshToken(refreshToken string) (*TokenResponse, error) {
 	ctx := context.Background()
 	tokenResponse, err := t.exchangeTokens(ctx, "refresh_token", refreshToken, "", "")
@@ -184,17 +153,179 @@ func (t *TraefikOidc) getNewTokenWithRefreshToken(refreshToken string) (*TokenRe
 	return tokenResponse, nil
 }
 
-// extractClaims decodes the payload (claims set) part of a JWT string.
+// handleExpiredToken manages token expiration by clearing the session
-// It splits the JWT into its three parts, base64 URL decodes the second part (payload),
+// and initiating a new authentication flow.
-// and unmarshals the resulting JSON into a map.
+func (t *TraefikOidc) handleExpiredToken(rw http.ResponseWriter, req *http.Request, session *SessionData, redirectURL string) {
-// Note: This function does *not* validate the token's signature or claims.
+	// Clear authentication data but preserve CSRF state
-//
+	session.SetAuthenticated(false)
-// Parameters:
+	session.SetAccessToken("")
-//   - tokenString: The raw JWT string.
+	session.SetRefreshToken("")
-//
+	session.SetEmail("")
-// Returns:
+
-//   - A map representing the JSON claims extracted from the token payload.
+	// Save the cleared session state
-//   - An error if the token format is invalid, decoding fails, or JSON unmarshaling fails.
+	if err := session.Save(req, rw); err != nil {
+		t.logger.Errorf("Failed to save cleared session: %v", err)
+		http.Error(rw, "Internal Server Error", http.StatusInternalServerError)
+		return
+	}
+
+	t.defaultInitiateAuthentication(rw, req, session, redirectURL)
+}
+
+// handleCallback processes the authentication callback from the OIDC provider.
+// It validates the callback parameters, exchanges the authorization code for
+// tokens, verifies the tokens, and establishes the user's session.
+func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request, redirectURL string) {
+	session, err := t.sessionManager.GetSession(req)
+	if err != nil {
+		t.logger.Errorf("Session error: %v", err)
+		http.Error(rw, "Session error", http.StatusInternalServerError)
+		return
+	}
+
+	t.logger.Debugf("Handling callback, URL: %s", req.URL.String())
+
+	// Check for errors in the callback
+	if req.URL.Query().Get("error") != "" {
+		errorDescription := req.URL.Query().Get("error_description")
+		t.logger.Errorf("Authentication error: %s - %s", req.URL.Query().Get("error"), errorDescription)
+		http.Error(rw, fmt.Sprintf("Authentication error: %s", errorDescription), http.StatusBadRequest)
+		return
+	}
+
+	// Validate CSRF state
+	state := req.URL.Query().Get("state")
+	if state == "" {
+		t.logger.Error("No state in callback")
+		http.Error(rw, "State parameter missing in callback", http.StatusBadRequest)
+		return
+	}
+
+	csrfToken := session.GetCSRF()
+	if csrfToken == "" {
+		t.logger.Error("CSRF token missing in session")
+		http.Error(rw, "CSRF token missing", http.StatusBadRequest)
+		return
+	}
+
+	if state != csrfToken {
+		t.logger.Error("State parameter does not match CSRF token in session")
+		http.Error(rw, "Invalid state parameter", http.StatusBadRequest)
+		return
+	}
+
+	// Exchange code for tokens
+	code := req.URL.Query().Get("code")
+	if code == "" {
+		t.logger.Error("No code in callback")
+		http.Error(rw, "No code in callback", http.StatusBadRequest)
+		return
+	}
+
+	// Get the code verifier from the session for PKCE flow
+	codeVerifier := session.GetCodeVerifier()
+
+	tokenResponse, err := t.exchangeCodeForTokenFunc(code, redirectURL, codeVerifier)
+	if err != nil {
+		t.logger.Errorf("Failed to exchange code for token: %v", err)
+		http.Error(rw, "Authentication failed", http.StatusInternalServerError)
+		return
+	}
+
+	// Verify tokens and claims
+	if err := t.verifyToken(tokenResponse.IDToken); err != nil {
+		t.logger.Errorf("Failed to verify id_token: %v", err)
+		http.Error(rw, "Authentication failed", http.StatusInternalServerError)
+		return
+	}
+
+	claims, err := t.extractClaimsFunc(tokenResponse.IDToken)
+	if err != nil {
+		t.logger.Errorf("Failed to extract claims: %v", err)
+		http.Error(rw, "Authentication failed", http.StatusInternalServerError)
+		return
+	}
+
+	// Verify nonce to prevent replay attacks
+	nonceClaim, ok := claims["nonce"].(string)
+	if !ok || nonceClaim == "" {
+		t.logger.Error("Nonce claim missing in id_token")
+		http.Error(rw, "Authentication failed", http.StatusInternalServerError)
+		return
+	}
+
+	sessionNonce := session.GetNonce()
+	if sessionNonce == "" {
+		t.logger.Error("Nonce not found in session")
+		http.Error(rw, "Authentication failed", http.StatusInternalServerError)
+		return
+	}
+
+	if nonceClaim != sessionNonce {
+		t.logger.Error("Nonce claim does not match session nonce")
+		http.Error(rw, "Authentication failed", http.StatusInternalServerError)
+		return
+	}
+
+	// Validate user's email domain
+	email, _ := claims["email"].(string)
+	if email == "" || !t.isAllowedDomain(email) {
+		t.logger.Errorf("Invalid or disallowed email: %s", email)
+		http.Error(rw, "Authentication failed: Invalid or disallowed email", http.StatusForbidden)
+		return
+	}
+
+	// Update session with authentication data
+	session.SetAuthenticated(true)
+	session.SetEmail(email)
+	session.SetAccessToken(tokenResponse.IDToken)
+	session.SetRefreshToken(tokenResponse.RefreshToken)
+
+	if err := session.Save(req, rw); err != nil {
+		t.logger.Errorf("Failed to save session: %v", err)
+		http.Error(rw, "Failed to save session", http.StatusInternalServerError)
+		return
+	}
+	// Redirect to original path or root
+	redirectPath := "/"
+	if incomingPath := session.GetIncomingPath(); incomingPath != "" && incomingPath != t.redirURLPath {
+		redirectPath = incomingPath
+	}
+
+	// For redirecting, we need to ensure URL fragments are preserved
+	// To do this, we'll use a small JavaScript snippet that preserves any URL fragments
+	// This is necessary because URL fragments are not sent to the server
+	rw.Header().Set("Content-Type", "text/html; charset=utf-8")
+	rw.WriteHeader(http.StatusOK)
+	fmt.Fprintf(rw, `<!DOCTYPE html>
+<html>
+<head>
+<title>Authentication Complete</title>
+<script>
+// Preserve URL fragments by combining the redirectPath with any fragment in the current URL
+(function() {
+ var redirectPath = %q;
+ var redirectUrl = new URL(redirectPath, window.location.href);
+ 
+ // If we have a hash in the current URL, and the redirect path doesn't already have one,
+ // append the hash to the redirect URL to preserve anchors
+ if (window.location.hash && !redirectPath.includes('#')) {
+   redirectUrl.hash = window.location.hash;
+ }
+ 
+ window.location.replace(redirectUrl.toString());
+})();
+</script>
+</head>
+<body>
+<p>Authentication successful. Redirecting...</p>
+</body>
+</html>`, redirectPath)
+	http.Redirect(rw, req, redirectPath, http.StatusFound)
+}
+
+// extractClaims parses a JWT token and extracts its claims.
+// It handles base64url decoding and JSON parsing of the token payload.
 func extractClaims(tokenString string) (map[string]interface{}, error) {
 	parts := strings.Split(tokenString, ".")
 	if len(parts) != 3 {
@@ -222,36 +353,21 @@ type TokenCache struct {
 	cache *Cache
 }
 
-// NewTokenCache creates and initializes a new TokenCache.
+// NewTokenCache creates a new TokenCache instance.
-// It internally creates a new generic Cache instance for storage.
 func NewTokenCache() *TokenCache {
 	return &TokenCache{
 		cache: NewCache(),
 	}
 }
 
-// Set stores the claims associated with a specific token string in the cache.
+// Set stores a token's claims in the cache with an expiration time.
-// It prefixes the token string to avoid potential collisions with other cache types
-// and sets the provided expiration duration.
-//
-// Parameters:
-//   - token: The raw token string (used as the key).
-//   - claims: The map of claims associated with the token.
-//   - expiration: The duration for which the cache entry should be valid.
 func (tc *TokenCache) Set(token string, claims map[string]interface{}, expiration time.Duration) {
 	token = "t-" + token
 	tc.cache.Set(token, claims, expiration)
 }
 
-// Get retrieves the cached claims for a given token string.
+// Get retrieves a token's claims from the cache.
-// It prefixes the token string before querying the underlying cache.
+// Returns the claims and a boolean indicating if the token was found.
-//
-// Parameters:
-//   - token: The raw token string to look up.
-//
-// Returns:
-//   - The cached claims map if found and valid.
-//   - A boolean indicating whether the token was found in the cache (true if found, false otherwise).
 func (tc *TokenCache) Get(token string) (map[string]interface{}, bool) {
 	token = "t-" + token
 	value, found := tc.cache.Get(token)
@@ -262,39 +378,20 @@ func (tc *TokenCache) Get(token string) (map[string]interface{}, bool) {
 	return claims, ok
 }
 
-// Delete removes the cached entry for a specific token string.
+// Delete removes a token from the cache.
-// It prefixes the token string before calling the underlying cache's Delete method.
-//
-// Parameters:
-//   - token: The raw token string to remove from the cache.
 func (tc *TokenCache) Delete(token string) {
 	token = "t-" + token
 	tc.cache.Delete(token)
 }
 
-// Cleanup triggers the cleanup process for the underlying generic cache,
+// Cleanup removes expired tokens from the cache.
-// removing expired token entries.
 func (tc *TokenCache) Cleanup() {
 	tc.cache.Cleanup()
 }
 
-// Close stops the cleanup goroutine in the underlying cache.
+// exchangeCodeForToken exchanges an authorization code for tokens.
-func (tc *TokenCache) Close() {
+// It handles PKCE (Proof Key for Code Exchange) based on middleware configuration.
-	tc.cache.Close()
+// The code verifier is only included in the token request if PKCE is enabled.
-}
-
-// exchangeCodeForToken is a convenience function that wraps exchangeTokens specifically
-// for the "authorization_code" grant type. It handles the conditional inclusion of the
-// PKCE code verifier based on the middleware's configuration (t.enablePKCE).
-//
-// Parameters:
-//   - code: The authorization code received from the OIDC provider.
-//   - redirectURL: The redirect URI used in the initial authorization request.
-//   - codeVerifier: The PKCE code verifier stored in the session (if PKCE is enabled).
-//
-// Returns:
-//   - A TokenResponse containing the obtained tokens.
-//   - An error if the code exchange fails.
 func (t *TraefikOidc) exchangeCodeForToken(code string, redirectURL string, codeVerifier string) (*TokenResponse, error) {
 	ctx := context.Background()
 
@@ -311,15 +408,8 @@ func (t *TraefikOidc) exchangeCodeForToken(code string, redirectURL string, code
 	return tokenResponse, nil
 }
 
-// createStringMap converts a slice of strings into a map[string]struct{} (a set).
+// createStringMap creates a map from a slice of strings.
-// This is useful for creating efficient lookups (O(1) average time complexity)
+// Used for efficient lookups in allowed domains and roles.
-// for checking the presence of items like allowed domains, roles, or groups.
-//
-// Parameters:
-//   - keys: A slice of strings to be added to the set.
-//
-// Returns:
-//   - A map where the keys are the strings from the input slice and the values are empty structs.
 func createStringMap(keys []string) map[string]struct{} {
 	result := make(map[string]struct{})
 	for _, key := range keys {
@@ -328,17 +418,9 @@ func createStringMap(keys []string) map[string]struct{} {
 	return result
 }
 
-// handleLogout processes requests to the configured logout path.
+// handleLogout manages the OIDC logout process.
-// It performs the following steps:
+// It clears the session and redirects either to the OIDC provider's
-//  1. Retrieves the current user session.
+// end session endpoint (if available) or to the configured post-logout URL.
-//  2. Gets the access token (ID token hint) from the session.
-//  3. Clears all authentication-related data from the session cookies.
-//  4. Determines the final post-logout redirect URI.
-//  5. If an OIDC end_session_endpoint is configured and an ID token hint is available,
-//     it builds the OIDC logout URL and redirects the user agent to the provider for logout.
-//  6. Otherwise, it redirects the user agent directly to the post-logout redirect URI.
-//
-// It handles potential errors during session retrieval or clearing.
 func (t *TraefikOidc) handleLogout(rw http.ResponseWriter, req *http.Request) {
 	session, err := t.sessionManager.GetSession(req)
 	if err != nil {
@@ -380,18 +462,11 @@ func (t *TraefikOidc) handleLogout(rw http.ResponseWriter, req *http.Request) {
 	http.Redirect(rw, req, postLogoutRedirectURI, http.StatusFound)
 }
 
-// BuildLogoutURL constructs the URL for redirecting the user agent to the OIDC provider's
+// BuildLogoutURL constructs the OIDC end session URL with appropriate parameters.
-// end_session_endpoint, including the required id_token_hint and optional
-// post_logout_redirect_uri parameters as query arguments.
-//
 // Parameters:
-//   - endSessionURL: The URL of the OIDC provider's end session endpoint.
+//   - endSessionURL: The OIDC provider's end session endpoint
-//   - idToken: The ID token previously issued to the user (used as id_token_hint).
+//   - idToken: The ID token to be invalidated
-//   - postLogoutRedirectURI: The optional URI where the provider should redirect the user agent after logout.
+//   - postLogoutRedirectURI: Where to redirect after logout completes
-//
-// Returns:
-//   - The fully constructed logout URL string.
-//   - An error if the provided endSessionURL is invalid.
 func BuildLogoutURL(endSessionURL, idToken, postLogoutRedirectURI string) (string, error) {
 	u, err := url.Parse(endSessionURL)
 	if err != nil {

helpers_test.go

@@ -1,17 +1,227 @@
 package traefikoidc
 
 import (
-	"crypto/rand"
+	"fmt"
-	"encoding/hex"
+	"runtime"
+	"testing"
+	"time"
 )
 
-// generateRandomString generates a random string of the specified length
+func TestTokenBlacklistSizeLimit(t *testing.T) {
-// This is used in tests to create unique identifiers
+	tb := NewTokenBlacklist()
-func generateRandomString(length int) string {
+
-	bytes := make([]byte, length/2)
+	// Add tokens up to maxSize
-	if _, err := rand.Read(bytes); err != nil {
+	for i := 0; i < 1000; i++ {
-		// In tests, fallback to a predictable string if random fails
+		tb.Add(fmt.Sprintf("token%d", i), time.Now().Add(time.Hour))
-		return "random-string-fallback"
+	}
+
+	// Verify size is at max
+	if tb.Count() != 1000 {
+		t.Errorf("Expected blacklist size to be 1000, got %d", tb.Count())
+	}
+
+	// Add one more token, should trigger cleanup/eviction
+	tb.Add("newtoken", time.Now().Add(time.Hour))
+
+	// Size should still be at max
+	if tb.Count() > 1000 {
+		t.Errorf("Blacklist exceeded max size: %d", tb.Count())
+	}
+}
+
+func TestTokenBlacklistExpiredCleanup(t *testing.T) {
+	tb := NewTokenBlacklist()
+
+	// Add some expired tokens
+	for i := 0; i < 500; i++ {
+		tb.Add(fmt.Sprintf("expired%d", i), time.Now().Add(-time.Hour))
+	}
+
+	// Add some valid tokens
+	for i := 0; i < 500; i++ {
+		tb.Add(fmt.Sprintf("valid%d", i), time.Now().Add(time.Hour))
+	}
+
+	// Force cleanup
+	tb.Cleanup()
+
+	// Only valid tokens should remain
+	if tb.Count() != 500 {
+		t.Errorf("Expected 500 valid tokens after cleanup, got %d", tb.Count())
+	}
+
+	// Verify only valid tokens remain
+	tb.mutex.RLock()
+	defer tb.mutex.RUnlock()
+	for token, expiry := range tb.tokens {
+		if time.Now().After(expiry) {
+			t.Errorf("Found expired token after cleanup: %s", token)
+		}
+	}
+}
+
+func TestTokenBlacklistOldestEviction(t *testing.T) {
+	tb := NewTokenBlacklist()
+
+	// Add tokens at capacity with different expiration times
+	baseTime := time.Now()
+	oldestToken := "oldest"
+
+	// Add oldest token first
+	tb.Add(oldestToken, baseTime.Add(time.Hour))
+
+	// Fill up to capacity with newer tokens
+	for i := 0; i < 999; i++ {
+		tb.Add(fmt.Sprintf("token%d", i), baseTime.Add(time.Hour*2))
+	}
+
+	// Add a new token that should evict the oldest
+	newToken := "newest"
+	tb.Add(newToken, baseTime.Add(time.Hour*3))
+
+	// Verify oldest token was evicted
+	if tb.IsBlacklisted(oldestToken) {
+		t.Error("Oldest token should have been evicted")
+	}
+
+	// Verify newest token is present
+	if !tb.IsBlacklisted(newToken) {
+		t.Error("Newest token should be present")
+	}
+}
+
+func TestTokenBlacklistMemoryUsage(t *testing.T) {
+	tb := NewTokenBlacklist()
+	iterations := 10000
+
+	// Force initial GC
+	runtime.GC()
+
+	// Record initial memory stats
+	var m1, m2 runtime.MemStats
+	runtime.ReadMemStats(&m1)
+
+	// Simulate heavy usage
+	for i := 0; i < iterations; i++ {
+		// Add new token
+		tb.Add(fmt.Sprintf("token%d", i), time.Now().Add(time.Hour))
+
+		// Periodically check blacklisted status
+		if i%100 == 0 {
+			tb.IsBlacklisted(fmt.Sprintf("token%d", i-50))
+		}
+
+		// Periodically cleanup
+		if i%1000 == 0 {
+			tb.Cleanup()
+		}
+	}
+
+	// Force GC and wait for it to complete
+	runtime.GC()
+	time.Sleep(100 * time.Millisecond)
+	runtime.ReadMemStats(&m2)
+
+	// Check memory growth (using HeapAlloc for more accurate measurement)
+	memoryGrowth := int64(m2.HeapAlloc - m1.HeapAlloc)
+	maxAllowedGrowth := int64(2 * 1024 * 1024) // 2MB max growth
+
+	if memoryGrowth > maxAllowedGrowth {
+		t.Logf("Initial HeapAlloc: %d, Final HeapAlloc: %d", m1.HeapAlloc, m2.HeapAlloc)
+		t.Errorf("Excessive memory growth: %d bytes", memoryGrowth)
+	}
+
+	// Verify size stayed within limits
+	if tb.Count() > 1000 {
+		t.Errorf("Blacklist exceeded max size: %d", tb.Count())
+	}
+}
+
+func TestConcurrentTokenBlacklistOperations(t *testing.T) {
+	tb := NewTokenBlacklist()
+	iterations := 1000
+	concurrency := 10
+	done := make(chan bool)
+
+	// Start multiple goroutines performing operations
+	for i := 0; i < concurrency; i++ {
+		go func(id int) {
+			for j := 0; j < iterations; j++ {
+				// Add tokens
+				token := fmt.Sprintf("token%d-%d", id, j)
+				tb.Add(token, time.Now().Add(time.Hour))
+
+				// Check blacklist status
+				tb.IsBlacklisted(token)
+
+				// Periodic cleanup
+				if j%100 == 0 {
+					tb.Cleanup()
+				}
+			}
+			done <- true
+		}(i)
+	}
+
+	// Wait for all goroutines to complete
+	for i := 0; i < concurrency; i++ {
+		<-done
+	}
+
+	// Verify size constraints were maintained
+	if tb.Count() > 1000 {
+		t.Errorf("Blacklist exceeded max size under concurrent operations: %d", tb.Count())
+	}
+}
+
+func TestTokenCacheMemoryUsage(t *testing.T) {
+	tc := NewTokenCache()
+	iterations := 10000
+
+	// Force initial GC
+	runtime.GC()
+
+	// Record initial memory stats
+	var m1, m2 runtime.MemStats
+	runtime.ReadMemStats(&m1)
+
+	// Simulate heavy cache usage
+	for i := 0; i < iterations; i++ {
+		claims := map[string]interface{}{
+			"sub": fmt.Sprintf("user%d", i),
+			"exp": time.Now().Add(time.Hour).Unix(),
+		}
+
+		// Add to cache
+		tc.Set(fmt.Sprintf("token%d", i), claims, time.Hour)
+
+		// Periodically retrieve
+		if i%100 == 0 {
+			tc.Get(fmt.Sprintf("token%d", i-50))
+		}
+
+		// Periodically cleanup
+		if i%1000 == 0 {
+			tc.Cleanup()
+		}
+	}
+
+	// Force GC and wait for it to complete
+	runtime.GC()
+	time.Sleep(100 * time.Millisecond)
+	runtime.ReadMemStats(&m2)
+
+	// Check memory growth (using HeapAlloc for more accurate measurement)
+	memoryGrowth := int64(m2.HeapAlloc - m1.HeapAlloc)
+	maxAllowedGrowth := int64(2 * 1024 * 1024) // 2MB max growth
+
+	if memoryGrowth > maxAllowedGrowth {
+		t.Logf("Initial HeapAlloc: %d, Final HeapAlloc: %d", m1.HeapAlloc, m2.HeapAlloc)
+		t.Errorf("Excessive cache memory growth: %d bytes", memoryGrowth)
+	}
+
+	// Verify cache size stayed within limits
+	if len(tc.cache.items) > tc.cache.maxSize {
+		t.Errorf("Cache exceeded max size: %d", len(tc.cache.items))
 	}
-	return hex.EncodeToString(bytes)
 }

jwk.go

@@ -38,48 +38,14 @@ type JWKCache struct {
 	mutex     sync.RWMutex
 	// CacheLifetime is configurable to determine how long the JWKS is cached.
 	CacheLifetime time.Duration
-	internalCache *Cache // To hold the closable Cache instance from cache.go
-	maxSize       int    // Maximum number of items in the cache
 }
 
 type JWKCacheInterface interface {
 	GetJWKS(ctx context.Context, jwksURL string, httpClient *http.Client) (*JWKSet, error)
 	Cleanup()
-	Close()
-}
-
-// GetJWKS retrieves the JSON Web Key Set (JWKS) from the cache or fetches it from the provider.
-// It first checks if a valid, non-expired JWKS is present in the cache. If so, it returns the cached version.
-// Otherwise, it attempts to fetch the JWKS from the specified jwksURL using the provided httpClient.
-// If the fetch is successful, the JWKS is stored in the cache with an expiration time based on CacheLifetime
-// (defaulting to 1 hour if not set) and returned.
-// This method uses double-checked locking to minimize contention when the cache needs refreshing.
-//
-// Parameters:
-//   - ctx: Context for the HTTP request if fetching is required.
-//   - jwksURL: The URL of the OIDC provider's JWKS endpoint.
-//   - httpClient: The HTTP client to use for fetching the JWKS.
-//
-// Returns:
-//   - A pointer to the JWKSet containing the keys.
-//   - An error if fetching fails or the response cannot be decoded.
-func NewJWKCache() *JWKCache {
-	cache := &JWKCache{
-		CacheLifetime: 1 * time.Hour,
-		maxSize:       100, // Default maximum size
-		internalCache: NewCache(),
-	}
-	return cache
 }
 
 func (c *JWKCache) GetJWKS(ctx context.Context, jwksURL string, httpClient *http.Client) (*JWKSet, error) {
-	// First check if we already have cached JWKS for this URL
-	if c.internalCache != nil {
-		if cachedJwks, found := c.internalCache.Get(jwksURL); found {
-			return cachedJwks.(*JWKSet), nil
-		}
-	}
-
 	c.mutex.RLock()
 	if c.jwks != nil && time.Now().Before(c.expiresAt) {
 		defer c.mutex.RUnlock()
@@ -105,16 +71,9 @@ func (c *JWKCache) GetJWKS(ctx context.Context, jwksURL string, httpClient *http
 	}
 	c.expiresAt = time.Now().Add(lifetime)
 
-	// Also store in the internalCache
-	if c.internalCache != nil {
-		c.internalCache.Set(jwksURL, jwks, lifetime)
-	}
-
 	return jwks, nil
 }
 
-// Cleanup removes the cached JWKS if it has expired.
-// This is intended to be called periodically to ensure stale JWKS data is cleared.
 func (c *JWKCache) Cleanup() {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
@@ -125,33 +84,6 @@ func (c *JWKCache) Cleanup() {
 	}
 }
 
-// Close shuts down the cache's auto-cleanup routine.
-func (c *JWKCache) Close() {
-	// Close shuts down the internal cache's auto-cleanup routine, if the cache exists.
-	if c.internalCache != nil {
-		c.internalCache.Close()
-	}
-}
-
-// SetMaxSize sets the maximum number of items in the cache
-func (c *JWKCache) SetMaxSize(size int) {
-	c.maxSize = size
-	if c.internalCache != nil {
-		c.internalCache.maxSize = size
-	}
-}
-
-// fetchJWKS retrieves the JSON Web Key Set (JWKS) from the specified URL.
-// It uses the provided context and HTTP client to make the request.
-//
-// Parameters:
-//   - ctx: Context for the HTTP request.
-//   - jwksURL: The URL of the OIDC provider's JWKS endpoint.
-//   - httpClient: The HTTP client to use for the request.
-//
-// Returns:
-//   - A pointer to the fetched JWKSet.
-//   - An error if the request fails, the status code is not OK, or the response body cannot be decoded.
 func fetchJWKS(ctx context.Context, jwksURL string, httpClient *http.Client) (*JWKSet, error) {
 	// Create a request with context to enforce timeout
 	req, err := http.NewRequestWithContext(ctx, "GET", jwksURL, nil)
@@ -177,16 +109,6 @@ func fetchJWKS(ctx context.Context, jwksURL string, httpClient *http.Client) (*J
 	return &jwks, nil
 }
 
-// jwkToPEM converts a JWK (JSON Web Key) object into PEM (Privacy-Enhanced Mail) format.
-// It selects the appropriate conversion function based on the JWK's key type ("kty").
-// Currently supports "RSA" and "EC" key types.
-//
-// Parameters:
-//   - jwk: A pointer to the JWK object to convert.
-//
-// Returns:
-//   - A byte slice containing the public key in PEM format.
-//   - An error if the key type is unsupported or conversion fails.
 func jwkToPEM(jwk *JWK) ([]byte, error) {
 	converter, ok := jwkConverters[jwk.Kty]
 	if !ok {
@@ -202,17 +124,6 @@ var jwkConverters = map[string]jwkToPEMConverter{
 	"EC":  ecJWKToPEM,
 }
 
-// rsaJWKToPEM converts an RSA JWK into PEM format.
-// It decodes the modulus (n) and exponent (e) from base64 URL encoding,
-// constructs an rsa.PublicKey, marshals it into PKIX format, and then
-// encodes it as a PEM block.
-//
-// Parameters:
-//   - jwk: A pointer to the RSA JWK object (must have "kty": "RSA").
-//
-// Returns:
-//   - A byte slice containing the RSA public key in PEM format.
-//   - An error if decoding parameters fails or key marshaling fails.
 func rsaJWKToPEM(jwk *JWK) ([]byte, error) {
 	nBytes, err := base64.RawURLEncoding.DecodeString(jwk.N)
 	if err != nil {
@@ -244,18 +155,6 @@ func rsaJWKToPEM(jwk *JWK) ([]byte, error) {
 	return pubKeyPEM, nil
 }
 
-// ecJWKToPEM converts an EC (Elliptic Curve) JWK into PEM format.
-// It decodes the X and Y coordinates from base64 URL encoding, determines the
-// elliptic curve based on the "crv" parameter (P-256, P-384, P-521),
-// constructs an ecdsa.PublicKey, marshals it into PKIX format, and then
-// encodes it as a PEM block.
-//
-// Parameters:
-//   - jwk: A pointer to the EC JWK object (must have "kty": "EC").
-//
-// Returns:
-//   - A byte slice containing the EC public key in PEM format.
-//   - An error if decoding parameters fails, the curve is unsupported, or key marshaling fails.
 func ecJWKToPEM(jwk *JWK) ([]byte, error) {
 	xBytes, err := base64.RawURLEncoding.DecodeString(jwk.X)
 	if err != nil {

jwt.go

@@ -15,15 +15,9 @@ import (
 	"time"
 )
 
-var (
+var replayCacheMu sync.Mutex
-	replayCacheMu sync.Mutex
+var replayCache = make(map[string]time.Time)
-	replayCache   = make(map[string]time.Time)
-)
 
-// cleanupReplayCache iterates through the replay cache and removes entries
-// whose expiration time is before the current time. This function should be
-// called periodically to prevent the cache from growing indefinitely.
-// It acquires a mutex to ensure thread safety during cleanup.
 func cleanupReplayCache() {
 	now := time.Now()
 	for token, expiry := range replayCache {
@@ -33,16 +27,8 @@ func cleanupReplayCache() {
 	}
 }
 
-// ClockSkewToleranceFuture defines the tolerance for future-based claims like 'exp'.
+// ClockSkewTolerance is configurable to adjust time-based validations.
-// Allows for more leniency with expiration checks.
+var ClockSkewTolerance = 2 * time.Minute
-var ClockSkewToleranceFuture = 2 * time.Minute
-
-// ClockSkewTolerancePast defines the tolerance for past-based claims like 'iat' and 'nbf'.
-// A smaller tolerance is typically used here to prevent accepting tokens issued too far in the future.
-var (
-	ClockSkewTolerancePast = 10 * time.Second
-	ClockSkewTolerance     = 2 * time.Minute
-)
 
 // JWT represents a JSON Web Token as defined in RFC 7519.
 type JWT struct {
@@ -52,18 +38,7 @@ type JWT struct {
 	Token     string
 }
 
-// parseJWT decodes a raw JWT string into its constituent parts: header, claims, and signature.
+// parseJWT parses a JWT token string into a JWT struct.
-// It splits the token string by '.', decodes each part using base64 URL decoding,
-// and unmarshals the header and claims JSON into maps. The raw signature bytes are stored.
-// It performs basic format validation (expecting 3 parts).
-// Note: This function does *not* validate the signature or the claims.
-//
-// Parameters:
-//   - tokenString: The raw JWT string.
-//
-// Returns:
-//   - A pointer to a JWT struct containing the decoded parts.
-//   - An error if the token format is invalid or decoding/unmarshaling fails.
 func parseJWT(tokenString string) (*JWT, error) {
 	parts := strings.Split(tokenString, ".")
 	if len(parts) != 3 {
@@ -99,24 +74,8 @@ func parseJWT(tokenString string) (*JWT, error) {
 	return jwt, nil
 }
 
-// Verify performs standard claim validation on the JWT according to RFC 7519.
+// Verify validates the standard JWT claims as defined in RFC 7519.
-// It checks the following:
+// Verify validates the standard JWT claims as defined in RFC 7519.
-// - Algorithm ('alg') is supported.
-// - Issuer ('iss') matches the expected issuerURL.
-// - Audience ('aud') contains the expected clientID.
-// - Expiration time ('exp') is in the future (within tolerance).
-// - Issued at time ('iat') is in the past (within tolerance).
-// - Not before time ('nbf'), if present, is in the past (within tolerance).
-// - Subject ('sub') claim exists and is not empty.
-// - JWT ID ('jti'), if present, is checked against a replay cache to prevent token reuse.
-//
-// Parameters:
-//   - issuerURL: The expected issuer URL (e.g., "https://accounts.google.com").
-//   - clientID: The expected audience value (the client ID of this application).
-//
-// Returns:
-//   - nil if all standard claims are valid.
-//   - An error describing the first validation failure encountered.
 func (j *JWT) Verify(issuerURL, clientID string) error {
 	// Validate algorithm to prevent algorithm switching attacks
 	alg, ok := j.Header["alg"].(string)
@@ -205,17 +164,6 @@ func (j *JWT) Verify(issuerURL, clientID string) error {
 
 	return nil
 }
-
-// verifyAudience checks if the expected audience is present in the token's 'aud' claim.
-// The 'aud' claim can be a single string or an array of strings.
-//
-// Parameters:
-//   - tokenAudience: The 'aud' claim value extracted from the token (can be string or []interface{}).
-//   - expectedAudience: The audience value expected for this application (client ID).
-//
-// Returns:
-//   - nil if the expected audience is found.
-//   - An error if the claim type is invalid or the expected audience is not present.
 func verifyAudience(tokenAudience interface{}, expectedAudience string) error {
 	switch aud := tokenAudience.(type) {
 	case string:
@@ -239,15 +187,6 @@ func verifyAudience(tokenAudience interface{}, expectedAudience string) error {
 	return nil
 }
 
-// verifyIssuer checks if the token's 'iss' claim matches the expected issuer URL.
-//
-// Parameters:
-//   - tokenIssuer: The 'iss' claim value from the token.
-//   - expectedIssuer: The expected issuer URL configured for the OIDC provider.
-//
-// Returns:
-//   - nil if the issuers match.
-//   - An error if the issuers do not match.
 func verifyIssuer(tokenIssuer, expectedIssuer string) error {
 	if tokenIssuer != expectedIssuer {
 		return fmt.Errorf("invalid issuer (token: %s, expected: %s)", tokenIssuer, expectedIssuer)
@@ -255,76 +194,54 @@ func verifyIssuer(tokenIssuer, expectedIssuer string) error {
 	return nil
 }
 
-// verifyTimeConstraint checks time-based claims ('exp', 'iat', 'nbf') against the current time,
+// verifyTimeConstraint is a generic function to verify time-based claims
-// allowing for configurable clock skew. It uses different tolerances for past and future checks.
-//
-// Parameters:
-//   - unixTime: The timestamp value from the claim (as a float64 Unix time).
-//   - claimName: The name of the claim being verified ("exp", "iat", "nbf").
-//   - future: A boolean indicating the direction of the check (true for 'exp', false for 'iat'/'nbf').
-//
-// Returns:
-//   - nil if the time constraint is met within the allowed tolerance.
-//   - An error describing the failure (e.g., "token has expired", "token used before issued").
 func verifyTimeConstraint(unixTime float64, claimName string, future bool) error {
 	claimTime := time.Unix(int64(unixTime), 0)
-	now := time.Now() // Use current time without truncation
+	now := time.Now().Truncate(time.Second)
 
-	var err error
+	// For expiration (future=true), we add skew to now (making now later)
-	if future { // 'exp' check
+	// For iat/nbf (future=false), we subtract skew from now (making now earlier)
-		// Token is expired if Now is after (ClaimTime + FutureTolerance)
+	skewDirection := 1
-		allowedExpiry := claimTime.Add(ClockSkewToleranceFuture)
+	if !future {
-		if now.After(allowedExpiry) {
+		skewDirection = -1
-			err = fmt.Errorf("token has expired (exp: %v, now: %v, allowed_until: %v)", claimTime.UTC(), now.UTC(), allowedExpiry.UTC())
+	}
-		}
+	skewedNow := now.Add(time.Duration(skewDirection) * ClockSkewTolerance)
-	} else { // 'iat' or 'nbf' check
+
-		// Token is invalid if Now is before (ClaimTime - PastTolerance)
+	if claimTime.Equal(now) {
-		allowedStart := claimTime.Add(-ClockSkewTolerancePast)
+		return nil
-		if now.Before(allowedStart) {
-			reason := "not yet valid"
-			if claimName == "iat" {
-				reason = "used before issued"
-			}
-			err = fmt.Errorf("token %s (%s: %v, now: %v, allowed_from: %v)", reason, claimName, claimTime.UTC(), now.UTC(), allowedStart.UTC())
-		}
 	}
 
-	return err
+	// For expiration: if skewedNow (later) is after expiration, token expired
+	// For iat/nbf: if skewedNow (earlier) is before claim time, token not yet valid
+	if (future && skewedNow.After(claimTime)) || (!future && skewedNow.Before(claimTime)) {
+		var reason string
+		if future {
+			reason = "has expired"
+		} else {
+			if claimName == "iat" {
+				reason = "used before issued"
+			} else {
+				reason = "not yet valid"
+			}
+		}
+		return fmt.Errorf("token %s (%s: %v, now: %v)", reason, claimName, claimTime.UTC(), now.UTC())
+	}
+
+	return nil
 }
 
-// verifyExpiration checks the 'exp' (Expiration Time) claim.
-// It calls verifyTimeConstraint with future=true.
 func verifyExpiration(expiration float64) error {
 	return verifyTimeConstraint(expiration, "exp", true)
 }
 
-// verifyIssuedAt checks the 'iat' (Issued At) claim.
-// It calls verifyTimeConstraint with future=false.
 func verifyIssuedAt(issuedAt float64) error {
 	return verifyTimeConstraint(issuedAt, "iat", false)
 }
 
-// verifyNotBefore checks the 'nbf' (Not Before) claim.
-// It calls verifyTimeConstraint with future=false.
 func verifyNotBefore(notBefore float64) error {
 	return verifyTimeConstraint(notBefore, "nbf", false)
 }
 
-// verifySignature validates the JWT's signature using the provided public key.
-// It parses the public key from PEM format, selects the appropriate hashing algorithm
-// based on the 'alg' parameter (SHA256/384/512), hashes the token's signing input
-// (header + "." + payload), and then verifies the signature against the hash using
-// the corresponding RSA (PKCS1v15 or PSS) or ECDSA verification method.
-//
-// Parameters:
-//   - tokenString: The raw, complete JWT string.
-//   - publicKeyPEM: The public key corresponding to the private key used for signing, in PEM format.
-//   - alg: The algorithm specified in the JWT header (e.g., "RS256", "ES384").
-//
-// Returns:
-//   - nil if the signature is valid.
-//   - An error if the token format is invalid, decoding fails, key parsing fails,
-//     the algorithm is unsupported, or the signature verification fails.
 func verifySignature(tokenString string, publicKeyPEM []byte, alg string) error {
 	parts := strings.Split(tokenString, ".")
 	if len(parts) != 3 {

metadata_cache.go

@@ -15,8 +15,6 @@ type MetadataCache struct {
 	stopCleanup         chan struct{}
 }
 
-// NewMetadataCache creates a new MetadataCache instance.
-// It initializes the cache structure and starts the background cleanup goroutine.
 func NewMetadataCache() *MetadataCache {
 	c := &MetadataCache{
 		autoCleanupInterval: 5 * time.Minute,
@@ -26,8 +24,7 @@ func NewMetadataCache() *MetadataCache {
 	return c
 }
 
-// Cleanup removes the cached provider metadata if it has expired.
+// Cleanup removes expired metadata from the cache.
-// This is called periodically by the auto-cleanup goroutine.
 func (c *MetadataCache) Cleanup() {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
@@ -37,32 +34,11 @@ func (c *MetadataCache) Cleanup() {
 		c.metadata = nil
 	}
 }
-
-// isCacheValid checks if the cached metadata is present and has not expired.
-// Note: This function assumes the read lock is held or it's called from a context
-// where the lock is already held (like within GetMetadata after locking).
 func (c *MetadataCache) isCacheValid() bool {
 	return c.metadata != nil && time.Now().Before(c.expiresAt)
 }
 
-// GetMetadata retrieves the OIDC provider metadata.
+// GetMetadata retrieves the metadata from cache or fetches it if expired
-// It first checks the cache for valid, non-expired metadata. If found, it's returned immediately.
-// If the cache is empty or expired, it attempts to fetch the metadata from the provider's
-// well-known endpoint using discoverProviderMetadata.
-// If fetching is successful, the new metadata is cached for 1 hour.
-// If fetching fails but valid metadata exists in the cache (even if expired), the cache expiry
-// is extended by 5 minutes, and the cached data is returned to prevent thundering herd issues.
-// If fetching fails and there's no cached data, an error is returned.
-// It employs double-checked locking for thread safety and performance.
-//
-// Parameters:
-//   - providerURL: The base URL of the OIDC provider.
-//   - httpClient: The HTTP client to use for fetching metadata.
-//   - logger: The logger instance for recording errors or warnings.
-//
-// Returns:
-//   - A pointer to the ProviderMetadata struct.
-//   - An error if metadata cannot be retrieved from cache or fetched from the provider.
 func (c *MetadataCache) GetMetadata(providerURL string, httpClient *http.Client, logger *Logger) (*ProviderMetadata, error) {
 	c.mutex.RLock()
 	if c.isCacheValid() {
@@ -91,21 +67,24 @@ func (c *MetadataCache) GetMetadata(providerURL string, httpClient *http.Client,
 	}
 
 	c.metadata = metadata
-	// Set a fixed cache lifetime (e.g., 1 hour)
+	// Calculate expiration time based on usage patterns
-	// TODO: Consider making this configurable or respecting HTTP cache headers
+	usageCount := 0 // This should be replaced with actual usage tracking logic
-	c.expiresAt = time.Now().Add(1 * time.Hour)
+	if usageCount < 10 {
+		c.expiresAt = time.Now().Add(30 * time.Minute)
+	} else if usageCount < 50 {
+		c.expiresAt = time.Now().Add(1 * time.Hour)
+	} else {
+		c.expiresAt = time.Now().Add(2 * time.Hour)
+	}
 
 	// End of GetMetadata
 	return metadata, nil
 }
 
-// startAutoCleanup starts the background goroutine that periodically calls Cleanup
-// to remove expired metadata from the cache.
 func (c *MetadataCache) startAutoCleanup() {
 	autoCleanupRoutine(c.autoCleanupInterval, c.stopCleanup, c.Cleanup)
 }
 
-// Close stops the automatic cleanup goroutine associated with this metadata cache.
 func (c *MetadataCache) Close() {
 	close(c.stopCleanup)
 }

session.go

@@ -16,15 +16,8 @@ import (
 	"github.com/gorilla/sessions"
 )
 
-// generateSecureRandomString creates a cryptographically secure, hex-encoded random string.
+// generateSecureRandomString creates a cryptographically secure random string of specified length.
-// It reads the specified number of bytes from crypto/rand and encodes them as a hexadecimal string.
+// It returns the generated string or an error if random generation fails.
-//
-// Parameters:
-//   - length: The number of random bytes to generate (the resulting hex string will be twice this length).
-//
-// Returns:
-//   - A hex-encoded random string.
-//   - An error if reading random bytes fails.
 func generateSecureRandomString(length int) (string, error) {
 	bytes := make([]byte, length)
 	if _, err := rand.Read(bytes); err != nil {
@@ -63,14 +56,7 @@ const (
 	minEncryptionKeyLength = 32
 )
 
-// compressToken compresses the input string using gzip and then encodes the result using standard base64 encoding.
+// compressToken compresses a token using gzip and base64 encodes it.
-// If any error occurs during compression, it returns the original uncompressed token as a fallback.
-//
-// Parameters:
-//   - token: The string to compress.
-//
-// Returns:
-//   - The base64 encoded, gzipped string, or the original string if compression fails.
 func compressToken(token string) string {
 	var b bytes.Buffer
 	gz := gzip.NewWriter(&b)
@@ -83,15 +69,7 @@ func compressToken(token string) string {
 	return base64.StdEncoding.EncodeToString(b.Bytes())
 }
 
-// decompressToken decodes a standard base64 encoded string and then decompresses the result using gzip.
+// decompressToken decompresses a base64 encoded gzipped token.
-// If base64 decoding or gzip decompression fails, it returns the original input string as a fallback,
-// assuming it might not have been compressed.
-//
-// Parameters:
-//   - compressed: The base64 encoded, gzipped string.
-//
-// Returns:
-//   - The decompressed original string, or the input string if decompression fails.
 func decompressToken(compressed string) string {
 	data, err := base64.StdEncoding.DecodeString(compressed)
 	if err != nil {
@@ -150,28 +128,25 @@ func NewSessionManager(encryptionKey string, forceHTTPS bool, logger *Logger) (*
 
 	// Initialize session pool.
 	sm.sessionPool.New = func() interface{} {
-		// Initialize SessionData with necessary fields and the mutex.
 		return &SessionData{
 			manager:            sm,
 			accessTokenChunks:  make(map[int]*sessions.Session),
 			refreshTokenChunks: make(map[int]*sessions.Session),
-			refreshMutex:       sync.Mutex{}, // Initialize the mutex
-			dirty:              false,        // Initialize dirty flag
 		}
 	}
 
 	return sm, nil
 }
 
-// getSessionOptions returns a sessions.Options struct configured with security best practices.
+// getSessionOptions returns secure session options configured for the current request.
-// It sets HttpOnly to true, Secure based on the request scheme or forceHTTPS setting,
-// SameSite to LaxMode, MaxAge to the absoluteSessionTimeout, and Path to "/".
-//
 // Parameters:
-//   - isSecure: A boolean indicating if the current request context is secure (HTTPS).
+//   - isSecure: Whether the current request is using HTTPS.
 //
-// Returns:
+// The options ensure cookies are:
-//   - A pointer to a configured sessions.Options struct.
+//   - HTTP-only (not accessible via JavaScript)
+//   - Secure when using HTTPS or when forceHTTPS is enabled
+//   - Using SameSite=Lax for CSRF protection
+//   - Set with appropriate timeout and path settings
 func (sm *SessionManager) getSessionOptions(isSecure bool) *sessions.Options {
 	return &sessions.Options{
 		HttpOnly: true,
@@ -190,38 +165,32 @@ func (sm *SessionManager) GetSession(r *http.Request) (*SessionData, error) {
 	// Get session from pool.
 	sessionData := sm.sessionPool.Get().(*SessionData)
 	sessionData.request = r
-	sessionData.dirty = false // Reset dirty flag when getting a session
-
-	// Function to properly handle errors and return the session to the pool
-	handleError := func(err error, message string) (*SessionData, error) {
-		if sessionData != nil {
-			sm.sessionPool.Put(sessionData)
-		}
-		return nil, fmt.Errorf("%s: %w", message, err)
-	}
 
 	var err error
 	sessionData.mainSession, err = sm.store.Get(r, mainCookieName)
 	if err != nil {
-		return handleError(err, "failed to get main session")
+		sm.sessionPool.Put(sessionData)
+		return nil, fmt.Errorf("failed to get main session: %w", err)
 	}
 
 	// Check for absolute session timeout.
 	if createdAt, ok := sessionData.mainSession.Values["created_at"].(int64); ok {
 		if time.Since(time.Unix(createdAt, 0)) > absoluteSessionTimeout {
 			sessionData.Clear(r, nil)
-			return handleError(fmt.Errorf("session timeout"), "session expired")
+			return nil, fmt.Errorf("session expired")
 		}
 	}
 
 	sessionData.accessSession, err = sm.store.Get(r, accessTokenCookie)
 	if err != nil {
-		return handleError(err, "failed to get access token session")
+		sm.sessionPool.Put(sessionData)
+		return nil, fmt.Errorf("failed to get access token session: %w", err)
 	}
 
 	sessionData.refreshSession, err = sm.store.Get(r, refreshTokenCookie)
 	if err != nil {
-		return handleError(err, "failed to get refresh token session")
+		sm.sessionPool.Put(sessionData)
+		return nil, fmt.Errorf("failed to get refresh token session: %w", err)
 	}
 
 	// Clear and reuse chunk maps.
@@ -239,14 +208,11 @@ func (sm *SessionManager) GetSession(r *http.Request) (*SessionData, error) {
 	return sessionData, nil
 }
 
-// getTokenChunkSessions retrieves all cookie chunks associated with a large token (access or refresh).
+// getTokenChunkSessions retrieves all session chunks for a given token type.
-// It iteratively attempts to load cookies named "{baseName}_0", "{baseName}_1", etc., until
-// a cookie is not found or returns an error. The loaded sessions are stored in the provided chunks map.
-//
 // Parameters:
-//   - r: The incoming HTTP request containing the cookies.
+//   - r: The HTTP request
-//   - baseName: The base name of the cookie (e.g., accessTokenCookie).
+//   - baseName: The base name for the token's session cookies
-//   - chunks: The map (typically SessionData.accessTokenChunks or SessionData.refreshTokenChunks) to populate with the found session chunks.
+//   - chunks: Map to store the chunks in
 func (sm *SessionManager) getTokenChunkSessions(r *http.Request, baseName string, chunks map[int]*sessions.Session) {
 	for i := 0; ; i++ {
 		sessionName := fmt.Sprintf("%s_%d", baseName, i)
@@ -285,36 +251,12 @@ type SessionData struct {
 	// refreshTokenChunks stores additional chunks of the refresh token
 	// when it exceeds the maximum cookie size.
 	refreshTokenChunks map[int]*sessions.Session
-
-	// refreshMutex protects refresh token operations within this session instance.
-	refreshMutex sync.Mutex
-
-	// dirty indicates whether the session data has changed and needs to be saved.
-	dirty bool
 }
 
-// IsDirty returns true if the session data has been modified since it was last loaded or saved.
+// Save persists all session data to cookies in the HTTP response.
-func (sd *SessionData) IsDirty() bool {
+// It saves the main session, token sessions, and any token chunks,
-	return sd.dirty
+// applying appropriate security options to each cookie. All cookies
-}
+// are saved with consistent security settings based on the request scheme.
-
-// MarkDirty explicitly sets the dirty flag to true.
-// This can be used when an operation doesn't change session data
-// but should still trigger a session save (e.g., to ensure the cookie is re-issued).
-func (sd *SessionData) MarkDirty() {
-	sd.dirty = true
-}
-
-// Save persists all parts of the session (main, access token, refresh token, and any chunks)
-// back to the client as cookies in the HTTP response. It applies secure cookie options
-// obtained via getSessionOptions based on the request's security context.
-//
-// Parameters:
-//   - r: The original HTTP request (used to determine security context for cookie options).
-//   - w: The HTTP response writer to which the Set-Cookie headers will be added.
-//
-// Returns:
-//   - An error if saving any of the session components fails.
 func (sd *SessionData) Save(r *http.Request, w http.ResponseWriter) error {
 	isSecure := strings.HasPrefix(r.URL.Scheme, "https") || sd.manager.forceHTTPS
 
@@ -324,125 +266,76 @@ func (sd *SessionData) Save(r *http.Request, w http.ResponseWriter) error {
 	sd.accessSession.Options = options
 	sd.refreshSession.Options = options
 
-	var firstErr error
+	// Save main session.
-	// Helper to record first error and log subsequent ones
+	if err := sd.mainSession.Save(r, w); err != nil {
-	saveOrLogError := func(s *sessions.Session, name string) {
+		return fmt.Errorf("failed to save main session: %w", err)
-		if s == nil { // Should not happen if initialized correctly
-			sd.manager.logger.Errorf("Attempted to save nil session: %s", name)
-			if firstErr == nil {
-				firstErr = fmt.Errorf("attempted to save nil session: %s", name)
-			}
-			return
-		}
-		if err := s.Save(r, w); err != nil {
-			errMsg := fmt.Errorf("failed to save %s session: %w", name, err)
-			sd.manager.logger.Error(errMsg.Error())
-			if firstErr == nil {
-				firstErr = errMsg
-			}
-		}
 	}
 
-	// Save main session.
-	saveOrLogError(sd.mainSession, "main")
-
 	// Save access token session.
-	saveOrLogError(sd.accessSession, "access token")
+	if err := sd.accessSession.Save(r, w); err != nil {
+		return fmt.Errorf("failed to save access token session: %w", err)
+	}
 
 	// Save refresh token session.
-	saveOrLogError(sd.refreshSession, "refresh token")
+	if err := sd.refreshSession.Save(r, w); err != nil {
+		return fmt.Errorf("failed to save refresh token session: %w", err)
+	}
 
 	// Save access token chunks.
-	for i, sessionChunk := range sd.accessTokenChunks {
+	for _, session := range sd.accessTokenChunks {
-		sessionChunk.Options = options
+		session.Options = options
-		saveOrLogError(sessionChunk, fmt.Sprintf("access token chunk %d", i))
+		if err := session.Save(r, w); err != nil {
+			return fmt.Errorf("failed to save access token chunk session: %w", err)
+		}
 	}
 
 	// Save refresh token chunks.
-	for i, sessionChunk := range sd.refreshTokenChunks {
+	for _, session := range sd.refreshTokenChunks {
-		sessionChunk.Options = options
+		session.Options = options
-		saveOrLogError(sessionChunk, fmt.Sprintf("refresh token chunk %d", i))
+		if err := session.Save(r, w); err != nil {
+			return fmt.Errorf("failed to save refresh token chunk session: %w", err)
+		}
 	}
 
-	if firstErr == nil {
+	return nil
-		sd.dirty = false // Reset dirty flag only if all saves were successful
-	}
-	return firstErr
 }
 
-// Clear removes all session data associated with this SessionData instance.
+// Clear removes all session data by expiring all cookies and clearing their values.
-// It clears the values map of the main, access, and refresh sessions, sets their MaxAge to -1
-// to expire the cookies immediately, and clears any associated token chunk cookies.
-// If a ResponseWriter is provided, it attempts to save the expired sessions to send the
-// expiring Set-Cookie headers. Finally, it clears internal fields and returns the SessionData
-// object to the pool.
-//
-// Parameters:
-//   - r: The HTTP request (required by the underlying session store).
-//   - w: The HTTP response writer (optional). If provided, expiring Set-Cookie headers will be sent.
-//
-// Returns:
-//   - An error if saving the expired sessions fails (only if w is not nil).
-//
-// Note: This method will always return the SessionData object to the pool, even if an error occurs.
 func (sd *SessionData) Clear(r *http.Request, w http.ResponseWriter) error {
-	sd.dirty = true // Clearing the session means its state is changing and needs to be saved.
-
 	// Clear and expire all sessions.
-	if sd.mainSession != nil {
+	sd.mainSession.Options.MaxAge = -1
-		sd.mainSession.Options.MaxAge = -1
+	sd.accessSession.Options.MaxAge = -1
-		for k := range sd.mainSession.Values {
+	sd.refreshSession.Options.MaxAge = -1
-			delete(sd.mainSession.Values, k)
+
-		}
+	for k := range sd.mainSession.Values {
+		delete(sd.mainSession.Values, k)
 	}
-	if sd.accessSession != nil {
+	for k := range sd.accessSession.Values {
-		sd.accessSession.Options.MaxAge = -1
+		delete(sd.accessSession.Values, k)
-		for k := range sd.accessSession.Values {
-			delete(sd.accessSession.Values, k)
-		}
 	}
-	if sd.refreshSession != nil {
+	for k := range sd.refreshSession.Values {
-		sd.refreshSession.Options.MaxAge = -1
+		delete(sd.refreshSession.Values, k)
-		for k := range sd.refreshSession.Values {
-			delete(sd.refreshSession.Values, k)
-		}
 	}
 
 	// Clear chunk sessions.
 	sd.clearTokenChunks(r, sd.accessTokenChunks)
 	sd.clearTokenChunks(r, sd.refreshTokenChunks)
 
-	// Create a guaranteed error when the response writer is set
-	// This is primarily for testing - in production w will often be nil
 	var err error
 	if w != nil {
-		// Intentionally create a test error in session
-		if r != nil && r.Header.Get("X-Test-Error") == "true" {
-			sd.mainSession.Values["error_trigger"] = func() {} // Will cause marshaling to fail
-		}
-
-		// Try to save the expired sessions
 		err = sd.Save(r, w)
 	}
 
 	// Clear transient per-request fields.
 	sd.request = nil
 
-	// Return session to pool, regardless of error.
+	// Return session to pool.
-	// This ensures the session is always returned to the pool,
-	// preventing memory leaks.
 	sd.manager.sessionPool.Put(sd)
 
-	// Return the error from Save, if any
 	return err
 }
 
-// clearTokenChunks iterates through a map of session chunks, clears their values,
+// clearTokenChunks removes all session chunks for a given token type.
-// and sets their MaxAge to -1 to expire them. This is used internally by Clear.
-//
-// Parameters:
-//   - r: The HTTP request (required by the underlying session store, though not directly used here).
-//   - chunks: The map of session chunks (e.g., sd.accessTokenChunks) to clear and expire.
 func (sd *SessionData) clearTokenChunks(r *http.Request, chunks map[int]*sessions.Session) {
 	for _, session := range chunks {
 		session.Options.MaxAge = -1
@@ -452,12 +345,7 @@ func (sd *SessionData) clearTokenChunks(r *http.Request, chunks map[int]*session
 	}
 }
 
-// GetAuthenticated checks if the session is marked as authenticated and has not exceeded
+// GetAuthenticated returns whether the current session is authenticated.
-// the absolute session timeout.
-//
-// Returns:
-//   - true if the "authenticated" flag is set to true and the session creation time is within the allowed timeout.
-//   - false otherwise.
 func (sd *SessionData) GetAuthenticated() bool {
 	auth, _ := sd.mainSession.Values["authenticated"].(bool)
 	if !auth {
@@ -472,74 +360,22 @@ func (sd *SessionData) GetAuthenticated() bool {
 	return time.Since(time.Unix(createdAt, 0)) <= absoluteSessionTimeout
 }
 
-// SetAuthenticated sets the authentication status of the session.
+// SetAuthenticated updates the session's authentication status and rotates session ID.
-// If setting to true, it generates a new secure session ID for the main session
+// Returns an error if generating a new session ID fails.
-// to prevent session fixation attacks and records the current time as the creation time.
-//
-// Parameters:
-//   - value: The boolean authentication status (true for authenticated, false otherwise).
-//
-// Returns:
-//   - An error if generating a new session ID fails when setting value to true.
 func (sd *SessionData) SetAuthenticated(value bool) error {
-	currentAuth := sd.GetAuthenticated() // This checks flag and expiry
-	changed := false
-
-	if currentAuth != value {
-		changed = true
-	}
-
 	if value {
-		// If we are setting to true, and either it wasn't true before,
-		// or if the session ID needs regeneration (e.g. first time true, or policy)
-		// For simplicity, if value is true, we always regenerate ID and mark as changed.
-		// This ensures session ID regeneration is always saved.
 		id, err := generateSecureRandomString(32)
 		if err != nil {
 			return fmt.Errorf("failed to generate secure session id: %w", err)
 		}
-		if sd.mainSession.ID != id { // ID actually changed
-			changed = true
-		}
 		sd.mainSession.ID = id
-		newCreationTime := time.Now().Unix()
+		sd.mainSession.Values["created_at"] = time.Now().Unix()
-		if oldTime, ok := sd.mainSession.Values["created_at"].(int64); !ok || oldTime != newCreationTime {
-			changed = true
-		}
-		sd.mainSession.Values["created_at"] = newCreationTime
-		if oldAuth, ok := sd.mainSession.Values["authenticated"].(bool); !ok || oldAuth != value {
-			changed = true
-		}
-	} else { // value is false
-		if oldAuth, ok := sd.mainSession.Values["authenticated"].(bool); !ok || oldAuth != value {
-			changed = true
-		}
 	}
-
 	sd.mainSession.Values["authenticated"] = value
-	if changed {
-		sd.dirty = true
-	}
 	return nil
 }
 
-// ReturnToPool explicitly returns this SessionData object to the pool.
+// GetAccessToken retrieves the complete access token from the session.
-// This should be called when you're done with a SessionData in any error path
-// where Clear() is not called, to prevent memory leaks.
-func (sd *SessionData) ReturnToPool() {
-	if sd != nil && sd.manager != nil {
-		// Clear request reference to avoid memory leaks
-		sd.request = nil
-		sd.manager.sessionPool.Put(sd)
-	}
-}
-
-// GetAccessToken retrieves the access token stored in the session.
-// It handles reassembling the token from multiple cookie chunks if necessary
-// and decompresses it if it was stored compressed.
-//
-// Returns:
-//   - The complete, decompressed access token string, or an empty string if not found.
 func (sd *SessionData) GetAccessToken() string {
 	token, _ := sd.accessSession.Values["token"].(string)
 	if token != "" {
@@ -573,23 +409,8 @@ func (sd *SessionData) GetAccessToken() string {
 	return token
 }
 
-// SetAccessToken stores the provided access token in the session.
+// SetAccessToken stores the access token in the session.
-// It first expires any existing access token chunk cookies.
-// It then compresses the token. If the compressed token fits within a single cookie (maxCookieSize),
-// it's stored directly in the primary access token session. Otherwise, the compressed token
-// is split into chunks, and each chunk is stored in a separate numbered cookie (_oidc_raczylo_a_0, _oidc_raczylo_a_1, etc.).
-//
-// Parameters:
-//   - token: The access token string to store.
 func (sd *SessionData) SetAccessToken(token string) {
-	currentAccessToken := sd.GetAccessToken()
-	if currentAccessToken == token {
-		// If token is empty, and current is also empty, it's not a change.
-		// This check handles both empty and non-empty identical cases.
-		return
-	}
-	sd.dirty = true
-
 	// Expire any existing chunk cookies first.
 	if sd.request != nil {
 		sd.expireAccessTokenChunks(nil) // Will be saved when Save() is called.
@@ -598,13 +419,6 @@ func (sd *SessionData) SetAccessToken(token string) {
 	// Clear and prepare chunks map for new token.
 	sd.accessTokenChunks = make(map[int]*sessions.Session)
 
-	if token == "" { // Clearing the token
-		sd.accessSession.Values["token"] = ""
-		sd.accessSession.Values["compressed"] = false
-		// sd.accessTokenChunks is already cleared
-		return
-	}
-
 	// Compress token.
 	compressed := compressToken(token)
 
@@ -613,30 +427,19 @@ func (sd *SessionData) SetAccessToken(token string) {
 		sd.accessSession.Values["compressed"] = true
 	} else {
 		// Split compressed token into chunks.
-		sd.accessSession.Values["token"] = ""        // Main cookie won't hold the token directly
+		sd.accessSession.Values["token"] = ""
-		sd.accessSession.Values["compressed"] = true // Data in chunks is compressed
+		sd.accessSession.Values["compressed"] = true
 		chunks := splitIntoChunks(compressed, maxCookieSize)
-		for i, chunkData := range chunks {
+		for i, chunk := range chunks {
 			sessionName := fmt.Sprintf("%s_%d", accessTokenCookie, i)
-			// Ensure sd.request is available, otherwise log warning or handle error
-			if sd.request == nil {
-				sd.manager.logger.Infof("SetAccessToken: sd.request is nil, cannot get/create chunk session %s", sessionName)
-				// Potentially skip this chunk or error out, depending on desired robustness
-				continue
-			}
 			session, _ := sd.manager.store.Get(sd.request, sessionName)
-			session.Values["token_chunk"] = chunkData
+			session.Values["token_chunk"] = chunk
 			sd.accessTokenChunks[i] = session
 		}
 	}
 }
 
-// GetRefreshToken retrieves the refresh token stored in the session.
+// GetRefreshToken retrieves the complete refresh token from the session.
-// It handles reassembling the token from multiple cookie chunks if necessary
-// and decompresses it if it was stored compressed.
-//
-// Returns:
-//   - The complete, decompressed refresh token string, or an empty string if not found.
 func (sd *SessionData) GetRefreshToken() string {
 	token, _ := sd.refreshSession.Values["token"].(string)
 	if token != "" {
@@ -670,21 +473,8 @@ func (sd *SessionData) GetRefreshToken() string {
 	return token
 }
 
-// SetRefreshToken stores the provided refresh token in the session.
+// SetRefreshToken stores the refresh token in the session.
-// It first expires any existing refresh token chunk cookies.
-// It then compresses the token. If the compressed token fits within a single cookie (maxCookieSize),
-// it's stored directly in the primary refresh token session. Otherwise, the compressed token
-// is split into chunks, and each chunk is stored in a separate numbered cookie (_oidc_raczylo_r_0, _oidc_raczylo_r_1, etc.).
-//
-// Parameters:
-//   - token: The refresh token string to store.
 func (sd *SessionData) SetRefreshToken(token string) {
-	currentRefreshToken := sd.GetRefreshToken()
-	if currentRefreshToken == token {
-		return
-	}
-	sd.dirty = true
-
 	// Expire any existing chunk cookies first.
 	if sd.request != nil {
 		sd.expireRefreshTokenChunks(nil) // Will be saved when Save() is called.
@@ -693,13 +483,6 @@ func (sd *SessionData) SetRefreshToken(token string) {
 	// Clear and prepare chunks map for new token.
 	sd.refreshTokenChunks = make(map[int]*sessions.Session)
 
-	if token == "" { // Clearing the token
-		sd.refreshSession.Values["token"] = ""
-		sd.refreshSession.Values["compressed"] = false
-		// sd.refreshTokenChunks is already cleared
-		return
-	}
-
 	// Compress token.
 	compressed := compressToken(token)
 
@@ -708,29 +491,19 @@ func (sd *SessionData) SetRefreshToken(token string) {
 		sd.refreshSession.Values["compressed"] = true
 	} else {
 		// Split compressed token into chunks.
-		sd.refreshSession.Values["token"] = ""        // Main cookie won't hold the token directly
+		sd.refreshSession.Values["token"] = ""
-		sd.refreshSession.Values["compressed"] = true // Data in chunks is compressed
+		sd.refreshSession.Values["compressed"] = true
 		chunks := splitIntoChunks(compressed, maxCookieSize)
-		for i, chunkData := range chunks {
+		for i, chunk := range chunks {
 			sessionName := fmt.Sprintf("%s_%d", refreshTokenCookie, i)
-			if sd.request == nil {
-				sd.manager.logger.Infof("SetRefreshToken: sd.request is nil, cannot get/create chunk session %s", sessionName)
-				continue
-			}
 			session, _ := sd.manager.store.Get(sd.request, sessionName)
-			session.Values["token_chunk"] = chunkData
+			session.Values["token_chunk"] = chunk
 			sd.refreshTokenChunks[i] = session
 		}
 	}
 }
 
-// expireAccessTokenChunks finds all existing access token chunk cookies (_oidc_raczylo_a_N)
+// expireAccessTokenChunks expires any existing access token chunk cookies.
-// associated with the current request, clears their values, and sets their MaxAge to -1.
-// If a ResponseWriter is provided, it attempts to save the expired chunk sessions to send
-// the expiring Set-Cookie headers. This is used internally when setting a new access token.
-//
-// Parameters:
-//   - w: The HTTP response writer (optional). If provided, expiring Set-Cookie headers will be sent.
 func (sd *SessionData) expireAccessTokenChunks(w http.ResponseWriter) {
 	for i := 0; ; i++ {
 		sessionName := fmt.Sprintf("%s_%d", accessTokenCookie, i)
@@ -748,13 +521,7 @@ func (sd *SessionData) expireAccessTokenChunks(w http.ResponseWriter) {
 	}
 }
 
-// expireRefreshTokenChunks finds all existing refresh token chunk cookies (_oidc_raczylo_r_N)
+// expireRefreshTokenChunks expires any existing refresh token chunk cookies.
-// associated with the current request, clears their values, and sets their MaxAge to -1.
-// If a ResponseWriter is provided, it attempts to save the expired chunk sessions to send
-// the expiring Set-Cookie headers. This is used internally when setting a new refresh token.
-//
-// Parameters:
-//   - w: The HTTP response writer (optional). If provided, expiring Set-Cookie headers will be sent.
 func (sd *SessionData) expireRefreshTokenChunks(w http.ResponseWriter) {
 	for i := 0; ; i++ {
 		sessionName := fmt.Sprintf("%s_%d", refreshTokenCookie, i)
@@ -772,15 +539,7 @@ func (sd *SessionData) expireRefreshTokenChunks(w http.ResponseWriter) {
 	}
 }
 
-// splitIntoChunks divides a string `s` into a slice of strings, where each element
+// splitIntoChunks splits a string into chunks of specified size.
-// has a maximum length of `chunkSize`.
-//
-// Parameters:
-//   - s: The string to split.
-//   - chunkSize: The maximum size of each chunk.
-//
-// Returns:
-//   - A slice of strings representing the chunks.
 func splitIntoChunks(s string, chunkSize int) []string {
 	var chunks []string
 	for len(s) > 0 {
@@ -795,161 +554,57 @@ func splitIntoChunks(s string, chunkSize int) []string {
 	return chunks
 }
 
-// GetCSRF retrieves the Cross-Site Request Forgery (CSRF) token stored in the main session.
+// GetCSRF retrieves the CSRF token from the session.
-//
-// Returns:
-//   - The CSRF token string, or an empty string if not set.
 func (sd *SessionData) GetCSRF() string {
 	csrf, _ := sd.mainSession.Values["csrf"].(string)
 	return csrf
 }
 
-// SetCSRF stores the provided CSRF token string in the main session.
+// SetCSRF stores a new CSRF token in the session.
-// This token is typically generated at the start of the authentication flow.
-//
-// Parameters:
-//   - token: The CSRF token to store.
 func (sd *SessionData) SetCSRF(token string) {
-	currentVal, _ := sd.mainSession.Values["csrf"].(string)
+	sd.mainSession.Values["csrf"] = token
-	if currentVal != token {
-		sd.mainSession.Values["csrf"] = token
-		sd.dirty = true
-	}
 }
 
-// GetNonce retrieves the OIDC nonce value stored in the main session.
+// GetNonce retrieves the nonce value from the session.
-// The nonce is used to associate an ID token with the specific authentication request.
-//
-// Returns:
-//   - The nonce string, or an empty string if not set.
 func (sd *SessionData) GetNonce() string {
 	nonce, _ := sd.mainSession.Values["nonce"].(string)
 	return nonce
 }
 
-// SetNonce stores the provided OIDC nonce string in the main session.
+// SetNonce stores a new nonce value in the session.
-// This nonce is typically generated at the start of the authentication flow.
-//
-// Parameters:
-//   - nonce: The nonce string to store.
 func (sd *SessionData) SetNonce(nonce string) {
-	currentVal, _ := sd.mainSession.Values["nonce"].(string)
+	sd.mainSession.Values["nonce"] = nonce
-	if currentVal != nonce {
-		sd.mainSession.Values["nonce"] = nonce
-		sd.dirty = true
-	}
 }
 
-// GetCodeVerifier retrieves the PKCE (Proof Key for Code Exchange) code verifier
+// GetCodeVerifier retrieves the PKCE code verifier from the session.
-// stored in the main session. This is only relevant if PKCE is enabled.
-//
-// Returns:
-//   - The code verifier string, or an empty string if not set or PKCE is disabled.
 func (sd *SessionData) GetCodeVerifier() string {
 	codeVerifier, _ := sd.mainSession.Values["code_verifier"].(string)
 	return codeVerifier
 }
 
-// SetCodeVerifier stores the provided PKCE code verifier string in the main session.
+// SetCodeVerifier stores the PKCE code verifier in the session.
-// This is typically called at the start of the authentication flow if PKCE is enabled.
-//
-// Parameters:
-//   - codeVerifier: The PKCE code verifier string to store.
 func (sd *SessionData) SetCodeVerifier(codeVerifier string) {
-	currentVal, _ := sd.mainSession.Values["code_verifier"].(string)
+	sd.mainSession.Values["code_verifier"] = codeVerifier
-	if currentVal != codeVerifier {
-		sd.mainSession.Values["code_verifier"] = codeVerifier
-		sd.dirty = true
-	}
 }
 
-// GetEmail retrieves the authenticated user's email address stored in the main session.
+// GetEmail retrieves the authenticated user's email address from the session.
-// This is typically extracted from the ID token claims after successful authentication.
-//
-// Returns:
-//   - The user's email address string, or an empty string if not set.
 func (sd *SessionData) GetEmail() string {
 	email, _ := sd.mainSession.Values["email"].(string)
 	return email
 }
 
-// SetEmail stores the provided user email address string in the main session.
+// SetEmail stores the user's email address in the session.
-// This is typically called after successful authentication and claim extraction.
-//
-// Parameters:
-//   - email: The user's email address to store.
 func (sd *SessionData) SetEmail(email string) {
-	currentVal, _ := sd.mainSession.Values["email"].(string)
+	sd.mainSession.Values["email"] = email
-	if currentVal != email {
-		sd.mainSession.Values["email"] = email
-		sd.dirty = true
-	}
 }
 
-// GetIncomingPath retrieves the original request URI (including query parameters)
+// GetIncomingPath retrieves the original request path that triggered the authentication flow.
-// that the user was trying to access before being redirected for authentication.
-// This is stored in the main session to allow redirection back after successful login.
-//
-// Returns:
-//   - The original request URI string, or an empty string if not set.
 func (sd *SessionData) GetIncomingPath() string {
 	path, _ := sd.mainSession.Values["incoming_path"].(string)
 	return path
 }
 
-// SetIncomingPath stores the original request URI (path and query parameters)
+// SetIncomingPath stores the original request path that triggered the authentication flow.
-// in the main session. This is typically called at the start of the authentication flow.
-//
-// Parameters:
-//   - path: The original request URI string (e.g., "/protected/resource?id=123").
 func (sd *SessionData) SetIncomingPath(path string) {
-	currentVal, _ := sd.mainSession.Values["incoming_path"].(string)
+	sd.mainSession.Values["incoming_path"] = path
-	if currentVal != path {
-		sd.mainSession.Values["incoming_path"] = path
-		sd.dirty = true
-	}
-}
-
-// GetIDToken retrieves the ID token stored in the session.
-// It handles reassembling the token from multiple cookie chunks if necessary
-// and decompresses it if it was stored compressed.
-//
-// Returns:
-//   - The complete, decompressed ID token string, or an empty string if not found.
-func (sd *SessionData) GetIDToken() string {
-	token, _ := sd.mainSession.Values["id_token"].(string)
-	if token != "" {
-		compressed, _ := sd.mainSession.Values["id_token_compressed"].(bool)
-		if compressed {
-			return decompressToken(token)
-		}
-		return token
-	}
-	return ""
-}
-
-// SetIDToken stores the provided ID token in the session.
-//
-// Parameters:
-//   - token: The ID token string to store.
-func (sd *SessionData) SetIDToken(token string) {
-	currentIDToken := sd.GetIDToken() // Gets fully reassembled, decompressed token
-	if currentIDToken == token {
-		// This handles cases where token is "" and currentIDToken is also "", no change.
-		// Or token is "abc" and currentIDToken is "abc", no change.
-		return
-	}
-
-	sd.dirty = true // Mark as dirty because a change is being made
-
-	if token == "" {
-		sd.mainSession.Values["id_token"] = ""
-		sd.mainSession.Values["id_token_compressed"] = false
-		return
-	}
-
-	// Compress token
-	compressed := compressToken(token)
-	sd.mainSession.Values["id_token"] = compressed
-	sd.mainSession.Values["id_token_compressed"] = true
 }

session_test.go

@@ -1,221 +1,382 @@
 package traefikoidc
 
 import (
-	"net/http"
+	"math/rand"
 	"net/http/httptest"
-	"runtime"
 	"strings"
 	"testing"
-	"time"
 )
 
-func TestSessionPoolMemoryLeak(t *testing.T) {
+// generateRandomString creates a random string of specified length
-	logger := NewLogger("debug")
+func generateRandomString(length int) string {
-	sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
+	const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
-	if err != nil {
+	b := make([]byte, length)
-		t.Fatalf("Failed to create session manager: %v", err)
+	for i := range b {
+		b[i] = charset[rand.Intn(len(charset))]
+	}
+	return string(b)
+}
+
+// TestTokenCompression tests the token compression functionality
+func TestTokenCompression(t *testing.T) {
+	tests := []struct {
+		name     string
+		token    string
+		wantSize int // Expected size after compression (approximate)
+	}{
+		{
+			name:     "Short token",
+			token:    "shorttoken",
+			wantSize: 50, // Base64 encoded gzip has overhead for small content
+		},
+		{
+			name:     "Repeating content",
+			token:    strings.Repeat("abcdef", 1000),
+			wantSize: 100, // Should compress well due to repetition
+		},
+		{
+			name:     "Random content",
+			token:    generateRandomString(1000),
+			wantSize: 2000, // Random content won't compress much
+		},
 	}
 
-	// Create a fake request
+	for _, tt := range tests {
-	req := httptest.NewRequest("GET", "http://example.com/foo", nil)
+		t.Run(tt.name, func(t *testing.T) {
+			compressed := compressToken(tt.token)
+			decompressed := decompressToken(compressed)
 
-	// Test 1: Successful session creation and return
+			// Only verify compression ratio for non-short tokens
+			if len(tt.token) > 100 {
+				compressionRatio := float64(len(compressed)) / float64(len(tt.token))
+				t.Logf("Compression ratio for %s: %.2f", tt.name, compressionRatio)
+
+				if compressionRatio > 1.1 { // Allow up to 10% size increase
+					t.Errorf("Compression increased size too much: original=%d, compressed=%d, ratio=%.2f",
+						len(tt.token), len(compressed), compressionRatio)
+				}
+			}
+
+			// Verify decompression restores original
+			if decompressed != tt.token {
+				t.Error("Decompression failed to restore original token")
+			}
+
+			// Verify approximate compression ratio
+			if len(compressed) > tt.wantSize*2 {
+				t.Errorf("Compression ratio worse than expected: got=%d, want<%d", len(compressed), tt.wantSize*2)
+			}
+		})
+	}
+}
+
+// TestSessionManager tests the SessionManager functionality
+
+func TestCookiePrefix(t *testing.T) {
+	// Create a session and verify cookie names
+	req := httptest.NewRequest("GET", "/test", nil)
+	rr := httptest.NewRecorder()
+
+	sm, _ := NewSessionManager("0123456789abcdef0123456789abcdef", true, NewLogger("debug"))
 	session, err := sm.GetSession(req)
 	if err != nil {
-		t.Fatalf("GetSession failed: %v", err)
+		t.Fatalf("Failed to get session: %v", err)
 	}
 
-	// Clear the session which should return it to the pool
+	// Set some data to ensure cookies are created
-	session.Clear(req, nil)
+	session.SetAuthenticated(true)
 
-	// Test 2: ReturnToPool explicit method
+	// Expire any existing cookies
-	session, err = sm.GetSession(req)
+	session.expireAccessTokenChunks(rr)
-	if err != nil {
+	session.expireRefreshTokenChunks(rr)
-		t.Fatalf("GetSession failed: %v", err)
+
+	// Set new tokens
+	session.SetAccessToken("test_token")
+	session.SetRefreshToken("test_refresh_token")
+
+	if err := session.Save(req, rr); err != nil {
+		t.Fatalf("Failed to save session: %v", err)
 	}
 
-	// Call ReturnToPool directly
+	// Check cookie prefixes
-	session.ReturnToPool()
+	cookies := rr.Result().Cookies()
-
+	for _, cookie := range cookies {
-	// Test 3: Error path in GetSession
+		if !strings.HasPrefix(cookie.Name, "_oidc_raczylo_") {
-	// Modify the session store to force an error - use a different encryption key
+			t.Errorf("Cookie %s does not have expected prefix '_oidc_raczylo_'", cookie.Name)
-	badSM, _ := NewSessionManager("different0123456789abcdef0123456789abcdef0123456789", false, logger)
+		}
-
-	// Get session using mismatched manager/request to force error
-	_, err = badSM.GetSession(req)
-	if err == nil {
-		// We don't test the exact error since it could vary, just that we get one
-		t.Log("Note: Expected error when using mismatched encryption keys")
-	}
-
-	// Force GC to ensure any objects are cleaned up
-	runtime.GC()
-
-	// Wait a moment for GC to complete
-	time.Sleep(100 * time.Millisecond)
-
-	// Check if we have objects in the pool
-	// This is just a simple check; in a real scenario, we'd have to
-	// consider that sync.Pool can discard objects at any time.
-	pooledCount := getPooledObjects(sm)
-	t.Logf("Pooled objects count: %d", pooledCount)
-}
-
-func TestSessionErrorHandling(t *testing.T) {
-	logger := NewLogger("debug")
-	sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
-	if err != nil {
-		t.Fatalf("Failed to create session manager: %v", err)
-	}
-
-	// Create a fake request
-	req := httptest.NewRequest("GET", "http://example.com/foo", nil)
-
-	// Call the GetSession method, corrupting the cookie to force an error
-	req.AddCookie(&http.Cookie{
-		Name:  mainCookieName,
-		Value: "corrupt-value",
-	})
-
-	_, err = sm.GetSession(req)
-	if err == nil {
-		t.Fatal("Expected error, got nil")
-	}
-
-	// Check that the error message contains our expected prefix
-	if err != nil && !strings.Contains(err.Error(), "failed to get main session:") {
-		t.Fatalf("Unexpected error message: %v", err)
 	}
 }
 
-func TestSessionClearAlwaysReturnsToPool(t *testing.T) {
+func TestTokenRefreshCleanup(t *testing.T) {
-	logger := NewLogger("debug")
+	req := httptest.NewRequest("GET", "/test", nil)
-	sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
+	rr := httptest.NewRecorder()
-	if err != nil {
-		t.Fatalf("Failed to create session manager: %v", err)
-	}
 
-	// Create a test request with the special header that will trigger an error
+	sm, _ := NewSessionManager("0123456789abcdef0123456789abcdef", true, NewLogger("debug"))
-	req := httptest.NewRequest("GET", "http://example.com/foo", nil)
-	req.Header.Set("X-Test-Error", "true") // This will trigger the error in session.Clear
-
-	// Get a session
 	session, err := sm.GetSession(req)
 	if err != nil {
-		t.Fatalf("GetSession failed: %v", err)
+		t.Fatalf("Failed to get session: %v", err)
 	}
 
-	// Create a response writer
+	// Set a large token that will be split into chunks
-	w := httptest.NewRecorder()
+	largeToken := strings.Repeat("x", 5000)
+	session.SetAccessToken(largeToken)
 
-	// Call Clear with the test request (with X-Test-Error header) and response writer
+	if err := session.Save(req, rr); err != nil {
-	// This should trigger the serialization error in Save
+		t.Fatalf("Failed to save session: %v", err)
-	clearErr := session.Clear(req, w)
-
-	// Verify that Clear returned the error from Save
-	if clearErr == nil {
-		t.Error("Expected an error from Clear with X-Test-Error header, but got nil")
-	} else {
-		t.Logf("Received expected error from Clear: %v", clearErr)
 	}
 
-	// Force GC to ensure any objects are cleaned up
+	// Get initial cookies
-	runtime.GC()
+	initialCookies := rr.Result().Cookies()
-	time.Sleep(100 * time.Millisecond)
 
-	// Create and clear another session (without the error header) to verify the pool is still working
+	// Create a new request with the initial cookies
-	normalReq := httptest.NewRequest("GET", "http://example.com/foo", nil)
+	newReq := httptest.NewRequest("GET", "/test", nil)
-	session2, err := sm.GetSession(normalReq)
+	for _, cookie := range initialCookies {
+		newReq.AddCookie(cookie)
+	}
+	newRr := httptest.NewRecorder()
+
+	// Get session with cookies and set a new token
+	newSession, err := sm.GetSession(newReq)
 	if err != nil {
-		t.Fatalf("Second GetSession failed: %v", err)
+		t.Fatalf("Failed to get new session: %v", err)
 	}
-	session2.Clear(normalReq, nil)
 
-	// If we got here without panics, the test is successful
+	// Create a response recorder for expired cookies
-	t.Log("Session returned to pool despite errors")
+	expiredRr := httptest.NewRecorder()
+
+	// Expire old chunk cookies
+	newSession.expireAccessTokenChunks(expiredRr)
+
+	// Set a smaller token that won't need chunks
+	newSession.SetAccessToken("small_token")
+
+	// Save session with new token
+	if err := newSession.Save(newReq, newRr); err != nil {
+		t.Fatalf("Failed to save new session: %v", err)
+	}
+
+	// Check cookies in response where old cookies are expired
+	intermediateResponse := expiredRr.Result()
+	intermediateCount := 0
+	chunkCount := 0
+	expiredCount := 0
+
+	for _, cookie := range intermediateResponse.Cookies() {
+		if strings.Contains(cookie.Name, "_oidc_raczylo_a_") && strings.Count(cookie.Name, "_") > 3 {
+			chunkCount++
+			if cookie.MaxAge < 0 {
+				expiredCount++
+				t.Logf("Found expired chunk cookie: %s (MaxAge=%d)", cookie.Name, cookie.MaxAge)
+			}
+		} else if cookie.MaxAge >= 0 {
+			intermediateCount++
+			t.Logf("Found active cookie: %s (MaxAge=%d)", cookie.Name, cookie.MaxAge)
+		}
+	}
+
+	// All chunk cookies should be expired
+	if chunkCount > 0 && chunkCount != expiredCount {
+		t.Errorf("Not all chunk cookies are expired: %d chunks, %d expired", chunkCount, expiredCount)
+	}
+
+	// Should have fewer active cookies after setting smaller token
+	if intermediateCount >= len(initialCookies) {
+		t.Errorf("Expected fewer active cookies after token refresh, got %d, want less than %d", intermediateCount, len(initialCookies))
+	}
 }
 
-// This placeholder comment is intentionally left empty since we're removing redundant code
+func TestSessionManager(t *testing.T) {
+	ts := &TestSuite{t: t}
+	ts.Setup()
 
-// Helper function to count objects in the session pool for a given manager
+	tests := []struct {
-func getPooledObjects(sm *SessionManager) int {
+		name                string
-	// Collect objects until we can't get any more from the pool
+		authenticated       bool
-	// Set a max limit to avoid potential infinite loops
+		email               string
-	var objects []*SessionData
+		accessToken         string
-	maxAttempts := 100 // Safety limit to prevent infinite loops
+		refreshToken        string
-
+		expectedCookieCount int
-	for i := 0; i < maxAttempts; i++ {
+		wantCompressed      bool // Whether tokens should be compressed
-		obj := sm.sessionPool.Get()
+	}{
-		if obj == nil {
+		{
-			break
+			name:                "Short tokens",
-		}
+			authenticated:       true,
-
+			email:               "[email protected]",
-		// Type assertion with validation
+			accessToken:         "shortaccesstoken",
-		sessionData, ok := obj.(*SessionData)
+			refreshToken:        "shortrefreshtoken",
-		if !ok {
+			expectedCookieCount: 3, // main, access, refresh
-			// Return the object even if it's not the right type to avoid leaks
+			wantCompressed:      true,
-			sm.sessionPool.Put(obj)
+		},
-			break
+		{
-		}
+			name:                "Long tokens exceeding 4096 bytes",
-
+			authenticated:       true,
-		objects = append(objects, sessionData)
+			email:               "[email protected]",
+			accessToken:         strings.Repeat("x", 5000),
+			refreshToken:        strings.Repeat("y", 6000),
+			expectedCookieCount: calculateExpectedCookieCount(strings.Repeat("x", 5000), strings.Repeat("y", 6000)),
+			wantCompressed:      true,
+		},
+		{
+			name:                "REALLY long tokens, exceeding 25000 bytes",
+			authenticated:       true,
+			email:               "[email protected]",
+			accessToken:         strings.Repeat("x", 25000),
+			refreshToken:        strings.Repeat("y", 25000),
+			expectedCookieCount: calculateExpectedCookieCount(strings.Repeat("x", 25000), strings.Repeat("y", 25000)),
+			wantCompressed:      true,
+		},
+		{
+			name:                "Unauthenticated session",
+			authenticated:       false,
+			email:               "",
+			accessToken:         "",
+			refreshToken:        "",
+			expectedCookieCount: 3, // main, access, refresh
+			wantCompressed:      false,
+		},
+		{
+			name:                "Random content tokens",
+			authenticated:       true,
+			email:               "[email protected]",
+			accessToken:         generateRandomString(5000),
+			refreshToken:        generateRandomString(5000),
+			expectedCookieCount: calculateExpectedCookieCount(generateRandomString(5000), generateRandomString(5000)),
+			wantCompressed:      true,
+		},
 	}
 
-	// Count how many objects we found
+	for _, tc := range tests {
-	count := len(objects)
+		tc := tc // Capture range variable
+		t.Run(tc.name, func(t *testing.T) {
+			req := httptest.NewRequest("GET", "/test", nil)
+			rr := httptest.NewRecorder()
 
-	// Return all objects back to the pool to preserve the pool state
+			session, err := ts.sessionManager.GetSession(req)
-	for _, obj := range objects {
+			if err != nil {
-		sm.sessionPool.Put(obj)
+				t.Fatalf("Failed to get session: %v", err)
+			}
+
+			// Set session values
+			session.SetAuthenticated(tc.authenticated)
+			session.SetEmail(tc.email)
+
+			// Expire any existing cookies
+			session.expireAccessTokenChunks(rr)
+			session.expireRefreshTokenChunks(rr)
+
+			// Set new tokens
+			session.SetAccessToken(tc.accessToken)
+			session.SetRefreshToken(tc.refreshToken)
+
+			// Save session
+			if err := session.Save(req, rr); err != nil {
+				t.Fatalf("Failed to save session: %v", err)
+			}
+
+			// Verify cookies are set and compression is used when appropriate
+			cookies := rr.Result().Cookies()
+			if len(cookies) != tc.expectedCookieCount {
+				t.Errorf("Expected %d cookies, got %d", tc.expectedCookieCount, len(cookies))
+			}
+
+			// Verify compression is working by checking token sizes
+			for _, cookie := range cookies {
+				if strings.Contains(cookie.Name, accessTokenCookie) {
+					// Get original and stored sizes
+					originalSize := len(tc.accessToken)
+					storedSize := len(cookie.Value)
+
+					if originalSize > 100 && tc.wantCompressed {
+						// For large tokens, verify some compression occurred
+						compressionRatio := float64(storedSize) / float64(originalSize)
+						t.Logf("Access token compression ratio: %.2f (original: %d, stored: %d)",
+							compressionRatio, originalSize, storedSize)
+
+						if compressionRatio > 0.9 { // Allow some overhead, but should see compression
+							t.Errorf("Expected compression for large token in cookie %s (ratio: %.2f)",
+								cookie.Name, compressionRatio)
+						}
+					}
+				} else if strings.Contains(cookie.Name, refreshTokenCookie) {
+					originalSize := len(tc.refreshToken)
+					storedSize := len(cookie.Value)
+
+					if originalSize > 100 && tc.wantCompressed {
+						compressionRatio := float64(storedSize) / float64(originalSize)
+						t.Logf("Refresh token compression ratio: %.2f (original: %d, stored: %d)",
+							compressionRatio, originalSize, storedSize)
+
+						if compressionRatio > 0.9 {
+							t.Errorf("Expected compression for large token in cookie %s (ratio: %.2f)",
+								cookie.Name, compressionRatio)
+						}
+					}
+				}
+			}
+
+			// Create a new request with the cookies
+			newReq := httptest.NewRequest("GET", "/test", nil)
+			for _, cookie := range cookies {
+				newReq.AddCookie(cookie)
+			}
+
+			// Get the session again and verify values
+			newSession, err := ts.sessionManager.GetSession(newReq)
+			if err != nil {
+				t.Fatalf("Failed to get new session: %v", err)
+			}
+
+			// Verify session values
+			if newSession.GetAuthenticated() != tc.authenticated {
+				t.Errorf("Authentication status not preserved")
+			}
+			if email := newSession.GetEmail(); email != tc.email {
+				t.Errorf("Expected email %s, got %s", tc.email, email)
+			}
+			if token := newSession.GetAccessToken(); token != tc.accessToken {
+				t.Errorf("Access token not preserved: got len=%d, want len=%d", len(token), len(tc.accessToken))
+			}
+			if token := newSession.GetRefreshToken(); token != tc.refreshToken {
+				t.Errorf("Refresh token not preserved: got len=%d, want len=%d", len(token), len(tc.refreshToken))
+			}
+
+			// Verify session pooling by checking if the session is reused
+			session2, _ := ts.sessionManager.GetSession(newReq)
+			if session2 == newSession {
+				t.Error("Session not properly pooled")
+			}
+		})
+	}
+}
+
+func calculateExpectedCookieCount(accessToken, refreshToken string) int {
+	count := 3 // main, access, refresh
+
+	// Helper to calculate chunks for compressed token
+	calculateChunks := func(token string) int {
+		// Compress token (matching the actual implementation)
+		compressed := compressToken(token)
+
+		// If compressed token fits in one cookie, no additional chunks needed
+		if len(compressed) <= maxCookieSize {
+			return 0
+		}
+
+		// Calculate chunks needed for compressed token
+		return len(splitIntoChunks(compressed, maxCookieSize))
+	}
+
+	// Add chunks for access token if needed
+	accessChunks := calculateChunks(accessToken)
+	if accessChunks > 0 {
+		count += accessChunks
+	}
+
+	// Add chunks for refresh token if needed
+	refreshChunks := calculateChunks(refreshToken)
+	if refreshChunks > 0 {
+		count += refreshChunks
 	}
 
 	return count
 }
-
-// TestSessionObjectTracking verifies that session objects are properly
-// returned to the pool in various scenarios including normal usage and error paths
-func TestSessionObjectTracking(t *testing.T) {
-	logger := NewLogger("debug")
-	sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
-	if err != nil {
-		t.Fatalf("Failed to create session manager: %v", err)
-	}
-
-	// Create a fake request
-	req := httptest.NewRequest("GET", "http://example.com/foo", nil)
-
-	// Test that the session pool is used as expected
-	hasNew := sm.sessionPool.New != nil
-	if !hasNew {
-		t.Error("Expected sessionPool.New function to be set")
-	}
-
-	// Create and discard 5 sessions
-	for i := 0; i < 5; i++ {
-		session, err := sm.GetSession(req)
-		if err != nil {
-			t.Fatalf("GetSession failed: %v", err)
-		}
-		session.ReturnToPool()
-	}
-
-	// Create a session and get an error when trying to clear it
-	session, err := sm.GetSession(req)
-	if err != nil {
-		t.Fatalf("GetSession failed: %v", err)
-	}
-
-	// Deliberately cause bad state in the session object
-	session.mainSession = nil // This will cause an error in Clear
-
-	// Even with an error, the pool should not leak
-	session.ReturnToPool()
-
-	runtime.GC()
-	time.Sleep(100 * time.Millisecond)
-
-	// Success - if we got here without crashing, the pool is working as expected
-	t.Log("Session pool handling verified")
-}
-
-// This is intentionally left empty to remove unused code

settings.go

@@ -10,18 +10,6 @@ import (
 	"strings"
 )
 
-// TemplatedHeader represents a custom HTTP header with a templated value.
-// The value can contain template expressions that will be evaluated for each
-// authenticated request, such as {{.claims.email}} or {{.accessToken}}.
-type TemplatedHeader struct {
-	// Name is the HTTP header name to set (e.g., "X-Forwarded-Email")
-	Name string `json:"name"`
-
-	// Value is the template string for the header value
-	// Example: "{{.claims.email}}", "Bearer {{.accessToken}}"
-	Value string `json:"value"`
-}
-
 // Config holds the configuration for the OIDC middleware.
 // It provides all necessary settings to configure OpenID Connect authentication
 // with various providers like Auth0, Logto, or any standard OIDC provider.
@@ -82,10 +70,6 @@ type Config struct {
 	// Example: ["company.com", "subsidiary.com"]
 	AllowedUserDomains []string `json:"allowedUserDomains"`
 
-	// AllowedUsers restricts access to specific email addresses (optional)
-	// Example: ["[email protected]", "[email protected]"]
-	AllowedUsers []string `json:"allowedUsers"`
-
 	// AllowedRolesAndGroups restricts access to users with specific roles or groups (optional)
 	// Example: ["admin", "developer"]
 	AllowedRolesAndGroups []string `json:"allowedRolesAndGroups"`
@@ -100,22 +84,6 @@ type Config struct {
 
 	// HTTPClient allows customizing the HTTP client used for OIDC operations (optional)
 	HTTPClient *http.Client
-
-	// RefreshGracePeriodSeconds defines how many seconds before a token expires
-	// the plugin should attempt to refresh it proactively (optional)
-	// Default: 60
-	RefreshGracePeriodSeconds int `json:"refreshGracePeriodSeconds"`
-	// Headers defines custom HTTP headers to set with templated values (optional)
-	// Values can reference tokens and claims using Go templates with the following variables:
-	// - {{.AccessToken}} - The access token (ID token)
-	// - {{.IdToken}} - Same as AccessToken (for consistency)
-	// - {{.RefreshToken}} - The refresh token
-	// - {{.Claims.email}} - Access token claims (use proper case for claim names)
-	// Examples:
-	//
-	//	[{Name: "X-Forwarded-Email", Value: "{{.Claims.email}}"}]
-	//	[{Name: "Authorization", Value: "Bearer {{.AccessToken}}"}]
-	Headers []TemplatedHeader `json:"headers"`
 }
 
 const (
@@ -141,35 +109,21 @@ const (
 //   - PostLogoutRedirectURI: "/"
 //   - ForceHTTPS: true (for security)
 //   - EnablePKCE: false (PKCE is opt-in)
-//
-// CreateConfig initializes a new Config struct with default values for optional fields.
-// It sets default scopes, log level, rate limit, enables ForceHTTPS, and sets the
-// default refresh grace period. Required fields like ProviderURL, ClientID, ClientSecret,
-// CallbackURL, and SessionEncryptionKey must be set explicitly after creation.
-//
-// Returns:
-//   - A pointer to a new Config struct with default settings applied.
 func CreateConfig() *Config {
 	c := &Config{
-		Scopes:                    []string{"openid", "profile", "email"},
+		Scopes:     []string{"openid", "profile", "email"},
-		LogLevel:                  DefaultLogLevel,
+		LogLevel:   DefaultLogLevel,
-		RateLimit:                 DefaultRateLimit,
+		RateLimit:  DefaultRateLimit,
-		ForceHTTPS:                true,  // Secure by default
+		ForceHTTPS: true,  // Secure by default
-		EnablePKCE:                false, // PKCE is opt-in
+		EnablePKCE: false, // PKCE is opt-in
-		RefreshGracePeriodSeconds: 60,    // Default grace period of 60 seconds
 	}
 
 	return c
 }
 
-// Validate checks the configuration settings for validity.
+// Validate performs validation checks on the Config.
-// It ensures that required fields (ProviderURL, CallbackURL, ClientID, ClientSecret, SessionEncryptionKey)
+// It ensures all required fields are set and have valid values.
-// are present and that URLs are well-formed (HTTPS where required). It also validates
+// Returns an error if any validation check fails.
-// the session key length, log level, rate limit, and refresh grace period.
-//
-// Returns:
-//   - nil if the configuration is valid.
-//   - An error describing the first validation failure encountered.
 func (c *Config) Validate() error {
 	// Validate provider URL
 	if c.ProviderURL == "" {
@@ -243,61 +197,16 @@ func (c *Config) Validate() error {
 		return fmt.Errorf("rateLimit must be at least %d", MinRateLimit)
 	}
 
-	// Validate refresh grace period
-	if c.RefreshGracePeriodSeconds < 0 {
-		return fmt.Errorf("refreshGracePeriodSeconds cannot be negative")
-	}
-
-	// Validate headers configuration
-	for _, header := range c.Headers {
-		if header.Name == "" {
-			return fmt.Errorf("header name cannot be empty")
-		}
-		if header.Value == "" {
-			return fmt.Errorf("header value template cannot be empty")
-		}
-		if !strings.Contains(header.Value, "{{") || !strings.Contains(header.Value, "}}") {
-			return fmt.Errorf("header value '%s' does not appear to be a valid template (missing {{ }})", header.Value)
-		}
-
-		// Provide more helpful guidance for common template errors
-		if strings.Contains(header.Value, "{{.claims") {
-			return fmt.Errorf("header template '%s' appears to use lowercase 'claims' - use '{{.Claims...' instead (case sensitive)", header.Value)
-		}
-		if strings.Contains(header.Value, "{{.accessToken") {
-			return fmt.Errorf("header template '%s' appears to use lowercase 'accessToken' - use '{{.AccessToken...' instead (case sensitive)", header.Value)
-		}
-		if strings.Contains(header.Value, "{{.idToken") {
-			return fmt.Errorf("header template '%s' appears to use lowercase 'idToken' - use '{{.IdToken...' instead (case sensitive)", header.Value)
-		}
-		if strings.Contains(header.Value, "{{.refreshToken") {
-			return fmt.Errorf("header template '%s' appears to use lowercase 'refreshToken' - use '{{.RefreshToken...' instead (case sensitive)", header.Value)
-		}
-	}
-
 	return nil
 }
 
-// isValidSecureURL checks if a given string represents a valid, absolute HTTPS URL.
+// isValidSecureURL checks if the provided string is a valid HTTPS URL
-// It uses url.Parse and checks for a nil error, an "https" scheme, and a non-empty host.
-//
-// Parameters:
-//   - s: The URL string to validate.
-//
-// Returns:
-//   - true if the string is a valid HTTPS URL, false otherwise.
 func isValidSecureURL(s string) bool {
 	u, err := url.Parse(s)
 	return err == nil && u.Scheme == "https" && u.Host != ""
 }
 
-// isValidLogLevel checks if the provided log level string is one of the supported values ("debug", "info", "error").
+// isValidLogLevel checks if the provided log level is valid
-//
-// Parameters:
-//   - level: The log level string to validate.
-//
-// Returns:
-//   - true if the log level is valid, false otherwise.
 func isValidLogLevel(level string) bool {
 	return level == "debug" || level == "info" || level == "error"
 }
@@ -314,20 +223,14 @@ type Logger struct {
 	logDebug *log.Logger
 }
 
-// NewLogger creates and configures a new Logger instance based on the provided log level.
+// NewLogger creates a new Logger with the specified log level.
-// It initializes loggers for ERROR (stderr), INFO (stdout), and DEBUG (stdout) levels,
+// The log level determines which messages are output:
-// enabling output based on the specified level:
+//   - "debug": Outputs all messages (debug, info, error)
-//   - "error": Only ERROR messages are output.
+//   - "info": Outputs info and error messages
-//   - "info": INFO and ERROR messages are output.
+//   - "error": Outputs only error messages
-//   - "debug": DEBUG, INFO, and ERROR messages are output.
 //
-// If an invalid level is provided, it defaults to behavior similar to "error".
+// Error messages are always written to stderr, while info and debug
-//
+// messages are written to stdout when enabled.
-// Parameters:
-//   - logLevel: The desired logging level ("debug", "info", or "error").
-//
-// Returns:
-//   - A pointer to the configured Logger instance.
 func NewLogger(logLevel string) *Logger {
 	logError := log.New(io.Discard, "ERROR: TraefikOidcPlugin: ", log.Ldate|log.Ltime)
 	logInfo := log.New(io.Discard, "INFO: TraefikOidcPlugin: ", log.Ldate|log.Ltime)
@@ -349,77 +252,51 @@ func NewLogger(logLevel string) *Logger {
 	}
 }
 
-// Info logs a message at the INFO level using Printf style formatting.
+// Info logs an informational message.
-// Output is directed to stdout if the configured log level is "info" or "debug".
+// These messages are intended for general operational information
-//
+// and are written to stdout.
-// Parameters:
-//   - format: The format string (as in fmt.Printf).
-//   - args: The arguments for the format string.
 func (l *Logger) Info(format string, args ...interface{}) {
 	l.logInfo.Printf(format, args...)
 }
 
-// Debug logs a message at the DEBUG level using Printf style formatting.
+// Debug logs a debug message.
-// Output is directed to stdout only if the configured log level is "debug".
+// These messages are only output when debug level logging is enabled
-//
+// and are intended for detailed troubleshooting information.
-// Parameters:
-//   - format: The format string (as in fmt.Printf).
-//   - args: The arguments for the format string.
 func (l *Logger) Debug(format string, args ...interface{}) {
 	l.logDebug.Printf(format, args...)
 }
 
-// Error logs a message at the ERROR level using Printf style formatting.
+// Error logs an error message.
-// Output is always directed to stderr, regardless of the configured log level.
+// These messages indicate problems that need attention and are
-//
+// always written to stderr regardless of the log level.
-// Parameters:
-//   - format: The format string (as in fmt.Printf).
-//   - args: The arguments for the format string.
 func (l *Logger) Error(format string, args ...interface{}) {
 	l.logError.Printf(format, args...)
 }
 
-// Infof logs a message at the INFO level using Printf style formatting.
+// Infof logs an informational message using Printf formatting.
-// Equivalent to calling l.Info(format, args...).
+// These messages are intended for general operational information
-// Output is directed to stdout if the configured log level is "info" or "debug".
+// and are written to stdout.
-//
-// Parameters:
-//   - format: The format string (as in fmt.Printf).
-//   - args: The arguments for the format string.
 func (l *Logger) Infof(format string, args ...interface{}) {
 	l.logInfo.Printf(format, args...)
 }
 
-// Debugf logs a message at the DEBUG level using Printf style formatting.
+// Debugf logs a debug message using Printf formatting.
-// Equivalent to calling l.Debug(format, args...).
+// These messages are only output when debug level logging is enabled
-// Output is directed to stdout only if the configured log level is "debug".
+// and are intended for detailed troubleshooting information.
-//
-// Parameters:
-//   - format: The format string (as in fmt.Printf).
-//   - args: The arguments for the format string.
 func (l *Logger) Debugf(format string, args ...interface{}) {
 	l.logDebug.Printf(format, args...)
 }
 
-// Errorf logs a message at the ERROR level using Printf style formatting.
+// Errorf logs an error message using Printf formatting.
-// Equivalent to calling l.Error(format, args...).
+// These messages indicate problems that need attention and are
-// Output is always directed to stderr, regardless of the configured log level.
+// always written to stderr regardless of the log level.
-//
-// Parameters:
-//   - format: The format string (as in fmt.Printf).
-//   - args: The arguments for the format string.
 func (l *Logger) Errorf(format string, args ...interface{}) {
 	l.logError.Printf(format, args...)
 }
 
-// handleError logs an error message using the provided logger and sends an HTTP error
+// handleError writes an error message to both the HTTP response and the error log.
-// response to the client with the specified message and status code.
+// It ensures consistent error handling across the middleware by logging the error
-//
+// and sending an appropriate HTTP response to the client.
-// Parameters:
-//   - w: The http.ResponseWriter to send the error response to.
-//   - message: The error message string.
-//   - code: The HTTP status code for the response.
-//   - logger: The Logger instance to use for logging the error.
 func handleError(w http.ResponseWriter, message string, code int, logger *Logger) {
 	logger.Error(message)
 	http.Error(w, message, code)

settings_test.go

@@ -202,20 +202,6 @@ func TestConfigValidate(t *testing.T) {
 			},
 			expectedError: "",
 		},
-		{
-			name: "Valid Config With AllowedUsers",
-			config: &Config{
-				ProviderURL:          "https://provider.com",
-				CallbackURL:          "/callback",
-				ClientID:             "client-id",
-				ClientSecret:         "client-secret",
-				SessionEncryptionKey: "this-is-a-long-enough-encryption-key",
-				LogLevel:             "debug",
-				RateLimit:            100,
-				AllowedUsers:         []string{"[email protected]", "[email protected]"},
-			},
-			expectedError: "",
-		},
 	}
 
 	for _, tc := range tests {

templated_header_config_test.go

@@ -1,197 +0,0 @@
-package traefikoidc
-
-import (
-	"testing"
-	"text/template"
-)
-
-func TestTemplatedHeaderValidation(t *testing.T) {
-	tests := []struct {
-		name          string
-		header        TemplatedHeader
-		expectedError string
-	}{
-		{
-			name:          "Empty Name",
-			header:        TemplatedHeader{Name: "", Value: "{{.Claims.email}}"},
-			expectedError: "header name cannot be empty",
-		},
-		{
-			name:          "Empty Value",
-			header:        TemplatedHeader{Name: "X-Email", Value: ""},
-			expectedError: "header value template cannot be empty",
-		},
-		{
-			name:          "Not a Template",
-			header:        TemplatedHeader{Name: "X-Email", Value: "static-value"},
-			expectedError: "header value 'static-value' does not appear to be a valid template (missing {{ }})",
-		},
-		{
-			name:          "Lowercase claims",
-			header:        TemplatedHeader{Name: "X-Email", Value: "{{.claims.email}}"},
-			expectedError: "header template '{{.claims.email}}' appears to use lowercase 'claims' - use '{{.Claims...' instead (case sensitive)",
-		},
-		{
-			name:          "Lowercase accessToken",
-			header:        TemplatedHeader{Name: "X-Token", Value: "Bearer {{.accessToken}}"},
-			expectedError: "header template 'Bearer {{.accessToken}}' appears to use lowercase 'accessToken' - use '{{.AccessToken...' instead (case sensitive)",
-		},
-		{
-			name:          "Lowercase idToken",
-			header:        TemplatedHeader{Name: "X-Token", Value: "Bearer {{.idToken}}"},
-			expectedError: "header template 'Bearer {{.idToken}}' appears to use lowercase 'idToken' - use '{{.IdToken...' instead (case sensitive)",
-		},
-		{
-			name:          "Lowercase refreshToken",
-			header:        TemplatedHeader{Name: "X-Refresh", Value: "Bearer {{.refreshToken}}"},
-			expectedError: "header template 'Bearer {{.refreshToken}}' appears to use lowercase 'refreshToken' - use '{{.RefreshToken...' instead (case sensitive)",
-		},
-		{
-			name:          "Valid Template",
-			header:        TemplatedHeader{Name: "X-Email", Value: "{{.Claims.email}}"},
-			expectedError: "",
-		},
-		{
-			name:          "Valid Bearer Token Template",
-			header:        TemplatedHeader{Name: "Authorization", Value: "Bearer {{.AccessToken}}"},
-			expectedError: "",
-		},
-	}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			config := &Config{
-				ProviderURL:          "https://provider.com",
-				CallbackURL:          "/callback",
-				ClientID:             "client-id",
-				ClientSecret:         "client-secret",
-				SessionEncryptionKey: "this-is-a-long-enough-encryption-key",
-				RateLimit:            10, // Adding minimum required rate limit
-				Headers:              []TemplatedHeader{tc.header},
-			}
-
-			err := config.Validate()
-			if tc.expectedError == "" {
-				if err != nil {
-					t.Errorf("Expected no error, got: %v", err)
-				}
-			} else {
-				if err == nil {
-					t.Errorf("Expected error: %s, got nil", tc.expectedError)
-				} else if err.Error() != tc.expectedError {
-					t.Errorf("Expected error: %s, got: %s", tc.expectedError, err.Error())
-				}
-			}
-		})
-	}
-}
-
-func TestTemplateParsingInNew(t *testing.T) {
-	// Test successful parsing of templates during middleware creation
-	tests := []struct {
-		name              string
-		headers           []TemplatedHeader
-		expectedTemplates int
-		expectError       bool
-	}{
-		{
-			name: "Single Valid Template",
-			headers: []TemplatedHeader{
-				{Name: "X-Email", Value: "{{.Claims.email}}"},
-			},
-			expectedTemplates: 1,
-			expectError:       false,
-		},
-		{
-			name: "Multiple Valid Templates",
-			headers: []TemplatedHeader{
-				{Name: "X-Email", Value: "{{.Claims.email}}"},
-				{Name: "X-User-ID", Value: "{{.Claims.sub}}"},
-				{Name: "Authorization", Value: "Bearer {{.AccessToken}}"},
-			},
-			expectedTemplates: 3,
-			expectError:       false,
-		},
-		{
-			name: "Invalid Template",
-			headers: []TemplatedHeader{
-				{Name: "X-Email", Value: "{{.Claims.email"}, // Missing closing braces
-			},
-			expectedTemplates: 0,
-			expectError:       true,
-		},
-		{
-			name: "Mix of Valid and Invalid Templates",
-			headers: []TemplatedHeader{
-				{Name: "X-Email", Value: "{{.Claims.email}}"},
-				{Name: "X-Invalid", Value: "{{if .Claims.admin}}Admin{{end"}, // Invalid template
-			},
-			expectedTemplates: 1,    // Only the valid template should be parsed
-			expectError:       true, // We expect an error for the invalid template, but we'll handle it
-		},
-	}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			// For testing template parsing, we'll directly try to parse the templates instead of using New()
-			// This avoids the provider discovery that would fail in tests
-			headerTemplates := make(map[string]*template.Template)
-
-			// Special handling for the mixed valid/invalid templates case
-			if tc.name == "Mix of Valid and Invalid Templates" {
-				// Process templates one at a time so we can still have valid templates
-				for _, header := range tc.headers {
-					tmpl, err := template.New(header.Name).Parse(header.Value)
-					if err != nil {
-						// We expect an error for the invalid template
-						if !tc.expectError {
-							t.Errorf("Unexpected error parsing template %s: %v", header.Name, err)
-						}
-						// Skip this template but continue processing others
-						continue
-					}
-					headerTemplates[header.Name] = tmpl
-				}
-			} else {
-				// Normal handling for other test cases
-				var parseErr error
-				for _, header := range tc.headers {
-					tmpl, err := template.New(header.Name).Parse(header.Value)
-					if err != nil {
-						parseErr = err
-						break
-					}
-					headerTemplates[header.Name] = tmpl
-				}
-
-				if tc.expectError {
-					if parseErr == nil {
-						t.Error("Expected error parsing templates but got nil")
-					}
-					return
-				}
-
-				if parseErr != nil {
-					t.Fatalf("Unexpected error: %v", parseErr)
-				}
-			}
-
-			// Check the number of parsed templates
-			if len(headerTemplates) != tc.expectedTemplates {
-				t.Errorf("Expected %d parsed templates, got %d", tc.expectedTemplates, len(headerTemplates))
-			}
-
-			// Check each template was parsed
-			for _, header := range tc.headers {
-				// Skip the known invalid templates
-				if header.Value == "{{.Claims.email" || header.Value == "{{if .Claims.admin}}Admin{{end" {
-					continue
-				}
-
-				if _, ok := headerTemplates[header.Name]; !ok {
-					t.Errorf("Template for header %s was not parsed", header.Name)
-				}
-			}
-		})
-	}
-}

templated_header_execution_test.go

@@ -1,237 +0,0 @@
-package traefikoidc
-
-import (
-	"bytes"
-	"testing"
-	"text/template"
-)
-
-// TestTemplateExecution tests that templates are executed correctly with different types of claims
-func TestTemplateExecution(t *testing.T) {
-	tests := []struct {
-		name          string
-		templateText  string
-		data          map[string]interface{}
-		expectedValue string
-		expectError   bool
-	}{
-		{
-			name:         "String Claim",
-			templateText: "{{.Claims.email}}",
-			data: map[string]interface{}{
-				"Claims": map[string]interface{}{
-					"email": "[email protected]",
-				},
-			},
-			expectedValue: "[email protected]",
-			expectError:   false,
-		},
-		{
-			name:         "Number Claim",
-			templateText: "{{.Claims.age}}",
-			data: map[string]interface{}{
-				"Claims": map[string]interface{}{
-					"age": 30,
-				},
-			},
-			expectedValue: "30",
-			expectError:   false,
-		},
-		{
-			name:         "Boolean Claim",
-			templateText: "{{.Claims.admin}}",
-			data: map[string]interface{}{
-				"Claims": map[string]interface{}{
-					"admin": true,
-				},
-			},
-			expectedValue: "true",
-			expectError:   false,
-		},
-		{
-			name:         "Array Claim",
-			templateText: "{{index .Claims.roles 0}}",
-			data: map[string]interface{}{
-				"Claims": map[string]interface{}{
-					"roles": []string{"admin", "user"},
-				},
-			},
-			expectedValue: "admin",
-			expectError:   false,
-		},
-		{
-			name:         "Nested Object Claim",
-			templateText: "{{.Claims.user.name}}",
-			data: map[string]interface{}{
-				"Claims": map[string]interface{}{
-					"user": map[string]interface{}{
-						"name": "John Doe",
-					},
-				},
-			},
-			expectedValue: "John Doe",
-			expectError:   false,
-		},
-		{
-			name:         "Access Token",
-			templateText: "Bearer {{.AccessToken}}",
-			data: map[string]interface{}{
-				"AccessToken": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.e30.Et9HFtf9R3GEMA0IICOfFMVXY7kkTX1wr4qCyhIf58U",
-			},
-			expectedValue: "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.e30.Et9HFtf9R3GEMA0IICOfFMVXY7kkTX1wr4qCyhIf58U",
-			expectError:   false,
-		},
-		{
-			name:         "ID Token",
-			templateText: "{{.IdToken}}",
-			data: map[string]interface{}{
-				"IdToken": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.e30.Et9HFtf9R3GEMA0IICOfFMVXY7kkTX1wr4qCyhIf58U",
-			},
-			expectedValue: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.e30.Et9HFtf9R3GEMA0IICOfFMVXY7kkTX1wr4qCyhIf58U",
-			expectError:   false,
-		},
-		{
-			name:         "Refresh Token",
-			templateText: "{{.RefreshToken}}",
-			data: map[string]interface{}{
-				"RefreshToken": "refresh-token-value",
-			},
-			expectedValue: "refresh-token-value",
-			expectError:   false,
-		},
-		{
-			name:         "Conditional Template",
-			templateText: "{{if .Claims.admin}}Admin User{{else}}Regular User{{end}}",
-			data: map[string]interface{}{
-				"Claims": map[string]interface{}{
-					"admin": true,
-				},
-			},
-			expectedValue: "Admin User",
-			expectError:   false,
-		},
-		{
-			name:         "Multiple Claims",
-			templateText: "{{.Claims.firstName}} {{.Claims.lastName}} <{{.Claims.email}}>",
-			data: map[string]interface{}{
-				"Claims": map[string]interface{}{
-					"firstName": "John",
-					"lastName":  "Doe",
-					"email":     "[email protected]",
-				},
-			},
-			expectedValue: "John Doe <[email protected]>",
-			expectError:   false,
-		},
-		{
-			name:         "Missing Claim",
-			templateText: "{{.Claims.missing}}",
-			data: map[string]interface{}{
-				"Claims": map[string]interface{}{},
-			},
-			expectedValue: "<no value>",
-			expectError:   false, // Go templates don't error on missing values
-		},
-		{
-			name:         "Invalid Template Syntax",
-			templateText: "{{.Claims.email",
-			data: map[string]interface{}{
-				"Claims": map[string]interface{}{
-					"email": "[email protected]",
-				},
-			},
-			expectedValue: "",
-			expectError:   true, // Parsing should fail
-		},
-	}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			tmpl, err := template.New("test").Parse(tc.templateText)
-
-			if tc.expectError {
-				if err == nil {
-					t.Fatal("Expected template parsing error, but got nil")
-				}
-				return
-			}
-
-			if err != nil {
-				t.Fatalf("Failed to parse template: %v", err)
-			}
-
-			var buf bytes.Buffer
-			err = tmpl.Execute(&buf, tc.data)
-			if err != nil {
-				t.Fatalf("Failed to execute template: %v", err)
-			}
-
-			result := buf.String()
-			if result != tc.expectedValue {
-				t.Errorf("Expected template output %q, got %q", tc.expectedValue, result)
-			}
-		})
-	}
-}
-
-// TestTemplateExecutionContext tests the specific template data context used in processAuthorizedRequest
-func TestTemplateExecutionContext(t *testing.T) {
-	// Define a test struct that matches the one used in processAuthorizedRequest
-	type templateData struct {
-		AccessToken  string
-		IdToken      string
-		RefreshToken string
-		Claims       map[string]interface{}
-	}
-
-	// Test cases
-	tests := []struct {
-		name          string
-		templateText  string
-		data          templateData
-		expectedValue string
-	}{
-		{
-			name:         "Access and ID token distinction",
-			templateText: "Access: {{.AccessToken}} ID: {{.IdToken}}",
-			data: templateData{
-				AccessToken: "access-token-value",
-				IdToken:     "id-token-value", // Now these should be distinct values
-				Claims:      map[string]interface{}{},
-			},
-			expectedValue: "Access: access-token-value ID: id-token-value",
-		},
-		{
-			name:         "Combining tokens and claims",
-			templateText: "User: {{.Claims.sub}} Token: {{.AccessToken}}",
-			data: templateData{
-				AccessToken: "access-token",
-				IdToken:     "access-token",
-				Claims: map[string]interface{}{
-					"sub": "user123",
-				},
-			},
-			expectedValue: "User: user123 Token: access-token",
-		},
-	}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			tmpl, err := template.New("test").Parse(tc.templateText)
-			if err != nil {
-				t.Fatalf("Failed to parse template: %v", err)
-			}
-
-			var buf bytes.Buffer
-			err = tmpl.Execute(&buf, tc.data)
-			if err != nil {
-				t.Fatalf("Failed to execute template: %v", err)
-			}
-
-			result := buf.String()
-			if result != tc.expectedValue {
-				t.Errorf("Expected template output %q, got %q", tc.expectedValue, result)
-			}
-		})
-	}
-}

templated_header_integration_test.go

@@ -1,597 +0,0 @@
-package traefikoidc
-
-import (
-	"errors"
-	"net/http"
-	"net/http/httptest"
-	"testing"
-	"text/template"
-	"time"
-
-	"golang.org/x/time/rate"
-)
-
-// TestTemplatedHeadersIntegration tests that templated headers are correctly added to requests
-// in the actual middleware flow
-func TestTemplatedHeadersIntegration(t *testing.T) {
-	// Create a TestSuite to use its helper methods and fields
-	ts := &TestSuite{t: t}
-	ts.Setup()
-
-	tests := []struct {
-		name               string
-		headers            []TemplatedHeader
-		sessionSetup       func(*SessionData)
-		claims             map[string]interface{}
-		expectedHeaders    map[string]string
-		interceptedHeaders map[string]string
-	}{
-		{
-			name: "Basic Email Header",
-			headers: []TemplatedHeader{
-				{Name: "X-User-Email", Value: "{{.Claims.email}}"},
-			},
-			claims: map[string]interface{}{
-				"email": "[email protected]",
-			},
-			expectedHeaders: map[string]string{
-				"X-User-Email": "[email protected]",
-			},
-		},
-		{
-			name: "Multiple Headers",
-			headers: []TemplatedHeader{
-				{Name: "X-User-Email", Value: "{{.Claims.email}}"},
-				{Name: "X-User-ID", Value: "{{.Claims.sub}}"},
-				{Name: "X-User-Name", Value: "{{.Claims.given_name}} {{.Claims.family_name}}"},
-			},
-			claims: map[string]interface{}{
-				"email":       "[email protected]",
-				"sub":         "user123",
-				"given_name":  "John",
-				"family_name": "Doe",
-			},
-			expectedHeaders: map[string]string{
-				"X-User-Email": "[email protected]",
-				"X-User-ID":    "user123",
-				"X-User-Name":  "John Doe",
-			},
-		},
-		{
-			name: "Authorization Header with Bearer Token",
-			headers: []TemplatedHeader{
-				{Name: "Authorization", Value: "Bearer {{.AccessToken}}"},
-			},
-			expectedHeaders: map[string]string{
-				// We'll update this dynamically after generating the token
-				"Authorization": "",
-			},
-		},
-		{
-			name: "ID Token Header",
-			headers: []TemplatedHeader{
-				{Name: "X-ID-Token", Value: "{{.IdToken}}"},
-			},
-			expectedHeaders: map[string]string{
-				// We'll update this dynamically after generating the token
-				"X-ID-Token": "",
-			},
-		},
-		{
-			name: "Both Token Types",
-			headers: []TemplatedHeader{
-				{Name: "X-Access-Token", Value: "{{.AccessToken}}"},
-				{Name: "X-ID-Token", Value: "{{.IdToken}}"},
-			},
-			expectedHeaders: map[string]string{
-				// We'll update these dynamically after generating the tokens
-				"X-Access-Token": "",
-				"X-ID-Token":     "",
-			},
-		},
-		{
-			name: "Missing Claim",
-			headers: []TemplatedHeader{
-				{Name: "X-User-Role", Value: "{{.Claims.role}}"},
-			},
-			claims: map[string]interface{}{
-				"email": "[email protected]",
-				// role claim is missing
-			},
-			expectedHeaders: map[string]string{
-				"X-User-Role": "<no value>", // Go templates provide <no value> for missing fields
-			},
-		},
-		{
-			name: "Conditional Header",
-			headers: []TemplatedHeader{
-				{Name: "X-User-Admin", Value: "{{if .Claims.is_admin}}true{{else}}false{{end}}"},
-			},
-			claims: map[string]interface{}{
-				"email":    "[email protected]",
-				"is_admin": true,
-			},
-			expectedHeaders: map[string]string{
-				"X-User-Admin": "true",
-			},
-		},
-		{
-			name: "Combined Token and Claim",
-			headers: []TemplatedHeader{
-				{Name: "X-Auth-Info", Value: "User={{.Claims.email}}, Token={{.AccessToken}}"},
-			},
-			claims: map[string]interface{}{
-				"email": "[email protected]",
-			},
-			expectedHeaders: map[string]string{
-				// We'll update this dynamically after generating the token
-				"X-Auth-Info": "",
-			},
-		},
-		{
-			name: "Opaque Access Token with AccessTokenField",
-			headers: []TemplatedHeader{
-				{Name: "X-User-AccessToken", Value: "{{.AccessToken}}"},
-			},
-			claims: map[string]interface{}{ // For ID Token
-				"email": "[email protected]",
-				"sub":   "opaque_sub_for_id_token",
-			},
-			expectedHeaders: map[string]string{
-				"X-User-AccessToken": "this_is_an_opaque_access_token",
-			},
-		},
-	}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			// Create token with the test claims
-			token := ts.token
-			if len(tc.claims) > 0 {
-				var err error
-				baseClaims := map[string]interface{}{
-					"iss":   "https://test-issuer.com",
-					"aud":   "test-client-id",
-					"exp":   float64(3000000000), // Far future timestamp
-					"iat":   float64(1000000000),
-					"nbf":   float64(1000000000),
-					"sub":   "test-subject",
-					"nonce": "test-nonce",
-					"jti":   generateRandomString(16),
-				}
-
-				// Add the test-specific claims
-				for k, v := range tc.claims {
-					baseClaims[k] = v
-				}
-
-				token, err = createTestJWT(ts.rsaPrivateKey, "RS256", "test-key-id", baseClaims)
-				if err != nil {
-					t.Fatalf("Failed to create test JWT: %v", err)
-				}
-			}
-
-			// Update expectedHeaders for the token-based tests after token generation
-			if tc.name == "Authorization Header with Bearer Token" {
-				tc.expectedHeaders["Authorization"] = "Bearer " + token
-			}
-
-			if tc.name == "Combined Token and Claim" {
-				// If this test case uses specific ID/Access tokens, 'token' here might be just the ID token.
-				// This part might need adjustment if AccessToken is different and opaque.
-				// For now, assuming 'token' is the one to be used if not overridden later.
-				// The specific test "Opaque Access Token with AccessTokenField" will handle its AccessToken.
-				// This generic 'token' is used as a fallback if specific logic isn't hit.
-				// Let's ensure this test case uses the JWT access token if not otherwise specified.
-				accessTokenForHeader := token                        // Default to the generated JWT 'token'
-				if sessionVal, ok := tc.claims["_accessToken"]; ok { // Check if a specific access token is provided for this test
-					accessTokenForHeader = sessionVal.(string)
-				}
-				tc.expectedHeaders["X-Auth-Info"] = "User=" + tc.claims["email"].(string) + ", Token=" + accessTokenForHeader
-			}
-
-			// Store intercepted headers for verification
-			interceptedHeaders := make(map[string]string)
-
-			// Create a test next handler that captures the headers
-			nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-				// Capture headers for verification
-				for name := range tc.expectedHeaders {
-					if value := r.Header.Get(name); value != "" {
-						interceptedHeaders[name] = value
-					}
-				}
-				w.WriteHeader(http.StatusOK)
-			})
-
-			tOidc := &TraefikOidc{
-				next:               nextHandler,
-				name:               "test",
-				redirURLPath:       "/callback",
-				logoutURLPath:      "/callback/logout",
-				issuerURL:          "https://test-issuer.com",
-				clientID:           "test-client-id",
-				clientSecret:       "test-client-secret",
-				jwkCache:           ts.mockJWKCache,
-				jwksURL:            "https://test-jwks-url.com",
-				tokenBlacklist:     NewCache(),
-				tokenCache:         NewTokenCache(),
-				limiter:            rate.NewLimiter(rate.Every(time.Second), 10),
-				logger:             NewLogger("debug"),
-				allowedUserDomains: map[string]struct{}{"example.com": {}, "[email protected]": {}}, // Ensure domain for opaque test is allowed
-				excludedURLs:       map[string]struct{}{"/favicon": {}},
-				httpClient:         &http.Client{},
-				initComplete:       make(chan struct{}),
-				sessionManager:     ts.sessionManager,
-				extractClaimsFunc:  extractClaims,
-				headerTemplates:    make(map[string]*template.Template),
-				// Default to true, which means PopulateSessionWithIdTokenClaims is true
-				// UseIdTokenForSession: true, // Explicitly can be set if needed
-			}
-			tOidc.tokenVerifier = tOidc
-			tOidc.jwtVerifier = tOidc
-			tOidc.tokenExchanger = tOidc
-
-			// Initialize and parse header templates
-			for _, header := range tc.headers {
-				tmpl, err := template.New(header.Name).Parse(header.Value)
-				if err != nil {
-					t.Fatalf("Failed to parse header template for %s: %v", header.Name, err)
-				}
-				tOidc.headerTemplates[header.Name] = tmpl
-			}
-
-			close(tOidc.initComplete)
-
-			req := httptest.NewRequest("GET", "/protected", nil)
-			req.Header.Set("X-Forwarded-Proto", "https")
-			req.Header.Set("X-Forwarded-Host", "example.com")
-			rr := httptest.NewRecorder()
-
-			session, err := tOidc.sessionManager.GetSession(req)
-			if err != nil {
-				t.Fatalf("Failed to get session: %v", err)
-			}
-
-			session.SetAuthenticated(true)
-			// Set a default email; specific tests might override or rely on ID token population
-			defaultEmail := "[email protected]"
-			if emailClaim, ok := tc.claims["email"].(string); ok {
-				defaultEmail = emailClaim // Use email from claims if available for initial setup
-			}
-			session.SetEmail(defaultEmail)
-
-			// Default token setup (can be overridden by specific test cases below)
-			session.SetIDToken(token)
-			session.SetAccessToken(token)
-			session.SetRefreshToken("test-refresh-token")
-
-			if tc.name == "ID Token Header" || tc.name == "Both Token Types" {
-				idTokenClaims := map[string]interface{}{
-					"iss": "https://test-issuer.com", "aud": "test-client-id", "exp": float64(3000000000),
-					"iat": float64(1000000000), "nbf": float64(1000000000), "sub": "test-subject",
-					"nonce": "test-nonce", "jti": generateRandomString(16), "type": "id_token",
-					"email": tc.claims["email"], // Ensure email from test case claims is in ID token
-				}
-				// Add other claims from tc.claims to idTokenClaims
-				for k, v := range tc.claims {
-					if _, exists := idTokenClaims[k]; !exists {
-						idTokenClaims[k] = v
-					}
-				}
-
-				idTokenForSession, idErr := createTestJWT(ts.rsaPrivateKey, "RS256", "test-key-id", idTokenClaims)
-				if idErr != nil {
-					t.Fatalf("Failed to create test ID JWT: %v", idErr)
-				}
-
-				accessTokenClaims := map[string]interface{}{
-					"iss": "https://test-issuer.com", "aud": "test-client-id", "exp": float64(3000000000),
-					"iat": float64(1000000000), "nbf": float64(1000000000), "sub": "test-subject",
-					"jti": generateRandomString(16), "type": "access_token", "scope": "openid email profile",
-					"email": tc.claims["email"], // Include email in access token too for these tests
-				}
-				accessTokenForSession, accessErr := createTestJWT(ts.rsaPrivateKey, "RS256", "test-key-id", accessTokenClaims)
-				if accessErr != nil {
-					t.Fatalf("Failed to create test access JWT: %v", accessErr)
-				}
-
-				session.SetIDToken(idTokenForSession)
-				session.SetAccessToken(accessTokenForSession)
-
-				tOidc.tokenExchanger = &MockTokenExchanger{
-					RefreshTokenFunc: func(refreshToken string) (*TokenResponse, error) {
-						return &TokenResponse{
-							IDToken: idTokenForSession, AccessToken: accessTokenForSession,
-							RefreshToken: refreshToken, ExpiresIn: 3600,
-						}, nil
-					},
-				}
-				tOidc.tokenVerifier = &MockTokenVerifier{VerifyFunc: func(token string) error { return nil }}
-
-				if tc.name == "ID Token Header" {
-					tc.expectedHeaders["X-ID-Token"] = idTokenForSession
-				} else if tc.name == "Both Token Types" {
-					tc.expectedHeaders["X-ID-Token"] = idTokenForSession
-					tc.expectedHeaders["X-Access-Token"] = accessTokenForSession
-				}
-			} else if tc.name == "Opaque Access Token with AccessTokenField" {
-				idTokenClaims := map[string]interface{}{
-					"iss": "https://test-issuer.com", "aud": "test-client-id", "exp": float64(3000000000),
-					"iat": float64(1000000000), "nbf": float64(1000000000), "sub": "test-subject", // Default sub
-					"nonce": "test-nonce", "jti": generateRandomString(16), "type": "id_token",
-				}
-				// Populate ID token claims from tc.claims
-				for k, v := range tc.claims {
-					idTokenClaims[k] = v
-				}
-				// Ensure email from tc.claims is used for the ID token
-				session.SetEmail(tc.claims["email"].(string)) // Also set it directly for initial session state
-
-				idTokenForSession, err := createTestJWT(ts.rsaPrivateKey, "RS256", "test-key-id", idTokenClaims)
-				if err != nil {
-					t.Fatalf("Failed to create test ID JWT for opaque test: %v", err)
-				}
-
-				opaqueAccessToken := "this_is_an_opaque_access_token"
-
-				session.SetIDToken(idTokenForSession)
-				session.SetAccessToken(opaqueAccessToken)
-
-				tOidc.tokenExchanger = &MockTokenExchanger{
-					RefreshTokenFunc: func(refreshToken string) (*TokenResponse, error) {
-						return &TokenResponse{
-							IDToken:      idTokenForSession,
-							AccessToken:  opaqueAccessToken,
-							RefreshToken: refreshToken,
-							ExpiresIn:    3600,
-						}, nil
-					},
-				}
-				tOidc.tokenVerifier = &MockTokenVerifier{
-					VerifyFunc: func(tokenToVerify string) error {
-						if tokenToVerify == idTokenForSession {
-							return nil // ID token is expected to be verified
-						}
-						if tokenToVerify == opaqueAccessToken {
-							t.Errorf("TokenVerifier was incorrectly called with the opaque access token.")
-							return errors.New("opaque access token should not be verified by this path")
-						}
-						t.Logf("TokenVerifier called with unexpected token: %s", tokenToVerify)
-						return errors.New("unexpected token passed to verifier for this test case")
-					},
-				}
-				// Expected header X-User-AccessToken is already set in tc.expectedHeaders
-			}
-
-			if err := session.Save(req, rr); err != nil {
-				t.Fatalf("Failed to save session: %v", err)
-			}
-
-			for _, cookie := range rr.Result().Cookies() {
-				req.AddCookie(cookie)
-			}
-
-			rr = httptest.NewRecorder()
-			tOidc.ServeHTTP(rr, req)
-
-			if rr.Code != http.StatusOK {
-				t.Errorf("Expected status code %d, got %d. Body: %s", http.StatusOK, rr.Code, rr.Body.String())
-			}
-
-			for name, expectedValue := range tc.expectedHeaders {
-				if value, exists := interceptedHeaders[name]; !exists {
-					// For <no value> case, it might not be set if template resolves to empty and header is omitted.
-					// However, Go templates usually insert "<no value>" string.
-					if expectedValue == "<no value>" && tc.name == "Missing Claim" { // Special handling for <no value>
-						// If the template {{.Claims.role}} results in an empty string because role is missing,
-						// and the header is not set, this is also acceptable for "<no value>".
-						// The current test expects the literal string "<no value>".
-						// Let's assume for now that if it's missing, it's an error unless specifically handled.
-						// The test as written expects "<no value>" to be present.
-					}
-					t.Errorf("Expected header %s was not set", name)
-
-				} else if value != expectedValue {
-					t.Errorf("Header %s expected value %q, got %q", name, expectedValue, value)
-				}
-			}
-
-			if tc.name == "Opaque Access Token with AccessTokenField" {
-				postReq := httptest.NewRequest("GET", "/protected", nil)
-				for _, cookie := range rr.Result().Cookies() {
-					postReq.AddCookie(cookie)
-				}
-				updatedSession, err := tOidc.sessionManager.GetSession(postReq)
-				if err != nil {
-					t.Fatalf("Failed to get updated session for opaque test: %v", err)
-				}
-
-				expectedEmail := tc.claims["email"].(string)
-				if updatedSession.GetEmail() != expectedEmail {
-					t.Errorf("Expected session email to be %q (from ID token), got %q", expectedEmail, updatedSession.GetEmail())
-				}
-				if !updatedSession.GetAuthenticated() {
-					t.Errorf("Session should be authenticated after successful flow for opaque test")
-				}
-			}
-		})
-	}
-}
-
-// TestEdgeCaseTemplatedHeaders tests edge cases for templated headers
-func TestEdgeCaseTemplatedHeaders(t *testing.T) {
-	// Create a TestSuite to use its helper methods and fields
-	ts := &TestSuite{t: t}
-	ts.Setup()
-
-	tests := []struct {
-		name               string
-		headers            []TemplatedHeader
-		claims             map[string]interface{}
-		shouldExecuteCheck bool
-	}{
-		{
-			name: "Very Large Template",
-			headers: []TemplatedHeader{
-				{
-					Name:  "X-Large-Header",
-					Value: createLargeTemplate(500), // Template with 500 variable references
-				},
-			},
-			claims:             createLargeClaims(500), // Map with 500 claims
-			shouldExecuteCheck: true,
-		},
-		{
-			name: "Array Claim Access",
-			headers: []TemplatedHeader{
-				{Name: "X-Roles", Value: "{{range $i, $e := .Claims.roles}}{{if $i}},{{end}}{{$e}}{{end}}"},
-			},
-			claims: map[string]interface{}{
-				"roles": []interface{}{"admin", "user", "manager"},
-			},
-			shouldExecuteCheck: true,
-		},
-	}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			// Create token with the test claims
-			claims := map[string]interface{}{
-				"iss":   "https://test-issuer.com",
-				"aud":   "test-client-id",
-				"exp":   float64(3000000000), // Far future timestamp
-				"iat":   float64(1000000000),
-				"nbf":   float64(1000000000),
-				"sub":   "test-subject",
-				"nonce": "test-nonce",
-				"jti":   generateRandomString(16),
-			}
-
-			// Add the test-specific claims
-			for k, v := range tc.claims {
-				claims[k] = v
-			}
-
-			token, err := createTestJWT(ts.rsaPrivateKey, "RS256", "test-key-id", claims)
-			if err != nil {
-				t.Fatalf("Failed to create test JWT: %v", err)
-			}
-
-			// Create a test next handler
-			nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-				w.WriteHeader(http.StatusOK)
-			})
-
-			tOidc := &TraefikOidc{
-				next:               nextHandler,
-				name:               "test",
-				redirURLPath:       "/callback",
-				logoutURLPath:      "/callback/logout",
-				issuerURL:          "https://test-issuer.com",
-				clientID:           "test-client-id",
-				clientSecret:       "test-client-secret",
-				jwkCache:           ts.mockJWKCache,
-				jwksURL:            "https://test-jwks-url.com",
-				tokenBlacklist:     NewCache(),
-				tokenCache:         NewTokenCache(),
-				limiter:            rate.NewLimiter(rate.Every(time.Second), 10),
-				logger:             NewLogger("debug"),
-				allowedUserDomains: map[string]struct{}{"example.com": {}},
-				excludedURLs:       map[string]struct{}{"/favicon": {}},
-				httpClient:         &http.Client{},
-				initComplete:       make(chan struct{}),
-				sessionManager:     ts.sessionManager,
-				extractClaimsFunc:  extractClaims,
-				headerTemplates:    make(map[string]*template.Template),
-			}
-			tOidc.tokenVerifier = tOidc
-			tOidc.jwtVerifier = tOidc
-
-			// Initialize and parse header templates
-			for _, header := range tc.headers {
-				tmpl, err := template.New(header.Name).Parse(header.Value)
-				if err != nil {
-					t.Fatalf("Failed to parse header template for %s: %v", header.Name, err)
-				}
-				tOidc.headerTemplates[header.Name] = tmpl
-			}
-
-			close(tOidc.initComplete)
-
-			req := httptest.NewRequest("GET", "/protected", nil)
-			req.Header.Set("X-Forwarded-Proto", "https")
-			req.Header.Set("X-Forwarded-Host", "example.com")
-			rr := httptest.NewRecorder()
-
-			session, err := tOidc.sessionManager.GetSession(req)
-			if err != nil {
-				t.Fatalf("Failed to get session: %v", err)
-			}
-
-			session.SetAuthenticated(true)
-			session.SetEmail("[email protected]")
-			session.SetIDToken(token)     // Use the new method
-			session.SetAccessToken(token) // Also set access token to match
-			session.SetRefreshToken("test-refresh-token")
-
-			tOidc.extractClaimsFunc = extractClaims
-			tOidc.tokenExchanger = &MockTokenExchanger{
-				RefreshTokenFunc: func(refreshToken string) (*TokenResponse, error) {
-					return &TokenResponse{
-						IDToken:      token,
-						AccessToken:  token,
-						RefreshToken: refreshToken,
-						ExpiresIn:    3600,
-					}, nil
-				},
-			}
-
-			if err := session.Save(req, rr); err != nil {
-				t.Fatalf("Failed to save session: %v", err)
-			}
-
-			for _, cookie := range rr.Result().Cookies() {
-				req.AddCookie(cookie)
-			}
-
-			rr = httptest.NewRecorder()
-			tOidc.ServeHTTP(rr, req)
-
-			if rr.Code != http.StatusOK {
-				t.Errorf("Expected status code %d, got %d", http.StatusOK, rr.Code)
-			}
-
-			// The "Array Claim Access" check previously here was problematic as it didn't correctly
-			// intercept headers in TestEdgeCaseTemplatedHeaders. The primary goal of this
-			// function is to test edge cases for panics/errors, and robust header value
-			// checking is already covered in TestTemplatedHeadersIntegration.
-			// Removing the ineffective check to resolve the "declared and not used" error.
-		})
-	}
-}
-
-// Helper functions for edge case tests
-
-// createLargeTemplate creates a template with many variable references
-func createLargeTemplate(size int) string {
-	template := "{{with .Claims}}"
-	for i := 0; i < size; i++ {
-		if i > 0 {
-			template += ","
-		}
-		template += "{{.field" + string(rune('a'+i%26)) + string(rune('0'+i%10)) + "}}"
-	}
-	template += "{{end}}"
-	return template
-}
-
-// createLargeClaims creates a map with many claims for testing large templates
-func createLargeClaims(size int) map[string]interface{} {
-	claims := make(map[string]interface{})
-	for i := 0; i < size; i++ {
-		key := "field" + string(rune('a'+i%26)) + string(rune('0'+i%10))
-		claims[key] = "value" + string(rune('a'+i%26)) + string(rune('0'+i%10))
-	}
-	return claims
-}

token_handling_test.go

@@ -1,311 +0,0 @@
-package traefikoidc
-
-import (
-	"bytes"
-	"net/http"
-	"net/http/httptest"
-	"testing"
-	"text/template"
-	"time"
-
-	"golang.org/x/time/rate"
-)
-
-// TestTokenTypeDistinction tests that AccessToken and IdToken are correctly distinguished in templates
-func TestTokenTypeDistinction(t *testing.T) {
-	// Define test data where AccessToken and IdToken are deliberately different
-	type templateData struct {
-		AccessToken  string
-		IdToken      string
-		RefreshToken string
-		Claims       map[string]interface{}
-	}
-
-	testData := templateData{
-		AccessToken:  "test-access-token-abc123",
-		IdToken:      "test-id-token-xyz789",
-		RefreshToken: "test-refresh-token",
-		Claims: map[string]interface{}{
-			"sub":   "test-subject",
-			"email": "[email protected]",
-		},
-	}
-
-	// Test cases
-	tests := []struct {
-		name          string
-		templateText  string
-		expectedValue string
-	}{
-		{
-			name:          "Access Token Only",
-			templateText:  "Bearer {{.AccessToken}}",
-			expectedValue: "Bearer test-access-token-abc123",
-		},
-		{
-			name:          "ID Token Only",
-			templateText:  "ID: {{.IdToken}}",
-			expectedValue: "ID: test-id-token-xyz789",
-		},
-		{
-			name:          "Both Tokens",
-			templateText:  "Access: {{.AccessToken}} ID: {{.IdToken}}",
-			expectedValue: "Access: test-access-token-abc123 ID: test-id-token-xyz789",
-		},
-		{
-			name:          "Both Tokens in Authorization Format",
-			templateText:  "Bearer {{.AccessToken}} and Bearer {{.IdToken}}",
-			expectedValue: "Bearer test-access-token-abc123 and Bearer test-id-token-xyz789",
-		},
-	}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			tmpl, err := template.New("test").Parse(tc.templateText)
-			if err != nil {
-				t.Fatalf("Failed to parse template: %v", err)
-			}
-
-			var buf bytes.Buffer
-			err = tmpl.Execute(&buf, testData)
-			if err != nil {
-				t.Fatalf("Failed to execute template: %v", err)
-			}
-
-			result := buf.String()
-			if result != tc.expectedValue {
-				t.Errorf("Expected template output %q, got %q", tc.expectedValue, result)
-			}
-		})
-	}
-}
-
-// TestTokenTypeIntegration tests the integration of ID and access tokens with the middleware
-func TestTokenTypeIntegration(t *testing.T) {
-	// Create a TestSuite to use its helper methods and fields
-	ts := &TestSuite{t: t}
-	ts.Setup()
-
-	// Create different tokens for ID and access tokens
-	idToken, err := createTestJWT(ts.rsaPrivateKey, "RS256", "test-key-id", map[string]interface{}{
-		"iss":        "https://test-issuer.com",
-		"aud":        "test-client-id",
-		"exp":        float64(3000000000),
-		"iat":        float64(1000000000),
-		"nbf":        float64(1000000000),
-		"sub":        "test-subject",
-		"nonce":      "test-nonce",
-		"jti":        generateRandomString(16),
-		"token_type": "id_token",
-		"email":      "[email protected]",
-	})
-	if err != nil {
-		t.Fatalf("Failed to create test ID JWT: %v", err)
-	}
-
-	accessToken, err := createTestJWT(ts.rsaPrivateKey, "RS256", "test-key-id", map[string]interface{}{
-		"iss":        "https://test-issuer.com",
-		"aud":        "test-client-id",
-		"exp":        float64(3000000000),
-		"iat":        float64(1000000000),
-		"nbf":        float64(1000000000),
-		"sub":        "test-subject",
-		"jti":        generateRandomString(16),
-		"token_type": "access_token",
-		"scope":      "openid profile email",
-		"email":      "[email protected]", // Add email to access token so it's available in claims
-	})
-	if err != nil {
-		t.Fatalf("Failed to create test access JWT: %v", err)
-	}
-
-	// Define test headers that use both token types
-	headers := []TemplatedHeader{
-		{Name: "X-ID-Token", Value: "{{.IdToken}}"},
-		{Name: "X-Access-Token", Value: "{{.AccessToken}}"},
-		{Name: "Authorization", Value: "Bearer {{.AccessToken}}"},
-		{Name: "X-Email-From-Claims", Value: "{{.Claims.email}}"},
-	}
-
-	// Store intercepted headers for verification
-	interceptedHeaders := make(map[string]string)
-
-	// Create a test next handler that captures the headers
-	nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		// Capture headers for verification
-		for _, header := range headers {
-			if value := r.Header.Get(header.Name); value != "" {
-				interceptedHeaders[header.Name] = value
-			}
-		}
-		w.WriteHeader(http.StatusOK)
-	})
-
-	// Create the TraefikOidc instance
-	tOidc := &TraefikOidc{
-		next:               nextHandler,
-		name:               "test",
-		redirURLPath:       "/callback",
-		logoutURLPath:      "/callback/logout",
-		issuerURL:          "https://test-issuer.com",
-		clientID:           "test-client-id",
-		clientSecret:       "test-client-secret",
-		jwkCache:           ts.mockJWKCache,
-		jwksURL:            "https://test-jwks-url.com",
-		tokenBlacklist:     NewCache(),
-		tokenCache:         NewTokenCache(),
-		limiter:            rate.NewLimiter(rate.Every(time.Second), 10),
-		logger:             NewLogger("debug"),
-		allowedUserDomains: map[string]struct{}{"example.com": {}},
-		excludedURLs:       map[string]struct{}{"/favicon": {}},
-		httpClient:         &http.Client{},
-		initComplete:       make(chan struct{}),
-		sessionManager:     ts.sessionManager,
-		extractClaimsFunc:  extractClaims,
-		headerTemplates:    make(map[string]*template.Template),
-	}
-	tOidc.tokenVerifier = tOidc
-	tOidc.jwtVerifier = tOidc
-
-	// Initialize and parse header templates
-	for _, header := range headers {
-		tmpl, err := template.New(header.Name).Parse(header.Value)
-		if err != nil {
-			t.Fatalf("Failed to parse header template for %s: %v", header.Name, err)
-		}
-		tOidc.headerTemplates[header.Name] = tmpl
-	}
-
-	// Close the initComplete channel to bypass the waiting
-	close(tOidc.initComplete)
-
-	// Create a test request
-	req := httptest.NewRequest("GET", "/protected", nil)
-	req.Header.Set("X-Forwarded-Proto", "https")
-	req.Header.Set("X-Forwarded-Host", "example.com")
-	rr := httptest.NewRecorder()
-
-	// Create a session
-	session, err := tOidc.sessionManager.GetSession(req)
-	if err != nil {
-		t.Fatalf("Failed to get session: %v", err)
-	}
-
-	// Setup the session with authentication data
-	session.SetAuthenticated(true)
-	session.SetEmail("[email protected]")
-	session.SetIDToken(idToken)         // Set the ID token
-	session.SetAccessToken(accessToken) // Set the access token
-	session.SetRefreshToken("test-refresh-token")
-
-	if err := session.Save(req, rr); err != nil {
-		t.Fatalf("Failed to save session: %v", err)
-	}
-
-	// Add session cookies to the request
-	for _, cookie := range rr.Result().Cookies() {
-		req.AddCookie(cookie)
-	}
-
-	// Reset the response recorder for the main test
-	rr = httptest.NewRecorder()
-
-	// Process the request
-	tOidc.ServeHTTP(rr, req)
-
-	// Check status code
-	if rr.Code != http.StatusOK {
-		t.Errorf("Expected status code %d, got %d", http.StatusOK, rr.Code)
-	}
-
-	// Verify headers were set correctly
-	expectedHeaders := map[string]string{
-		"X-ID-Token":          idToken,
-		"X-Access-Token":      accessToken,
-		"Authorization":       "Bearer " + accessToken,
-		"X-Email-From-Claims": "[email protected]",
-	}
-
-	for name, expectedValue := range expectedHeaders {
-		if value, exists := interceptedHeaders[name]; !exists {
-			t.Errorf("Expected header %s was not set", name)
-		} else if value != expectedValue {
-			t.Errorf("Header %s expected value %q, got %q", name, expectedValue, value)
-		}
-	}
-}
-
-// TestSessionIDTokenAccessToken tests that the SessionData correctly stores and retrieves
-// both ID tokens and access tokens separately
-func TestSessionIDTokenAccessToken(t *testing.T) {
-	// Create a logger for the session manager
-	logger := NewLogger("debug")
-
-	// Create a session manager
-	sessionManager, err := NewSessionManager("test-session-encryption-key-at-least-32-bytes", false, logger)
-	if err != nil {
-		t.Fatalf("Failed to create session manager: %v", err)
-	}
-
-	// Create a test request
-	req := httptest.NewRequest("GET", "/test", nil)
-	rr := httptest.NewRecorder()
-
-	// Get a session
-	session, err := sessionManager.GetSession(req)
-	if err != nil {
-		t.Fatalf("Failed to get session: %v", err)
-	}
-
-	// Set test tokens
-	idToken := "test-id-token-123"
-	accessToken := "test-access-token-456"
-	refreshToken := "test-refresh-token-789"
-
-	// Store tokens in session
-	session.SetIDToken(idToken)
-	session.SetAccessToken(accessToken)
-	session.SetRefreshToken(refreshToken)
-
-	// Save the session
-	if err := session.Save(req, rr); err != nil {
-		t.Fatalf("Failed to save session: %v", err)
-	}
-
-	// Get cookies from response
-	cookies := rr.Result().Cookies()
-
-	// Create a new request with those cookies
-	req2 := httptest.NewRequest("GET", "/test", nil)
-	for _, cookie := range cookies {
-		req2.AddCookie(cookie)
-	}
-
-	// Get the session again
-	session2, err := sessionManager.GetSession(req2)
-	if err != nil {
-		t.Fatalf("Failed to get session from request with cookies: %v", err)
-	}
-
-	// Verify that the tokens were correctly stored and retrieved
-	retrievedIDToken := session2.GetIDToken()
-	retrievedAccessToken := session2.GetAccessToken()
-	retrievedRefreshToken := session2.GetRefreshToken()
-
-	if retrievedIDToken != idToken {
-		t.Errorf("ID token mismatch: expected %q, got %q", idToken, retrievedIDToken)
-	}
-
-	if retrievedAccessToken != accessToken {
-		t.Errorf("Access token mismatch: expected %q, got %q", accessToken, retrievedAccessToken)
-	}
-
-	if retrievedRefreshToken != refreshToken {
-		t.Errorf("Refresh token mismatch: expected %q, got %q", refreshToken, retrievedRefreshToken)
-	}
-
-	// Verify that the tokens are distinct
-	if retrievedIDToken == retrievedAccessToken {
-		t.Errorf("ID token and Access token should be different, but both are %q", retrievedIDToken)
-	}
-}