fixup! Dummy change of readme.

.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
@@ -157,9 +152,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 +211,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: |

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
 
@@ -73,7 +73,6 @@ The middleware supports the following configuration options:
 | `rateLimit` | Sets the maximum number of requests per second | `100` | `500` |
 | `excludedURLs` | Lists paths that bypass authentication | none | `["/health", "/metrics", "/public"]` |
 | `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"]` |
 | `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` |
@@ -160,67 +159,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
@@ -359,7 +297,7 @@ spec:
 
 ### Google OIDC Configuration Example
 
-This example shows a configuration specifically tailored for Google OIDC:
+This example shows a configuration specifically tailored for Google OIDC, including necessary scopes for session extension:
 
 ```yaml
 apiVersion: traefik.io/v1alpha1
@@ -380,14 +318,11 @@ spec:
         - openid
         - email
         - profile
-        # Note: DO NOT manually add offline_access scope for Google
-        # The middleware automatically handles Google-specific requirements
+        - offline_access # Required for refresh tokens / long sessions with Google
       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 +449,6 @@ http:
             - profile
           allowedUserDomains:
             - company.com
-          allowedUsers:
-            - [email protected]
-            - [email protected]
           allowedRolesAndGroups:
             - admin
             - developer
@@ -573,21 +505,6 @@ This middleware aims to provide long-lived user sessions, typically up to 24 hou
 - 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.
@@ -674,12 +591,10 @@ logLevel: debug
 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`).
+   - The `offline_access` scope is included in your configuration (the middleware adds this automatically now, but verify if manually configured).
    - 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).
+   - The fix involving automatic `offline_access` scope and `prompt=consent` for Google is active in your middleware version. Check the plugin version corresponds to when this fix was implemented. Enhanced logging around refresh token failures can provide more clues if issues persist.
 
 ## Contributing
 
-Contributions are welcome! Please feel free to submit a Pull Request.
+Contributions are welcome! Please feel free to submit a Pull Request and it'll be reviewed as soon as possible.

TODO.txt

@@ -1,5 +0,0 @@
-### TODO / wishlist
-
-- [] Improve test coverage
-- [x] Improve caching mechanism
-- [x] Add automatic release and semver generation

cache.go

@@ -149,8 +149,8 @@ func (c *Cache) Cleanup() {
 
 	now := time.Now()
 	for key, item := range c.items {
-		// Remove items that are expired
-		if now.After(item.ExpiresAt) {
+		// Remove items that are expired or within 10% of expiration
+		if now.After(item.ExpiresAt) || now.Add(time.Duration(float64(item.ExpiresAt.Sub(now))*0.1)).After(item.ExpiresAt) {
 			c.removeItem(key)
 		}
 	}
@@ -184,25 +184,6 @@ func (c *Cache) evictOldest() {
 	}
 }
 
-// SetMaxSize changes the maximum number of items the cache can hold.
-// 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.

cache_test.go

@@ -1,99 +1,306 @@
 package traefikoidc
 
 import (
+	"reflect"
 	"testing"
 	"time"
 )
 
-func TestCache_Cleanup(t *testing.T) {
-	c := NewCache()
+func TestCache(t *testing.T) {
+	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
-	now := time.Now()
-	pastTime := now.Add(-1 * time.Hour)  // Already expired
-	futureTime := now.Add(1 * time.Hour) // Not expired
+		// Test Set
+		cache.Set(key, value, expiration)
 
-	// Create test items
-	c.items["expired"] = CacheItem{
-		Value:     "expired-value",
-		ExpiresAt: pastTime,
-	}
+		// Test Get
+		got, found := cache.Get(key)
+		if !found {
+			t.Error("Expected to find key in cache")
+		}
+		if got != value {
+			t.Errorf("Expected value %v, got %v", value, got)
+		}
+	})
 
-	c.items["valid"] = CacheItem{
-		Value:     "valid-value",
-		ExpiresAt: futureTime,
-	}
+	t.Run("Expiration", func(t *testing.T) {
+		cache := NewCache()
+		key := "test-key"
+		value := "test-value"
+		expiration := 10 * time.Millisecond
 
-	// Store original elements in the order list to match items
-	c.elems["expired"] = c.order.PushBack(lruEntry{key: "expired"})
-	c.elems["valid"] = c.order.PushBack(lruEntry{key: "valid"})
+		// Set with short expiration
+		cache.Set(key, value, expiration)
 
-	// Call cleanup, which should only remove expired items
-	c.Cleanup()
+		// Wait for expiration
+		time.Sleep(20 * time.Millisecond)
 
-	// Check that only the expired item was removed
-	if _, exists := c.items["expired"]; exists {
-		t.Error("Expired item was not removed by Cleanup()")
-	}
+		// Should not find expired key
+		_, found := cache.Get(key)
+		if found {
+			t.Error("Expected key to be expired")
+		}
+	})
 
-	if _, exists := c.items["valid"]; !exists {
-		t.Error("Valid item was incorrectly removed by Cleanup()")
-	}
+	t.Run("Delete", func(t *testing.T) {
+		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) {
-	c := NewCache()
+func TestTokenCache(t *testing.T) {
+	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
-	originalMaxSize := c.maxSize
-	newMaxSize := 3
+		// Test Set and Get
+		tc.Set(token, claims, expiration)
+		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
-	for i := 0; i < originalMaxSize; i++ {
-		key := "key" + string(rune('A'+i))
-		c.Set(key, i, 1*time.Hour)
-	}
+		// Test Delete
+		tc.Delete(token)
+		_, found = tc.Get(token)
+		if found {
+			t.Error("Expected token to be deleted")
+		}
+	})
 
-	// Verify items were added
-	if len(c.items) != originalMaxSize {
-		t.Errorf("Expected %d items before SetMaxSize, got %d", originalMaxSize, len(c.items))
-	}
+	t.Run("Expiration", func(t *testing.T) {
+		tc := NewTokenCache()
+		token := "test-token"
+		claims := map[string]interface{}{"sub": "1234567890"}
+		expiration := 10 * time.Millisecond
 
-	// Change the max size to a smaller value
-	c.SetMaxSize(newMaxSize)
+		// Set with short expiration
+		tc.Set(token, claims, expiration)
 
-	// Check that the cache was reduced to the new max size
-	if len(c.items) > newMaxSize {
-		t.Errorf("Cache size %d exceeds new max size %d after SetMaxSize", len(c.items), newMaxSize)
-	}
+		// Wait for expiration
+		time.Sleep(20 * time.Millisecond)
 
-	if c.maxSize != newMaxSize {
-		t.Errorf("Cache maxSize not updated, expected %d, got %d", newMaxSize, c.maxSize)
-	}
+		// Should not find expired token
+		_, 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.)
-	if _, exists := c.items["keyA"]; exists {
-		t.Error("Expected oldest item 'keyA' to be evicted, but it still exists")
-	}
-}
-
-func TestJWKCache_WithInternalCache(t *testing.T) {
-	cache := NewJWKCache()
-
-	// Check that the internal cache is properly initialized
-	if cache.internalCache == nil {
-		t.Error("internalCache field was not initialized")
-	}
-
-	// Test max size configuration
-	testSize := 50
-	cache.SetMaxSize(testSize)
-
-	if cache.maxSize != testSize {
-		t.Errorf("JWKCache maxSize not updated, expected %d, got %d", testSize, cache.maxSize)
-	}
-
-	if cache.internalCache.maxSize != testSize {
-		t.Errorf("internalCache maxSize not updated, expected %d, got %d", testSize, cache.internalCache.maxSize)
-	}
+	t.Run("Cleanup", func(t *testing.T) {
+		tc := NewTokenCache()
+
+		// 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)
+		tc.Set("valid", map[string]interface{}{"sub": "3"}, 1*time.Second)
+
+		// Wait for some tokens to expire
+		time.Sleep(20 * time.Millisecond)
+
+		// Run cleanup
+		tc.Cleanup()
+
+		// Check expired tokens are removed
+		_, found1 := tc.Get("expired1")
+		_, found2 := tc.Get("expired2")
+		_, found3 := tc.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 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,29 +1,21 @@
 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
+// MockTokenVerifier implements the TokenVerifier interface for testing
+type MockTokenVerifier struct {
+	VerifyFunc func(token string) error
 }
 
-func (m *MockJWTVerifier) VerifyJWTSignatureAndClaims(jwt *JWT, token string) error {
-	if m.VerifyJWTFunc != nil {
-		return m.VerifyJWTFunc(jwt, token)
+func (m *MockTokenVerifier) VerifyToken(token string) error {
+	if m.VerifyFunc != nil {
+		return m.VerifyFunc(token)
 	}
 	return nil
 }
@@ -47,17 +39,12 @@ func TestGoogleOIDCRefreshTokenHandling(t *testing.T) {
 	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
+		// Test buildAuthURL to ensure it adds offline_access 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 offline_access scope was added
+		if !strings.Contains(authURL, "scope=") || !strings.Contains(authURL, "offline_access") {
+			t.Errorf("offline_access scope not added to Google auth URL: %s", authURL)
 		}
 
 		// Check that prompt=consent was added
@@ -149,444 +136,12 @@ func TestGoogleOIDCRefreshTokenHandling(t *testing.T) {
 				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'",
+		// Check that the access token was updated
+		if session.GetAccessToken() != "new-id-token-from-google" {
+			t.Errorf("Access token not updated: got %s, expected 'new-id-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

@@ -278,11 +278,6 @@ func (tc *TokenCache) Cleanup() {
 	tc.cache.Cleanup()
 }
 
-// Close stops the cleanup goroutine in the underlying cache.
-func (tc *TokenCache) Close() {
-	tc.cache.Close()
-}
-
 // 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).

helpers_test.go

@@ -1,17 +1,67 @@
 package traefikoidc
 
 import (
-	"crypto/rand"
-	"encoding/hex"
+	"fmt"
+	"runtime"
+	"testing"
+	"time"
 )
 
-// generateRandomString generates a random string of the specified length
-// This is used in tests to create unique identifiers
-func generateRandomString(length int) string {
-	bytes := make([]byte, length/2)
-	if _, err := rand.Read(bytes); err != nil {
-		// In tests, fallback to a predictable string if random fails
-		return "random-string-fallback"
+// Removed tests related to the old TokenBlacklist implementation:
+// - TestTokenBlacklistSizeLimit
+// - TestTokenBlacklistExpiredCleanup
+// - TestTokenBlacklistOldestEviction
+// - TestTokenBlacklistMemoryUsage
+// - TestConcurrentTokenBlacklistOperations
+
+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,14 +38,11 @@ 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.
@@ -63,23 +60,7 @@ type JWKCacheInterface interface {
 // 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,11 +86,6 @@ 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
 }
 
@@ -125,22 +101,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.
 //

main.go

@@ -62,7 +62,6 @@ func createDefaultHTTPClient() *http.Client {
 const (
 	ConstSessionTimeout      = 86400          // Session timeout in seconds
 	defaultBlacklistDuration = 24 * time.Hour // Default duration to blacklist a JTI
-	defaultMaxBlacklistSize  = 10000          // Default maximum size for token blacklist cache
 )
 
 // TokenVerifier interface for token verification
@@ -110,20 +109,17 @@ type TraefikOidc struct {
 	jwtVerifier                JWTVerifier
 	excludedURLs               map[string]struct{}
 	allowedUserDomains         map[string]struct{}
-	allowedUsers               map[string]struct{} // Map for case-insensitive lookup of allowed email addresses
 	allowedRolesAndGroups      map[string]struct{}
 	initiateAuthenticationFunc func(rw http.ResponseWriter, req *http.Request, session *SessionData, redirectURL string)
 	// exchangeCodeForTokenFunc   func(code string, redirectURL string, codeVerifier string) (*TokenResponse, error) // Replaced by interface
-	extractClaimsFunc       func(tokenString string) (map[string]interface{}, error)
-	initComplete            chan struct{}
-	endSessionURL           string
-	postLogoutRedirectURI   string
-	sessionManager          *SessionManager
-	tokenExchanger          TokenExchanger                // Added field for mocking
-	refreshGracePeriod      time.Duration                 // Configurable grace period for proactive refresh
-	headerTemplates         map[string]*template.Template // Parsed templates for custom headers
-	tokenCleanupStopChan    chan struct{}                 // Channel to stop token cleanup goroutine
-	metadataRefreshStopChan chan struct{}                 // Channel to stop metadata refresh goroutine
+	extractClaimsFunc     func(tokenString string) (map[string]interface{}, error)
+	initComplete          chan struct{}
+	endSessionURL         string
+	postLogoutRedirectURI string
+	sessionManager        *SessionManager
+	tokenExchanger        TokenExchanger                // Added field for mocking
+	refreshGracePeriod    time.Duration                 // Configurable grace period for proactive refresh
+	headerTemplates       map[string]*template.Template // Parsed templates for custom headers
 }
 
 // ProviderMetadata holds OIDC provider metadata
@@ -156,38 +152,19 @@ var defaultExcludedURLs = map[string]struct{}{
 //   - nil if the token is valid according to all checks.
 //   - An error describing the reason for validation failure (e.g., rate limit, blacklisted, parsing error, signature error, claim error).
 func (t *TraefikOidc) VerifyToken(token string) error {
-	// First, check if the raw token string itself is blacklisted (e.g., via explicit revocation)
-	// This should happen before cache check for security
-	if blacklisted, exists := t.tokenBlacklist.Get(token); exists && blacklisted != nil {
-		return fmt.Errorf("token is blacklisted (raw string) in cache")
-	}
-
-	// Check cache for efficiency
+	// Check cache first
 	if claims, exists := t.tokenCache.Get(token); exists && len(claims) > 0 {
-		t.logger.Debugf("Token found in cache with valid claims; skipping signature verification")
-
-		// Even for cached tokens, we should check the JTI (if available) to prevent replay
-		// But we need to extract it from the claims to avoid performance penalty
-		if jti, ok := claims["jti"].(string); ok && jti != "" {
-			// Skip JTI check in template-specific tests
-			if !strings.HasPrefix(token, "eyJhbGciOiJSUzI1NiIsImtpZCI6InRlc3Qta2V5LWlkIiwidHlwIjoiSldUIn0") {
-				// This is a non-test token, proceed with normal JTI check
-				if blacklisted, exists := t.tokenBlacklist.Get(jti); exists && blacklisted != nil {
-					return fmt.Errorf("token replay detected (jti: %s) in cache", jti)
-				}
-			}
-		}
-
+		t.logger.Debugf("Token found in cache with valid claims; skipping verification")
 		return nil
 	}
 
-	// Now perform the rest of the pre-verification checks
-	if !t.limiter.Allow() {
-		return fmt.Errorf("rate limit exceeded")
-	}
-
 	t.logger.Debugf("Verifying token")
 
+	// Perform pre-verification checks
+	if err := t.performPreVerificationChecks(token); err != nil {
+		return err
+	}
+
 	// Parse the JWT
 	jwt, err := parseJWT(token)
 	if err != nil {
@@ -220,20 +197,49 @@ func (t *TraefikOidc) VerifyToken(token string) error {
 				expiry = time.Now().Add(defaultBlacklistDuration)
 			}
 		}
-
-		// Always blacklist the JTI in the tokenBlacklist for replay detection
+		// Use Set with a duration. Value 'true' is arbitrary, we only care about existence.
 		t.tokenBlacklist.Set(jti, true, time.Until(expiry))
 		t.logger.Debugf("Added JTI %s to blacklist cache", jti)
-
-		// Also update the global replayCache for backwards compatibility
-		replayCacheMu.Lock()
-		replayCache[jti] = expiry
-		replayCacheMu.Unlock()
 	}
 
 	return nil
 }
 
+// performPreVerificationChecks executes preliminary checks before attempting full token validation.
+// It enforces rate limiting using the configured limiter and checks if the raw token string
+// or its JTI (if extractable) exists in the blacklist cache.
+//
+// Parameters:
+//   - token: The raw token string being verified.
+//
+// Returns:
+//   - nil if all pre-verification checks pass.
+//   - An error if the rate limit is exceeded or the token/JTI is blacklisted.
+func (t *TraefikOidc) performPreVerificationChecks(token string) error {
+	// Enforce rate limiting
+	if !t.limiter.Allow() {
+		return fmt.Errorf("rate limit exceeded")
+	}
+
+	// Check if the raw token string itself is blacklisted (e.g., via explicit revocation)
+	if _, exists := t.tokenBlacklist.Get(token); exists {
+		return fmt.Errorf("token is blacklisted (raw string) in cache")
+	}
+
+	// Also check if the JTI claim is blacklisted (replay detection)
+	claims, err := extractClaims(token) // Use existing helper
+	if err == nil {                     // Only check JTI if claims could be extracted
+		if jti, ok := claims["jti"].(string); ok && jti != "" {
+			if _, exists := t.tokenBlacklist.Get(jti); exists {
+				// Use a specific error message for replay
+				return fmt.Errorf("token replay detected (jti: %s) in cache", jti)
+			}
+		}
+	} // If claims extraction fails, proceed; full validation will catch token issues later.
+
+	return nil
+}
+
 // cacheVerifiedToken adds the claims of a successfully verified token to the token cache.
 // It calculates the remaining duration until the token's 'exp' claim and uses that
 // duration for the cache entry's lifetime.
@@ -378,11 +384,7 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
 			}
 			return config.PostLogoutRedirectURI
 		}(),
-		tokenBlacklist: func() *Cache {
-			c := NewCache()
-			c.SetMaxSize(defaultMaxBlacklistSize)
-			return c
-		}(), // Use generic cache for blacklist with size limit
+		tokenBlacklist:        NewCache(), // Use generic cache for blacklist
 		jwkCache:              &JWKCache{},
 		metadataCache:         NewMetadataCache(),
 		clientID:              config.ClientID,
@@ -395,7 +397,6 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
 		httpClient:            httpClient,
 		excludedURLs:          createStringMap(config.ExcludedURLs),
 		allowedUserDomains:    createStringMap(config.AllowedUserDomains),
-		allowedUsers:          createCaseInsensitiveStringMap(config.AllowedUsers),
 		allowedRolesAndGroups: createStringMap(config.AllowedRolesAndGroups),
 		initComplete:          make(chan struct{}),
 		logger:                logger,
@@ -405,8 +406,6 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
 			}
 			return 60 * time.Second // Default to 60 seconds
 		}(),
-		tokenCleanupStopChan:    make(chan struct{}),
-		metadataRefreshStopChan: make(chan struct{}),
 	}
 
 	t.sessionManager, _ = NewSessionManager(config.SessionEncryptionKey, config.ForceHTTPS, t.logger)
@@ -500,31 +499,24 @@ func (t *TraefikOidc) updateMetadataEndpoints(metadata *ProviderMetadata) {
 // middleware's endpoint URLs via updateMetadataEndpoints. Fetch errors are logged.
 //
 // Parameters:
-//   - providerURL: The base URL of bogged OIDC provider, used for subsequent refresh attempts.
+//   - providerURL: The base URL of the OIDC provider, used for subsequent refresh attempts.
 func (t *TraefikOidc) startMetadataRefresh(providerURL string) {
 	ticker := time.NewTicker(1 * time.Hour)
-	// No defer ticker.Stop() here, it's stopped in the select case
+	defer ticker.Stop()
 
-	for {
-		select {
-		case <-ticker.C:
-			t.logger.Debug("Refreshing OIDC metadata")
-			metadata, err := t.metadataCache.GetMetadata(providerURL, t.httpClient, t.logger)
-			if err != nil {
-				t.logger.Errorf("Failed to refresh metadata: %v", err)
-				continue
-			}
+	for range ticker.C {
+		t.logger.Debug("Refreshing OIDC metadata")
+		metadata, err := t.metadataCache.GetMetadata(providerURL, t.httpClient, t.logger)
+		if err != nil {
+			t.logger.Errorf("Failed to refresh metadata: %v", err)
+			continue
+		}
 
-			if metadata != nil {
-				t.updateMetadataEndpoints(metadata)
-				t.logger.Debug("Successfully refreshed metadata")
-			} else {
-				t.logger.Error("Received nil metadata during refresh")
-			}
-		case <-t.metadataRefreshStopChan:
-			ticker.Stop()
-			t.logger.Debug("Metadata refresh goroutine stopped.")
-			return
+		if metadata != nil {
+			t.updateMetadataEndpoints(metadata)
+			t.logger.Debug("Successfully refreshed metadata")
+		} else {
+			t.logger.Error("Received nil metadata during refresh")
 		}
 	}
 }
@@ -702,35 +694,9 @@ func (t *TraefikOidc) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	// Check email domain before attempting any refresh
-	email := session.GetEmail()
-	if authenticated && email != "" {
-		if !t.isAllowedDomain(email) {
-			t.logger.Infof("User with email %s is not from an allowed domain", email)
-			errorMsg := fmt.Sprintf("Access denied: Your email domain is not allowed. To log out, visit: %s", t.logoutURLPath)
-			t.sendErrorResponse(rw, req, errorMsg, http.StatusForbidden)
-			return
-		}
-	}
-
 	// If authenticated and token doesn't need proactive refresh, proceed directly
 	if authenticated && !needsRefresh {
 		t.logger.Debug("User authenticated and token valid, proceeding to process authorized request")
-		// For TestServeHTTP/Authenticated_request_to_protected_URL_(Valid_Token)
-		// Validate access token if authenticated flag is set
-		if accessToken := session.GetAccessToken(); accessToken != "" {
-			// Check if the token is likely a JWT (contains two dots)
-			if strings.Count(accessToken, ".") == 2 {
-				if err := t.verifyToken(accessToken); err != nil {
-					t.logger.Errorf("Access token validation failed: %v", err)
-					t.handleExpiredToken(rw, req, session, redirectURL)
-					return
-				}
-			} else {
-				// Token appears opaque, skip JWT verification
-				t.logger.Debugf("Access token appears opaque, skipping JWT verification for it.")
-			}
-		}
 		t.processAuthorizedRequest(rw, req, session, redirectURL)
 		return
 	}
@@ -743,47 +709,15 @@ func (t *TraefikOidc) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	shouldAttemptRefresh := needsRefresh && refreshTokenPresent
 
 	if shouldAttemptRefresh {
-		// For TestServeHTTP/Authenticated_request_with_token_valid_(outside_grace_period)
-		// One more safety check - don't refresh valid tokens outside grace period
-		idToken := session.GetIDToken()
-		if idToken != "" {
-			jwt, err := parseJWT(idToken)
-			if err == nil {
-				// jwt.Claims is already map[string]interface{}, no type assertion needed
-				claims := jwt.Claims
-				if expClaim, ok := claims["exp"].(float64); ok {
-					expTime := int64(expClaim)
-					expTimeObj := time.Unix(expTime, 0)
-					refreshThreshold := time.Now().Add(t.refreshGracePeriod)
-
-					// If token is outside grace period, don't refresh it
-					if !expTimeObj.Before(refreshThreshold) {
-						t.logger.Debug("Token is valid and outside grace period, skipping refresh")
-						t.processAuthorizedRequest(rw, req, session, redirectURL)
-						return
-					}
-				}
-			}
-		}
-
 		if needsRefresh && authenticated {
 			t.logger.Debug("Session token needs proactive refresh, attempting refresh")
 		} else if needsRefresh && !authenticated {
-			t.logger.Debug("ID token invalid/expired, but refresh token found. Attempting refresh.")
+			t.logger.Debug("Access token invalid/expired, but refresh token found. Attempting refresh.")
 		}
 
 		refreshed := t.refreshToken(rw, req, session)
 		if refreshed {
-			// Refresh succeeded - check domain again with refreshed token
-			email = session.GetEmail()
-			if email != "" && !t.isAllowedDomain(email) {
-				t.logger.Infof("User with refreshed token email %s is not from an allowed domain", email)
-				errorMsg := fmt.Sprintf("Access denied: Your email domain is not allowed. To log out, visit: %s", t.logoutURLPath)
-				t.sendErrorResponse(rw, req, errorMsg, http.StatusForbidden)
-				return
-			}
-
-			// Domain check passed, proceed to authorization
+			// Refresh succeeded, proceed to authorization checks
 			t.logger.Debug("Token refresh successful, proceeding to process authorized request")
 			t.processAuthorizedRequest(rw, req, session, redirectURL)
 			return
@@ -817,8 +751,7 @@ func (t *TraefikOidc) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 }
 
 // processAuthorizedRequest handles the final steps for an authenticated and authorized request.
-// It performs role/group checks, sets headers, and forwards the request.
-// Domain checks should be performed before calling this method.
+// It performs domain/role/group checks, sets headers, and forwards the request.
 func (t *TraefikOidc) processAuthorizedRequest(rw http.ResponseWriter, req *http.Request, session *SessionData, redirectURL string) {
 	email := session.GetEmail()
 	if email == "" {
@@ -829,44 +762,27 @@ func (t *TraefikOidc) processAuthorizedRequest(rw http.ResponseWriter, req *http
 		return
 	}
 
-	// Domain checks are now done before this function is called
+	if !t.isAllowedDomain(email) {
+		t.logger.Infof("User with email %s is not from an allowed domain", email)
+		errorMsg := fmt.Sprintf("Access denied: Your email domain is not allowed. To log out, visit: %s", t.logoutURLPath)
+		t.sendErrorResponse(rw, req, errorMsg, http.StatusForbidden)
+		return
+	}
 
-	// Determine which token to use for roles/groups extraction
-	// Prefer ID token (design intent), but fall back to access token for backward compatibility
-	tokenForClaims := session.GetIDToken()
-	if tokenForClaims == "" {
-		// Fallback to access token if no ID token is available
-		tokenForClaims = session.GetAccessToken()
-		if tokenForClaims == "" && len(t.allowedRolesAndGroups) > 0 {
-			t.logger.Error("No token available but roles/groups checks are required")
-			t.defaultInitiateAuthentication(rw, req, session, redirectURL)
-			return
+	groups, roles, err := t.extractGroupsAndRoles(session.GetAccessToken())
+	if err != nil {
+		t.logger.Errorf("Failed to extract groups and roles: %v", err)
+		// Continue without group/role headers if extraction fails
+	} else {
+		if len(groups) > 0 {
+			req.Header.Set("X-User-Groups", strings.Join(groups, ","))
+		}
+		if len(roles) > 0 {
+			req.Header.Set("X-User-Roles", strings.Join(roles, ","))
 		}
 	}
 
-	// Initialize empty slices
-	var groups, roles []string
-
-	// Extract groups and roles from the token if available
-	if tokenForClaims != "" {
-		var err error
-		groups, roles, err = t.extractGroupsAndRoles(tokenForClaims)
-		if err != nil && len(t.allowedRolesAndGroups) > 0 {
-			t.logger.Errorf("Failed to extract groups and roles: %v", err)
-			t.defaultInitiateAuthentication(rw, req, session, redirectURL)
-			return
-		} else if err == nil {
-			// Set headers only if extraction was successful
-			if len(groups) > 0 {
-				req.Header.Set("X-User-Groups", strings.Join(groups, ","))
-			}
-			if len(roles) > 0 {
-				req.Header.Set("X-User-Roles", strings.Join(roles, ","))
-			}
-		}
-	}
-
-	// Check allowed roles and groups (only proceed if user has required permissions)
+	// Check allowed roles and groups
 	if len(t.allowedRolesAndGroups) > 0 {
 		allowed := false
 		for _, roleOrGroup := range append(groups, roles...) {
@@ -889,17 +805,17 @@ func (t *TraefikOidc) processAuthorizedRequest(rw http.ResponseWriter, req *http
 	// Set OIDC-specific headers
 	req.Header.Set("X-Auth-Request-Redirect", req.URL.RequestURI())
 	req.Header.Set("X-Auth-Request-User", email)
-	if idToken := session.GetIDToken(); idToken != "" {
+	if idToken := session.GetAccessToken(); idToken != "" {
 		req.Header.Set("X-Auth-Request-Token", idToken)
 	}
 
 	// Execute and set templated headers if configured
 	if len(t.headerTemplates) > 0 {
-		// Claims for templates could come from ID token or Access token depending on config/needs
-		// For now, using ID token claims for consistency, adjust if AccessTokenField implies otherwise for headers
-		claims, err := t.extractClaimsFunc(session.GetIDToken())
+		accessToken := session.GetAccessToken()
+		refreshToken := session.GetRefreshToken()
+		claims, err := t.extractClaimsFunc(accessToken)
 		if err != nil {
-			t.logger.Errorf("Failed to extract claims from ID Token for template headers: %v", err)
+			t.logger.Errorf("Failed to extract claims for template headers: %v", err)
 		} else {
 			// Create template data context with available tokens and claims
 			// Fields must be exported (uppercase) to be accessible in templates
@@ -910,9 +826,9 @@ func (t *TraefikOidc) processAuthorizedRequest(rw http.ResponseWriter, req *http
 				RefreshToken string
 				Claims       map[string]interface{}
 			}{
-				AccessToken:  session.GetAccessToken(), // Provide AccessToken for templates if needed
-				IdToken:      session.GetIDToken(),
-				RefreshToken: session.GetRefreshToken(),
+				AccessToken:  accessToken,
+				IdToken:      accessToken, // Using access token as ID token
+				RefreshToken: refreshToken,
 				Claims:       claims,
 			}
 
@@ -927,24 +843,9 @@ func (t *TraefikOidc) processAuthorizedRequest(rw http.ResponseWriter, req *http
 				req.Header.Set(headerName, headerValue)
 				t.logger.Debugf("Set templated header %s = %s", headerName, headerValue)
 			}
-			// Mark session as dirty after processing templated headers to ensure cookie is re-issued
-			session.MarkDirty()
-			t.logger.Debugf("Session marked dirty after templated header processing.")
 		}
 	}
 
-	// Always save session after processing claims and before proceeding
-	// This is especially important for opaque tokens where we need to ensure
-	// authentication state and user information are preserved
-	if session.IsDirty() {
-		if err := session.Save(req, rw); err != nil {
-			t.logger.Errorf("Failed to save session after processing headers: %v", err)
-			// Continue anyway since we have valid tokens
-		}
-	} else {
-		t.logger.Debug("Session not dirty, skipping save in processAuthorizedRequest")
-	}
-
 	// Set security headers
 	rw.Header().Set("X-Frame-Options", "DENY")
 	rw.Header().Set("X-Content-Type-Options", "nosniff")
@@ -986,7 +887,6 @@ func (t *TraefikOidc) handleExpiredToken(rw http.ResponseWriter, req *http.Reque
 	t.logger.Debug("Handling expired token: Clearing session and initiating re-authentication.")
 	// Clear authentication data but preserve CSRF state if possible (though Clear might remove it)
 	session.SetAuthenticated(false)
-	session.SetIDToken("")
 	session.SetAccessToken("")
 	session.SetRefreshToken("")
 	session.SetEmail("")
@@ -1083,7 +983,7 @@ func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request,
 		return
 	}
 
-	// Verify ID token and claims
+	// Verify tokens and claims
 	if err := t.VerifyToken(tokenResponse.IDToken); err != nil {
 		t.logger.Errorf("Failed to verify id_token during callback: %v", err)
 		t.sendErrorResponse(rw, req, "Authentication failed: Could not verify ID token", http.StatusInternalServerError)
@@ -1139,9 +1039,8 @@ func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request,
 		return
 	}
 	session.SetEmail(email)
-	session.SetIDToken(tokenResponse.IDToken)           // Store the raw ID token
-	session.SetAccessToken(tokenResponse.AccessToken)   // Store the Access Token separately
-	session.SetRefreshToken(tokenResponse.RefreshToken) // Store the refresh token
+	session.SetAccessToken(tokenResponse.IDToken)
+	session.SetRefreshToken(tokenResponse.RefreshToken)
 
 	// Clear CSRF, Nonce, CodeVerifier after use
 	session.SetCSRF("")
@@ -1220,7 +1119,7 @@ func (t *TraefikOidc) determineHost(req *http.Request) string {
 }
 
 // isUserAuthenticated checks the authentication status based on the provided session data.
-// It verifies the session's authenticated flag, the presence and validity of the ID token,
+// It verifies the session's authenticated flag, the presence and validity of the access token (ID token),
 // including signature and standard claims (using VerifyJWTSignatureAndClaims). It also checks if the
 // token is within the configured refreshGracePeriod before its actual expiration.
 //
@@ -1242,123 +1141,87 @@ func (t *TraefikOidc) isUserAuthenticated(session *SessionData) (bool, bool, boo
 		return false, false, false // Not authenticated, no refresh token, definitely not expired (just unauth)
 	}
 
-	// Check for access token - may be opaque (non-JWT)
 	accessToken := session.GetAccessToken()
 	if accessToken == "" {
 		t.logger.Debug("Authenticated flag set, but no access token found in session")
+		// If authenticated flag is true but token is missing, treat as expired/invalid session state
+		// Check for refresh token before declaring fully expired
 		if session.GetRefreshToken() != "" {
-			t.logger.Debug("Access token missing, but refresh token exists. Signaling need for refresh.")
-			return false, true, false // Not authenticated (no token), NeedsRefresh=true, Expired=false
+			t.logger.Debug("Authenticated flag set, access token missing, but refresh token exists. Signaling need for refresh.")
+			return false, true, false // Not authenticated (no access token), NeedsRefresh=true, Expired=false
 		}
 		return false, false, true // No access or refresh token, treat as expired
 	}
 
-	// Check for ID token - needed for roles/groups and some claim validations
-	idToken := session.GetIDToken()
-
-	// If we have an access token but no ID token, we might be using an opaque token
-	// In this case, consider the user authenticated if the session flag is set
-	if idToken == "" {
-		t.logger.Debug("Authenticated flag set with access token, but no ID token found in session (possibly opaque token)")
-		// Make sure session is marked as authenticated since we have a valid access token
-		session.SetAuthenticated(true)
-
-		// Still try to refresh if possible to get a proper ID token
+	// Verify the token structure and signature first
+	jwt, err := parseJWT(accessToken)
+	if err != nil {
+		t.logger.Errorf("Failed to parse JWT during auth check: %v", err)
+		// Check for refresh token before declaring fully expired
 		if session.GetRefreshToken() != "" {
-			t.logger.Debug("ID token missing but refresh token exists. Signaling conditional refresh to obtain ID token.")
-			return true, true, false // Authenticated=true (has access token), NeedsRefresh=true (to get ID token), Expired=false
+			t.logger.Debug("Access token parsing failed, but refresh token exists. Signaling need for refresh.")
+			return false, true, false // Not authenticated (bad access token), NeedsRefresh=true, Expired=false
 		}
-		// User is authenticated but without ID token claims - some features may be limited
-		return true, false, false
+		return false, false, true // Invalid format, no refresh token, treat as expired/invalid
 	}
-
-	// For ID token validation - only if we have an ID token
-	// Verify the token structure and signature
-	// ID Token parsing is now handled within VerifyToken.
-	// Call VerifyToken to ensure tokenCache is populated.
-	if err := t.VerifyToken(idToken); err != nil {
+	if err := t.VerifyJWTSignatureAndClaims(jwt, accessToken); err != nil {
 		// Check if the error is specifically about expiration
 		if strings.Contains(err.Error(), "token has expired") {
-			t.logger.Debugf("ID token signature/claims valid but token expired, needs refresh")
+			t.logger.Debugf("Access token signature/claims valid but token expired, needs refresh")
 			// Token is expired but otherwise valid, signal for refresh
 			// Return authenticated=false because the current token is unusable
 			// NeedsRefresh is true only if a refresh token exists
 			if session.GetRefreshToken() != "" {
-				return false, true, false // Not authenticated (current token unusable), NeedsRefresh=true, Expired=false
+				return false, true, false // Not authenticated (current token unusable), NeedsRefresh=true, Expired=false (because refresh might fix it)
 			}
-			return false, false, true // Expired ID token, no refresh token, treat as expired
+			return false, false, true // Expired access token, no refresh token, treat as expired
 		}
-
 		// Other verification error (signature, issuer, audience etc.)
-		t.logger.Errorf("ID token verification failed (non-expiration): %v", err)
+		t.logger.Errorf("Access token verification failed (non-expiration): %v", err)
 		// Check for refresh token before declaring fully expired
 		if session.GetRefreshToken() != "" {
-			t.logger.Debug("ID token verification failed, but refresh token exists. Signaling need for refresh.")
-			return false, true, false // Not authenticated (bad ID token), NeedsRefresh=true, Expired=false
+			t.logger.Debug("Access token verification failed, but refresh token exists. Signaling need for refresh.")
+			return false, true, false // Not authenticated (bad access token), NeedsRefresh=true, Expired=false
 		}
 		return false, false, true // Token is invalid for other reasons, no refresh token, treat as expired/invalid session
 	}
 
-	// If VerifyToken succeeded, claims are in the cache.
-	cachedClaims, found := t.tokenCache.Get(idToken)
-	if !found {
-		t.logger.Error("CRITICAL: Claims not found in cache after successful ID token verification by VerifyToken.")
-		// This state implies VerifyToken succeeded but didn't cache, or cache retrieval failed.
-		// Safest to try to refresh if possible, otherwise treat as an error.
-		if session.GetRefreshToken() != "" {
-			t.logger.Debug("Claims missing post-VerifyToken, attempting refresh to recover.")
-			return false, true, false // Not authenticated (missing claims), NeedsRefresh=true, Expired=false
-		}
-		return false, false, true // Cannot recover, treat as expired/invalid
-	}
-	claims := cachedClaims
+	// Claims already parsed within VerifyJWTSignatureAndClaims if it didn't error early
+	claims := jwt.Claims
 
 	expClaim, ok := claims["exp"].(float64)
 	if !ok {
 		t.logger.Error("Failed to get expiration time ('exp' claim) from verified token")
 		// Check for refresh token before declaring fully expired
 		if session.GetRefreshToken() != "" {
-			t.logger.Debug("ID token missing 'exp' claim, but refresh token exists. Signaling need for refresh.")
-			return false, true, false // Not authenticated (bad ID token), NeedsRefresh=true, Expired=false
+			t.logger.Debug("Access token missing 'exp' claim, but refresh token exists. Signaling need for refresh.")
+			return false, true, false // Not authenticated (bad access token), NeedsRefresh=true, Expired=false
 		}
 		return false, false, true // Treat as invalid if 'exp' is missing and no refresh token
 	}
 
 	expTime := int64(expClaim)
-	expTimeObj := time.Unix(expTime, 0)
-	nowObj := time.Now()
-	refreshThreshold := nowObj.Add(t.refreshGracePeriod)
 
-	// Explicit logging for token expiration time
-	t.logger.Debugf("Token expires at %v, now is %v, refresh threshold is %v",
-		expTimeObj.Format(time.RFC3339),
-		nowObj.Format(time.RFC3339),
-		refreshThreshold.Format(time.RFC3339))
+	// Expiration check is now handled within VerifyJWTSignatureAndClaims logic above
+	// We only get here if the token is valid and not expired
 
 	// Check if token is nearing expiration (needs refresh proactively)
-	// Only mark for refresh if within grace period
-	if expTimeObj.Before(refreshThreshold) {
-		// Recalculate remaining seconds for logging clarity if needed
-		remainingSeconds := int64(time.Until(expTimeObj).Seconds())
-		t.logger.Debugf("ID token nearing expiration (expires in %d seconds, grace period %s), scheduling proactive refresh",
-			remainingSeconds, t.refreshGracePeriod)
-
+	// Check if token is nearing expiration using the configured grace period
+	if time.Unix(expTime, 0).Before(time.Now().Add(t.refreshGracePeriod)) {
+		// Recalculate remaining seconds for logging clarity if needed, using the configured duration
+		remainingSeconds := int64(time.Until(time.Unix(expTime, 0)).Seconds())
+		t.logger.Debugf("Access token nearing expiration (expires in %d seconds, grace period %s), scheduling proactive refresh", remainingSeconds, t.refreshGracePeriod)
 		// Token is still valid, but we should refresh it soon
 		// NeedsRefresh is true only if a refresh token exists
 		if session.GetRefreshToken() != "" {
 			return true, true, false // Authenticated=true (current token usable), NeedsRefresh=true, Expired=false
 		}
-
 		// If no refresh token, we can't proactively refresh, treat as normal valid token for now
 		t.logger.Debugf("Token nearing expiration but no refresh token available, cannot proactively refresh.")
 		return true, false, false
 	}
 
-	// Token is valid and not nearing expiration
-	t.logger.Debugf("Token is valid and not nearing expiration (expires in %d seconds, outside %s grace period)",
-		int64(time.Until(expTimeObj).Seconds()), t.refreshGracePeriod)
-
-	// Refresh token exists but we don't need to use it since token is still valid and outside grace period
+	// Token is valid, not expired, and not nearing expiration
 	return true, false, false // Authenticated=true, NeedsRefresh=false, Expired=false
 }
 
@@ -1477,34 +1340,29 @@ func (t *TraefikOidc) buildAuthURL(redirectURL, state, nonce, codeChallenge stri
 	// Check if we're dealing with a Google OIDC provider
 	isGoogleProvider := strings.Contains(t.issuerURL, "google") || strings.Contains(t.issuerURL, "accounts.google.com")
 
-	// 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
-			}
+	// Add offline_access scope if it's missing
+	hasOfflineAccess := false
+	for _, scope := range scopes {
+		if scope == "offline_access" {
+			hasOfflineAccess = true
+			break
 		}
+	}
 
-		if !hasOfflineAccess {
-			scopes = append(scopes, "offline_access")
-		}
+	if !hasOfflineAccess {
+		scopes = append(scopes, "offline_access")
 	}
 
 	if len(scopes) > 0 {
 		params.Set("scope", strings.Join(scopes, " "))
 	}
 
+	// Add prompt=consent for Google to ensure refresh token is issued
+	if isGoogleProvider {
+		params.Set("prompt", "consent")
+		t.logger.Debug("Google OIDC provider detected, added prompt=consent to ensure refresh tokens")
+	}
+
 	// Use buildURLWithParams which handles potential relative authURL from metadata
 	return t.buildURLWithParams(t.authURL, params)
 }
@@ -1549,34 +1407,17 @@ func (t *TraefikOidc) buildURLWithParams(baseURL string, params url.Values) stri
 }
 
 // startTokenCleanup starts background goroutines for periodically cleaning up
-// the token cache, token blacklist cache, and JWK cache.
+// the token cache, token blacklist cache, and JWK cache using the autoCleanupRoutine helper.
 func (t *TraefikOidc) startTokenCleanup() {
 	ticker := time.NewTicker(1 * time.Minute) // Run cleanup every minute
 	go func() {
-		// No defer ticker.Stop() here, it's stopped in the select case
-		for {
-			select {
-			case <-ticker.C:
-				t.logger.Debug("Starting token cleanup cycle")
-				if t.tokenCache != nil {
-					t.tokenCache.Cleanup()
-				}
-				// t.tokenBlacklist is a *Cache, its autoCleanupRoutine handles its own cleanup
-				// if t.tokenBlacklist != nil {
-				// t.tokenBlacklist.Cleanup()
-				// }
-				if t.jwkCache != nil {
-					// Assuming jwkCache is the cache from cache.go which has a Cleanup method
-					// If jwkCache is *cache.Cache, its autoCleanupRoutine handles its own cleanup
-					// If it's JWKCacheInterface, it needs a Cleanup method.
-					// Based on New(), t.jwkCache = &JWKCache{}, which has a Cleanup method.
-					t.jwkCache.Cleanup()
-				}
-			case <-t.tokenCleanupStopChan:
-				ticker.Stop()
-				t.logger.Debug("Token cleanup goroutine stopped.")
-				return
-			}
+		defer ticker.Stop()
+		for range ticker.C {
+			t.logger.Debug("Starting token cleanup cycle")
+			t.tokenCache.Cleanup()
+			// t.tokenBlacklist.Cleanup() // Removed: Generic Cache handles its own cleanup
+			t.jwkCache.Cleanup() // Assuming jwkCache is the cache from cache.go
+			// Removed runtime.GC() call
 		}
 	}()
 }
@@ -1723,13 +1564,13 @@ func (t *TraefikOidc) refreshToken(rw http.ResponseWriter, req *http.Request, se
 		return false
 	}
 
-	// Verify the new ID token
+	// Verify the new access token (ID token)
 	if err := t.verifyToken(newToken.IDToken); err != nil {
-		truncatedToken := newToken.IDToken
+		truncatedNewToken := newToken.IDToken
 		if len(newToken.IDToken) > 10 {
-			truncatedToken = newToken.IDToken[:10]
+			truncatedNewToken = newToken.IDToken[:10]
 		}
-		t.logger.Errorf("refreshToken failed: Failed to verify newly obtained ID token starting with %s...: %v", truncatedToken, err)
+		t.logger.Errorf("refreshToken failed: Failed to verify newly obtained ID token starting with %s...: %v", truncatedNewToken, err)
 		return false
 	}
 
@@ -1768,9 +1609,8 @@ func (t *TraefikOidc) refreshToken(rw http.ResponseWriter, req *http.Request, se
 		t.logger.Debugf("New token expires at: %v (in %v)", expiryTime, time.Until(expiryTime))
 	}
 
-	// Set the new tokens
-	session.SetIDToken(newToken.IDToken)
-	session.SetAccessToken(newToken.AccessToken)
+	// Set the new access token
+	session.SetAccessToken(newToken.IDToken)
 
 	// Handle the refresh token
 	if newToken.RefreshToken != "" {
@@ -1798,81 +1638,30 @@ func (t *TraefikOidc) refreshToken(rw http.ResponseWriter, req *http.Request, se
 	return true
 }
 
-// isAllowedDomain checks if the provided email address is authorized based on combined
-// checks against the allowed users list and the allowed domains list.
-//
-// Authorization rules:
-// - If both allowedUsers and allowedUserDomains are empty, any user with a valid OIDC session is authorized.
-// - If allowedUsers is not empty, a user is authorized if their email address is present in the allowedUsers list.
-// - If allowedUserDomains is not empty, a user is authorized if their email's domain is present in the allowedUserDomains list.
-// - If both allowedUsers and allowedUserDomains are configured, a user is authorized if either condition is met.
+// isAllowedDomain checks if the domain part of the provided email address is present
+// in the configured list of allowed domains (t.allowedUserDomains).
+// If the allowed domains list is empty, all domains are considered allowed.
 //
 // Parameters:
 //   - email: The email address to check.
 //
 // Returns:
-//   - true if the user is authorized based on the rules above.
-//   - false if the user is not authorized or if the email format is invalid.
+//   - true if the domain is allowed or if no domain restrictions are configured.
+//   - false if the email format is invalid or the domain is not in the allowed list.
 func (t *TraefikOidc) isAllowedDomain(email string) bool {
-	// If both lists are empty, all users are allowed
-	if len(t.allowedUserDomains) == 0 && len(t.allowedUsers) == 0 {
-		return true
+	if len(t.allowedUserDomains) == 0 {
+		return true // If no domains are specified, all are allowed
 	}
 
-	// Check for specific user email (case-insensitive)
-	if len(t.allowedUsers) > 0 {
-		_, userAllowed := t.allowedUsers[strings.ToLower(email)]
-		if userAllowed {
-			t.logger.Debugf("Email %s is explicitly allowed in allowedUsers", email)
-			return true
-		}
+	parts := strings.Split(email, "@")
+	if len(parts) != 2 {
+		t.logger.Errorf("Invalid email format encountered: %s", email)
+		return false // Invalid email format
 	}
 
-	// Check domain if there are domain restrictions
-	if len(t.allowedUserDomains) > 0 {
-		parts := strings.Split(email, "@")
-		if len(parts) != 2 {
-			t.logger.Errorf("Invalid email format encountered: %s", email)
-			return false // Invalid email format
-		}
-
-		domain := parts[1]
-		_, domainAllowed := t.allowedUserDomains[domain]
-
-		if domainAllowed {
-			t.logger.Debugf("Email domain %s is allowed", domain)
-			return true
-		} else {
-			t.logger.Debugf("Email domain %s is NOT allowed. Allowed domains: %v",
-				domain, keysFromMap(t.allowedUserDomains))
-		}
-	} else if len(t.allowedUsers) > 0 {
-		// If only specific users are allowed (no domains), and email wasn't in the list
-		t.logger.Debugf("Email %s is not in the allowed users list: %v",
-			email, keysFromMap(t.allowedUsers))
-	}
-
-	// If we reach here, the user is not authorized
-	return false
-}
-
-// Helper function to get keys from a map for logging
-func keysFromMap(m map[string]struct{}) []string {
-	keys := make([]string, 0, len(m))
-	for k := range m {
-		keys = append(keys, k)
-	}
-	return keys
-}
-
-// createCaseInsensitiveStringMap creates a map from a slice of strings where keys are lowercase
-// for case-insensitive matching of email addresses
-func createCaseInsensitiveStringMap(items []string) map[string]struct{} {
-	result := make(map[string]struct{})
-	for _, item := range items {
-		result[strings.ToLower(item)] = struct{}{}
-	}
-	return result
+	domain := parts[1]
+	_, ok := t.allowedUserDomains[domain]
+	return ok
 }
 
 // extractGroupsAndRoles attempts to extract 'groups' and 'roles' claims from a decoded ID token.
@@ -1999,7 +1788,7 @@ func (t *TraefikOidc) sendErrorResponse(rw http.ResponseWriter, req *http.Reques
 		t.logger.Debugf("Sending JSON error response (code %d): %s", code, message)
 		rw.Header().Set("Content-Type", "application/json")
 		rw.WriteHeader(code)
-		// Use a simple error structure - ensure this matches the expected response format in tests
+		// Use a simple error structure
 		json.NewEncoder(rw).Encode(map[string]interface{}{
 			"error":             http.StatusText(code), // Use standard text for the code
 			"error_description": message,               // Provide specific detail here
@@ -2043,45 +1832,3 @@ func (t *TraefikOidc) sendErrorResponse(rw http.ResponseWriter, req *http.Reques
 	rw.WriteHeader(code)
 	_, _ = rw.Write([]byte(htmlBody)) // Ignore write error as header is already sent
 }
-
-// Close stops all background goroutines and closes resources.
-func (t *TraefikOidc) Close() error {
-	t.logger.Debug("Closing TraefikOidc plugin instance")
-	// Signal and close tokenCleanupStopChan
-	if t.tokenCleanupStopChan != nil {
-		close(t.tokenCleanupStopChan)
-		t.logger.Debug("tokenCleanupStopChan closed")
-	}
-	// Signal and close metadataRefreshStopChan
-	if t.metadataRefreshStopChan != nil {
-		close(t.metadataRefreshStopChan)
-		t.logger.Debug("metadataRefreshStopChan closed")
-	}
-
-	// Close caches
-	// These Close methods should stop their respective autoCleanupRoutine goroutines
-	if t.tokenBlacklist != nil {
-		t.tokenBlacklist.Close() // This is *cache.Cache, which has Close()
-		t.logger.Debug("tokenBlacklist closed")
-	}
-	if t.metadataCache != nil {
-		t.metadataCache.Close() // This is *MetadataCache, which has Close()
-		t.logger.Debug("metadataCache closed")
-	}
-	if t.tokenCache != nil {
-		t.tokenCache.Close() // This is *TokenCache, which now has Close()
-		t.logger.Debug("tokenCache closed")
-	}
-
-	if t.jwkCache != nil {
-		// Reverting to the original explicit instruction to call t.jwkCache.Close().
-		// This will cause a compile error if JWKCacheInterface (and its implementation *JWKCache)
-		// is not updated in jwk.go to include and implement a Close() method
-		// that properly closes the internal *cache.Cache instance.
-		t.jwkCache.Close()
-		t.logger.Debug("t.jwkCache.Close() called as per original instruction.")
-	}
-
-	t.logger.Info("TraefikOidc plugin instance closed successfully.")
-	return nil
-}

main_test.go

@@ -149,11 +149,6 @@ type MockJWKCache struct {
 	Err  error
 }
 
-// Close is a no-op for the mock.
-func (m *MockJWKCache) Close() {
-	// No operation needed for the mock.
-}
-
 func (m *MockJWKCache) GetJWKS(ctx context.Context, jwksURL string, httpClient *http.Client) (*JWKSet, error) {
 	return m.JWKS, m.Err
 }
@@ -164,18 +159,6 @@ func (m *MockJWKCache) Cleanup() {
 	m.Err = nil
 }
 
-// MockTokenVerifier implements TokenVerifier for testing, allowing interception of VerifyToken calls.
-type MockTokenVerifier struct {
-	VerifyFunc func(token string) error
-}
-
-func (m *MockTokenVerifier) VerifyToken(token string) error {
-	if m.VerifyFunc != nil {
-		return m.VerifyFunc(token)
-	}
-	return fmt.Errorf("VerifyFunc not implemented in mock")
-}
-
 // MockTokenExchanger implements TokenExchanger for testing
 type MockTokenExchanger struct {
 	ExchangeCodeFunc func(ctx context.Context, grantType, codeOrToken, redirectURL, codeVerifier string) (*TokenResponse, error)
@@ -462,7 +445,6 @@ func TestServeHTTP(t *testing.T) {
 					"jti": generateRandomString(16), // Unique JTI
 				})
 				session.SetAccessToken(freshToken)
-				session.SetIDToken(freshToken) // Ensure ID token is also set
 				session.SetRefreshToken("valid-refresh-token")
 			},
 			expectedStatus: http.StatusOK,
@@ -630,7 +612,6 @@ func TestServeHTTP(t *testing.T) {
 				session.SetAuthenticated(true)
 				session.SetEmail("[email protected]")
 				session.SetAccessToken(validToken)
-				session.SetIDToken(validToken) // Ensure ID token is also set
 				session.SetRefreshToken("should-not-be-used-refresh-token")
 			},
 			mockRefreshTokenFunc: func(originalFunc func(refreshToken string) (*TokenResponse, error)) func(refreshToken string) (*TokenResponse, error) {
@@ -656,7 +637,6 @@ func TestServeHTTP(t *testing.T) {
 					"jti": generateRandomString(16), // Unique JTI
 				})
 				session.SetAccessToken(freshToken)
-				session.SetIDToken(freshToken) // Ensure ID token is also set
 				session.SetRefreshToken("valid-refresh-token")
 			},
 			requestHeaders: map[string]string{
@@ -678,7 +658,6 @@ func TestServeHTTP(t *testing.T) {
 					"jti": generateRandomString(16), // Unique JTI
 				})
 				session.SetAccessToken(freshToken)
-				session.SetIDToken(freshToken) // Ensure ID token is also set
 				session.SetRefreshToken("valid-refresh-token")
 			},
 			requestHeaders: map[string]string{
@@ -691,45 +670,6 @@ func TestServeHTTP(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			// Reset token blacklist and cache for each test to prevent token replay detection errors
-			ts.tOidc.tokenBlacklist = NewCache() // Use generic cache for blacklist
-			ts.tOidc.tokenCache = NewTokenCache()
-
-			// Reset the global replayCache to prevent "token replay detected" errors
-			replayCacheMu.Lock()
-			replayCache = make(map[string]time.Time) // Reset the global cache
-			replayCacheMu.Unlock()
-
-			// Store original tokenVerifier to restore later
-			origTokenVerifier := ts.tOidc.tokenVerifier
-
-			// Create a mock tokenVerifier that clears the replay cache before verification
-			// This prevents replay detection when the same token is verified multiple times within a test
-			mockTokenVerifier := &MockTokenVerifier{
-				VerifyFunc: func(token string) error {
-					// Clear replay cache before token verification
-					replayCacheMu.Lock()
-					replayCache = make(map[string]time.Time)
-					replayCacheMu.Unlock()
-
-					// Call the original verifier's VerifyToken method
-					// Ensure origTokenVerifier is not nil and is the correct type if necessary,
-					// though in this context it should be the *TraefikOidc instance.
-					if origTokenVerifier != nil {
-						return origTokenVerifier.VerifyToken(token)
-					}
-					return fmt.Errorf("original token verifier is nil")
-				},
-			}
-
-			// Replace tokenVerifier with our mock
-			ts.tOidc.tokenVerifier = mockTokenVerifier
-
-			// Restore original tokenVerifier after test
-			defer func() {
-				ts.tOidc.tokenVerifier = origTokenVerifier
-			}()
-
 			req := httptest.NewRequest("GET", tc.requestPath, nil)
 			// Set common headers needed by the logic (determineScheme, determineHost)
 			req.Header.Set("X-Forwarded-Proto", "http") // Or https if testing that
@@ -1210,79 +1150,30 @@ func TestIsAllowedDomain(t *testing.T) {
 	ts.Setup()
 
 	tests := []struct {
-		name              string
-		email             string
-		allowedDomains    map[string]struct{}
-		allowedUsers      map[string]struct{}
-		allowed           bool
-		expectedLogOutput string // For testing log messages
+		name    string
+		email   string
+		allowed bool
 	}{
 		{
-			name:           "Allowed domain",
-			email:          "[email protected]",
-			allowedDomains: map[string]struct{}{"example.com": {}},
-			allowedUsers:   map[string]struct{}{},
-			allowed:        true,
+			name:    "Allowed domain",
+			email:   "[email protected]",
+			allowed: true,
 		},
 		{
-			name:           "Disallowed domain",
-			email:          "[email protected]",
-			allowedDomains: map[string]struct{}{"example.com": {}},
-			allowedUsers:   map[string]struct{}{},
-			allowed:        false,
+			name:    "Disallowed domain",
+			email:   "[email protected]",
+			allowed: false,
 		},
 		{
-			name:           "Invalid email",
-			email:          "invalid-email",
-			allowedDomains: map[string]struct{}{"example.com": {}},
-			allowedUsers:   map[string]struct{}{},
-			allowed:        false,
-		},
-		{
-			name:           "Specific user is allowed regardless of domain",
-			email:          "[email protected]",
-			allowedDomains: map[string]struct{}{"example.com": {}},
-			allowedUsers:   map[string]struct{}{"[email protected]": {}},
-			allowed:        true,
-		},
-		{
-			name:           "Case-insensitive email matching for specific user",
-			email:          "[email protected]", // Mixed case
-			allowedDomains: map[string]struct{}{"example.com": {}},
-			allowedUsers:   map[string]struct{}{"[email protected]": {}}, // Lowercase
-			allowed:        true,
-		},
-		{
-			name:           "Only allowed users configured (no domains)",
-			email:          "[email protected]",
-			allowedDomains: map[string]struct{}{}, // Empty allowed domains
-			allowedUsers:   map[string]struct{}{"[email protected]": {}},
-			allowed:        true,
-		},
-		{
-			name:           "User not in allowed list when only specific users configured",
-			email:          "[email protected]",
-			allowedDomains: map[string]struct{}{}, // Empty allowed domains
-			allowedUsers:   map[string]struct{}{"[email protected]": {}},
-			allowed:        false,
-		},
-		{
-			name:           "No restrictions (both empty)",
-			email:          "[email protected]",
-			allowedDomains: map[string]struct{}{},
-			allowedUsers:   map[string]struct{}{},
-			allowed:        true,
+			name:    "Invalid email",
+			email:   "invalid-email",
+			allowed: false,
 		},
 	}
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			// Configure TraefikOidc instance for this test case
-			tOidc := ts.tOidc
-			tOidc.allowedUserDomains = tc.allowedDomains
-			tOidc.allowedUsers = tc.allowedUsers
-
-			allowed := tOidc.isAllowedDomain(tc.email)
+			allowed := ts.tOidc.isAllowedDomain(tc.email)
 			if allowed != tc.allowed {
 				t.Errorf("Expected allowed=%v, got %v", tc.allowed, allowed)
 			}

session.go

@@ -156,7 +156,6 @@ func NewSessionManager(encryptionKey string, forceHTTPS bool, logger *Logger) (*
 			accessTokenChunks:  make(map[int]*sessions.Session),
 			refreshTokenChunks: make(map[int]*sessions.Session),
 			refreshMutex:       sync.Mutex{}, // Initialize the mutex
-			dirty:              false,        // Initialize dirty flag
 		}
 	}
 
@@ -190,38 +189,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.
@@ -288,21 +281,6 @@ type SessionData struct {
 
 	// 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.
-func (sd *SessionData) IsDirty() bool {
-	return sd.dirty
-}
-
-// 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)
@@ -324,50 +302,38 @@ func (sd *SessionData) Save(r *http.Request, w http.ResponseWriter) error {
 	sd.accessSession.Options = options
 	sd.refreshSession.Options = options
 
-	var firstErr error
-	// Helper to record first error and log subsequent ones
-	saveOrLogError := func(s *sessions.Session, name string) {
-		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.
+	if err := sd.mainSession.Save(r, w); err != nil {
+		return fmt.Errorf("failed to save main session: %w", err)
 	}
 
-	// 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 {
-		sessionChunk.Options = options
-		saveOrLogError(sessionChunk, fmt.Sprintf("access token chunk %d", i))
+	for _, session := range sd.accessTokenChunks {
+		session.Options = options
+		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 {
-		sessionChunk.Options = options
-		saveOrLogError(sessionChunk, fmt.Sprintf("refresh token chunk %d", i))
+	for _, session := range sd.refreshTokenChunks {
+		session.Options = options
+		if err := session.Save(r, w); err != nil {
+			return fmt.Errorf("failed to save refresh token chunk session: %w", err)
+		}
 	}
 
-	if firstErr == nil {
-		sd.dirty = false // Reset dirty flag only if all saves were successful
-	}
-	return firstErr
+	return nil
 }
 
 // Clear removes all session data associated with this SessionData instance.
@@ -383,57 +349,37 @@ func (sd *SessionData) Save(r *http.Request, w http.ResponseWriter) error {
 //
 // 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
-		for k := range sd.mainSession.Values {
-			delete(sd.mainSession.Values, k)
-		}
+	sd.mainSession.Options.MaxAge = -1
+	sd.accessSession.Options.MaxAge = -1
+	sd.refreshSession.Options.MaxAge = -1
+
+	for k := range sd.mainSession.Values {
+		delete(sd.mainSession.Values, k)
 	}
-	if sd.accessSession != nil {
-		sd.accessSession.Options.MaxAge = -1
-		for k := range sd.accessSession.Values {
-			delete(sd.accessSession.Values, k)
-		}
+	for k := range sd.accessSession.Values {
+		delete(sd.accessSession.Values, k)
 	}
-	if sd.refreshSession != nil {
-		sd.refreshSession.Options.MaxAge = -1
-		for k := range sd.refreshSession.Values {
-			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.
-	// This ensures the session is always returned to the pool,
-	// preventing memory leaks.
+	// Return session to pool.
 	sd.manager.sessionPool.Put(sd)
 
-	// Return the error from Save, if any
 	return err
 }
 
@@ -482,58 +428,18 @@ func (sd *SessionData) GetAuthenticated() bool {
 // 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()
-		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["created_at"] = time.Now().Unix()
 	}
-
 	sd.mainSession.Values["authenticated"] = value
-	if changed {
-		sd.dirty = true
-	}
 	return nil
 }
 
-// ReturnToPool explicitly returns this SessionData object to the pool.
-// 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.
@@ -582,14 +488,6 @@ func (sd *SessionData) GetAccessToken() string {
 // 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 +496,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,19 +504,13 @@ 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["compressed"] = true // Data in chunks is compressed
+		sd.accessSession.Values["token"] = ""
+		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
 		}
 	}
@@ -679,12 +564,6 @@ func (sd *SessionData) GetRefreshToken() string {
 // 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 +572,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,17 +580,13 @@ 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["compressed"] = true // Data in chunks is compressed
+		sd.refreshSession.Values["token"] = ""
+		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
 		}
 	}
@@ -810,11 +678,7 @@ func (sd *SessionData) GetCSRF() string {
 // Parameters:
 //   - token: The CSRF token to store.
 func (sd *SessionData) SetCSRF(token string) {
-	currentVal, _ := sd.mainSession.Values["csrf"].(string)
-	if currentVal != token {
-		sd.mainSession.Values["csrf"] = token
-		sd.dirty = true
-	}
+	sd.mainSession.Values["csrf"] = token
 }
 
 // GetNonce retrieves the OIDC nonce value stored in the main session.
@@ -833,11 +697,7 @@ func (sd *SessionData) GetNonce() string {
 // Parameters:
 //   - nonce: The nonce string to store.
 func (sd *SessionData) SetNonce(nonce string) {
-	currentVal, _ := sd.mainSession.Values["nonce"].(string)
-	if currentVal != nonce {
-		sd.mainSession.Values["nonce"] = nonce
-		sd.dirty = true
-	}
+	sd.mainSession.Values["nonce"] = nonce
 }
 
 // GetCodeVerifier retrieves the PKCE (Proof Key for Code Exchange) code verifier
@@ -856,11 +716,7 @@ func (sd *SessionData) GetCodeVerifier() string {
 // Parameters:
 //   - codeVerifier: The PKCE code verifier string to store.
 func (sd *SessionData) SetCodeVerifier(codeVerifier string) {
-	currentVal, _ := sd.mainSession.Values["code_verifier"].(string)
-	if currentVal != codeVerifier {
-		sd.mainSession.Values["code_verifier"] = codeVerifier
-		sd.dirty = true
-	}
+	sd.mainSession.Values["code_verifier"] = codeVerifier
 }
 
 // GetEmail retrieves the authenticated user's email address stored in the main session.
@@ -879,11 +735,7 @@ func (sd *SessionData) GetEmail() string {
 // Parameters:
 //   - email: The user's email address to store.
 func (sd *SessionData) SetEmail(email string) {
-	currentVal, _ := sd.mainSession.Values["email"].(string)
-	if currentVal != email {
-		sd.mainSession.Values["email"] = email
-		sd.dirty = true
-	}
+	sd.mainSession.Values["email"] = email
 }
 
 // GetIncomingPath retrieves the original request URI (including query parameters)
@@ -903,53 +755,5 @@ func (sd *SessionData) GetIncomingPath() string {
 // 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)
-	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
+	sd.mainSession.Values["incoming_path"] = path
 }

session_test.go

@@ -1,221 +1,389 @@
 package traefikoidc
 
 import (
-	"net/http"
+	"crypto/rand"
+	"fmt"
+	"math/big"
 	"net/http/httptest"
-	"runtime"
 	"strings"
 	"testing"
-	"time"
 )
 
-func TestSessionPoolMemoryLeak(t *testing.T) {
-	logger := NewLogger("debug")
-	sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
-	if err != nil {
-		t.Fatalf("Failed to create session manager: %v", err)
+// generateRandomString creates a random string of specified length
+func generateRandomString(length int) string {
+	const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+	b := make([]byte, length)
+	for i := range b {
+		num, err := rand.Int(rand.Reader, big.NewInt(int64(len(charset))))
+		if err != nil {
+			// Handle error appropriately in a real application, maybe panic in test helper
+			panic(fmt.Sprintf("crypto/rand failed: %v", err))
+		}
+		b[i] = charset[num.Int64()]
+	}
+	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
-	req := httptest.NewRequest("GET", "http://example.com/foo", nil)
+	for _, tt := range tests {
+		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
-	session.Clear(req, nil)
+	// Set some data to ensure cookies are created
+	session.SetAuthenticated(true)
 
-	// Test 2: ReturnToPool explicit method
-	session, err = sm.GetSession(req)
-	if err != nil {
-		t.Fatalf("GetSession failed: %v", err)
+	// Expire any existing cookies
+	session.expireAccessTokenChunks(rr)
+	session.expireRefreshTokenChunks(rr)
+
+	// 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
-	session.ReturnToPool()
-
-	// Test 3: Error path in GetSession
-	// Modify the session store to force an error - use a different encryption key
-	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)
+	// Check cookie prefixes
+	cookies := rr.Result().Cookies()
+	for _, cookie := range cookies {
+		if !strings.HasPrefix(cookie.Name, "_oidc_raczylo_") {
+			t.Errorf("Cookie %s does not have expected prefix '_oidc_raczylo_'", cookie.Name)
+		}
 	}
 }
 
-func TestSessionClearAlwaysReturnsToPool(t *testing.T) {
-	logger := NewLogger("debug")
-	sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
-	if err != nil {
-		t.Fatalf("Failed to create session manager: %v", err)
-	}
+func TestTokenRefreshCleanup(t *testing.T) {
+	req := httptest.NewRequest("GET", "/test", nil)
+	rr := httptest.NewRecorder()
 
-	// Create a test request with the special header that will trigger an error
-	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
+	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)
 	}
 
-	// Create a response writer
-	w := httptest.NewRecorder()
+	// Set a large token that will be split into chunks
+	largeToken := strings.Repeat("x", 5000)
+	session.SetAccessToken(largeToken)
 
-	// Call Clear with the test request (with X-Test-Error header) and response writer
-	// This should trigger the serialization error in Save
-	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)
+	if err := session.Save(req, rr); err != nil {
+		t.Fatalf("Failed to save session: %v", err)
 	}
 
-	// Force GC to ensure any objects are cleaned up
-	runtime.GC()
-	time.Sleep(100 * time.Millisecond)
+	// Get initial cookies
+	initialCookies := rr.Result().Cookies()
 
-	// Create and clear another session (without the error header) to verify the pool is still working
-	normalReq := httptest.NewRequest("GET", "http://example.com/foo", nil)
-	session2, err := sm.GetSession(normalReq)
+	// Create a new request with the initial cookies
+	newReq := httptest.NewRequest("GET", "/test", nil)
+	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
-	t.Log("Session returned to pool despite errors")
+	// Create a response recorder for expired cookies
+	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
-func getPooledObjects(sm *SessionManager) int {
-	// Collect objects until we can't get any more from the pool
-	// Set a max limit to avoid potential infinite loops
-	var objects []*SessionData
-	maxAttempts := 100 // Safety limit to prevent infinite loops
-
-	for i := 0; i < maxAttempts; i++ {
-		obj := sm.sessionPool.Get()
-		if obj == nil {
-			break
-		}
-
-		// Type assertion with validation
-		sessionData, ok := obj.(*SessionData)
-		if !ok {
-			// Return the object even if it's not the right type to avoid leaks
-			sm.sessionPool.Put(obj)
-			break
-		}
-
-		objects = append(objects, sessionData)
+	tests := []struct {
+		name                string
+		authenticated       bool
+		email               string
+		accessToken         string
+		refreshToken        string
+		expectedCookieCount int
+		wantCompressed      bool // Whether tokens should be compressed
+	}{
+		{
+			name:                "Short tokens",
+			authenticated:       true,
+			email:               "[email protected]",
+			accessToken:         "shortaccesstoken",
+			refreshToken:        "shortrefreshtoken",
+			expectedCookieCount: 3, // main, access, refresh
+			wantCompressed:      true,
+		},
+		{
+			name:                "Long tokens exceeding 4096 bytes",
+			authenticated:       true,
+			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
-	count := len(objects)
+	for _, tc := range tests {
+		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
-	for _, obj := range objects {
-		sm.sessionPool.Put(obj)
+			session, err := ts.sessionManager.GetSession(req)
+			if err != nil {
+				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

@@ -82,10 +82,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"`

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_execution_test.go

@@ -192,14 +192,14 @@ func TestTemplateExecutionContext(t *testing.T) {
 		expectedValue string
 	}{
 		{
-			name:         "Access and ID token distinction",
+			name:         "Access and ID token identity",
 			templateText: "Access: {{.AccessToken}} ID: {{.IdToken}}",
 			data: templateData{
-				AccessToken: "access-token-value",
-				IdToken:     "id-token-value", // Now these should be distinct values
+				AccessToken: "access-token",
+				IdToken:     "access-token", // Same as AccessToken in processAuthorizedRequest
 				Claims:      map[string]interface{}{},
 			},
-			expectedValue: "Access: access-token-value ID: id-token-value",
+			expectedValue: "Access: access-token ID: access-token",
 		},
 		{
 			name:         "Combining tokens and claims",

templated_header_integration_test.go

@@ -1,7 +1,6 @@
 package traefikoidc
 
 import (
-	"errors"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -67,28 +66,6 @@ func TestTemplatedHeadersIntegration(t *testing.T) {
 				"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{
@@ -128,19 +105,6 @@ func TestTemplatedHeadersIntegration(t *testing.T) {
 				"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 {
@@ -149,7 +113,7 @@ func TestTemplatedHeadersIntegration(t *testing.T) {
 			token := ts.token
 			if len(tc.claims) > 0 {
 				var err error
-				baseClaims := map[string]interface{}{
+				claims := map[string]interface{}{
 					"iss":   "https://test-issuer.com",
 					"aud":   "test-client-id",
 					"exp":   float64(3000000000), // Far future timestamp
@@ -162,10 +126,10 @@ func TestTemplatedHeadersIntegration(t *testing.T) {
 
 				// Add the test-specific claims
 				for k, v := range tc.claims {
-					baseClaims[k] = v
+					claims[k] = v
 				}
 
-				token, err = createTestJWT(ts.rsaPrivateKey, "RS256", "test-key-id", baseClaims)
+				token, err = createTestJWT(ts.rsaPrivateKey, "RS256", "test-key-id", claims)
 				if err != nil {
 					t.Fatalf("Failed to create test JWT: %v", err)
 				}
@@ -177,17 +141,7 @@ func TestTemplatedHeadersIntegration(t *testing.T) {
 			}
 
 			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
+				tc.expectedHeaders["X-Auth-Info"] = "[email protected], Token=" + token
 			}
 
 			// Store intercepted headers for verification
@@ -204,6 +158,8 @@ func TestTemplatedHeadersIntegration(t *testing.T) {
 				w.WriteHeader(http.StatusOK)
 			})
 
+			// Instead of using New(), we'll directly create a TraefikOidc instance
+			// similar to how it's done in TestSuite.Setup()
 			tOidc := &TraefikOidc{
 				next:               nextHandler,
 				name:               "test",
@@ -218,19 +174,14 @@ func TestTemplatedHeadersIntegration(t *testing.T) {
 				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
+				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),
-				// 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 {
@@ -241,180 +192,55 @@ func TestTemplatedHeadersIntegration(t *testing.T) {
 				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)
-			// 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.SetEmail("[email protected]")
 			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)
 			}
 
+			// 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. Body: %s", http.StatusOK, rr.Code, rr.Body.String())
+				t.Errorf("Expected status code %d, got %d", http.StatusOK, rr.Code)
 			}
 
+			// Verify headers were set correctly
 			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")
-				}
-			}
 		})
 	}
 }
@@ -483,6 +309,8 @@ func TestEdgeCaseTemplatedHeaders(t *testing.T) {
 				w.WriteHeader(http.StatusOK)
 			})
 
+			// Instead of using New(), we'll directly create a TraefikOidc instance
+			// similar to how it's done in TestSuite.Setup()
 			tOidc := &TraefikOidc{
 				next:               nextHandler,
 				name:               "test",
@@ -505,8 +333,6 @@ func TestEdgeCaseTemplatedHeaders(t *testing.T) {
 				extractClaimsFunc:  extractClaims,
 				headerTemplates:    make(map[string]*template.Template),
 			}
-			tOidc.tokenVerifier = tOidc
-			tOidc.jwtVerifier = tOidc
 
 			// Initialize and parse header templates
 			for _, header := range tc.headers {
@@ -517,56 +343,57 @@ func TestEdgeCaseTemplatedHeaders(t *testing.T) {
 				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(token)     // Use the new method
-			session.SetAccessToken(token) // Also set access token to match
+			session.SetAccessToken(token)
 			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)
 			}
 
+			// 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)
 			}
 
-			// 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.
+			// We are primarily checking that these edge cases don't cause panics or errors
+			// For the array test, we can verify the content
+			if tc.name == "Array Claim Access" {
+				// Check if the header was set
+				headerValue := req.Header.Get("X-Roles")
+				expectedValue := "admin,user,manager"
+				if headerValue != expectedValue {
+					t.Errorf("Expected X-Roles header to be %q, got %q", expectedValue, headerValue)
+				}
+			}
 		})
 	}
 }

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)
-	}
-}