Compare commits

...
43 changed files with 4473 additions and 353 deletions
+26
View File
@@ -77,6 +77,7 @@ testData:
revocationURL: https://accounts.google.com/revoke # Endpoint for revoking tokens
oidcEndSessionURL: https://accounts.google.com/logout # Provider's end session endpoint
enablePKCE: false # Enables PKCE (Proof Key for Code Exchange) for additional security
cookieDomain: "" # Explicit domain for session cookies (e.g., ".example.com" for multi-subdomain setups)
# --- Provider Specific Configuration Examples ---
#
@@ -373,6 +374,31 @@ configuration:
Default: false
required: false
cookieDomain:
type: string
description: |
Explicit domain for session cookies. This is important for multi-subdomain setups
and reverse proxy deployments to ensure consistent cookie handling.
When set, all session cookies will use this domain. When not set, the domain
is auto-detected from the request headers (X-Forwarded-Host or Host).
Use a leading dot for subdomain-wide cookies (e.g., ".example.com" allows
cookies to be shared between app.example.com, api.example.com, etc.).
Use a specific domain for host-only cookies (e.g., "app.example.com" restricts
cookies to that exact domain).
This setting is crucial to prevent authentication issues like "CSRF token missing
in session" errors that can occur when cookies are created with inconsistent domains.
Examples:
- ".example.com" - Allows all subdomains to share cookies
- "app.example.com" - Restricts cookies to this specific host
Default: "" (auto-detected from request headers)
required: false
headers:
type: array
description: |
+25
View File
@@ -82,6 +82,7 @@ The middleware supports the following configuration options:
| `oidcEndSessionURL` | The provider's end session endpoint | auto-discovered | `https://accounts.google.com/logout` |
| `enablePKCE` | Enables PKCE (Proof Key for Code Exchange) for authorization code flow | `false` | `true`, `false` |
| `refreshGracePeriodSeconds` | Seconds before token expiry to attempt proactive refresh | `60` | `120` |
| `cookieDomain` | Explicit domain for session cookies (important for multi-subdomain setups) | auto-detected | `.example.com`, `app.example.com` |
| `headers` | Custom HTTP headers with templates that can access OIDC claims and tokens | none | See "Templated Headers" section |
## Scope Configuration
@@ -311,6 +312,30 @@ spec:
- developer
```
### With Cookie Domain Configuration (Multi-Subdomain Setup)
```yaml
apiVersion: traefik.io/v1alpha1
kind: Middleware
metadata:
name: oidc-multi-subdomain
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
cookieDomain: .example.com # Allows cookies to be shared across all subdomains
scopes:
- roles # Appended to defaults: ["openid", "profile", "email", "roles"]
```
**Important**: The `cookieDomain` parameter is crucial when running behind a reverse proxy or when your application serves multiple subdomains. Without it, cookies may be created with inconsistent domains, leading to authentication issues like "CSRF token missing in session" errors.
### With Custom Logging and Rate Limiting
```yaml
+59 -13
View File
@@ -1,16 +1,24 @@
package traefikoidc
import "time"
import (
"sync"
"time"
)
// BackgroundTask represents a managed recurring task that runs in the background.
// It provides a clean interface for starting and stopping periodic operations
// with proper lifecycle management and logging.
type BackgroundTask struct {
stopChan chan struct{}
taskFunc func()
logger *Logger
name string
interval time.Duration
stopChan chan struct{}
taskFunc func()
logger *Logger
name string
interval time.Duration
externalWG *sync.WaitGroup // External WaitGroup (optional, for tracking by parent)
internalWG sync.WaitGroup // Internal WaitGroup for this task's goroutine
stopOnce sync.Once // Ensures Stop() can be called multiple times safely
stopped bool // Track if task has been stopped
mu sync.Mutex // Protects stopped flag
}
// NewBackgroundTask creates a new background task with the specified parameters.
@@ -20,35 +28,73 @@ type BackgroundTask struct {
// - interval: Duration between task executions.
// - taskFunc: The function to execute periodically.
// - logger: Logger instance for task lifecycle events.
// - wg: Optional WaitGroup for synchronizing goroutine completion.
//
// Returns:
// - A configured BackgroundTask ready to be started.
func NewBackgroundTask(name string, interval time.Duration, taskFunc func(), logger *Logger) *BackgroundTask {
func NewBackgroundTask(name string, interval time.Duration, taskFunc func(), logger *Logger, wg ...*sync.WaitGroup) *BackgroundTask {
var externalWG *sync.WaitGroup
if len(wg) > 0 {
externalWG = wg[0]
}
return &BackgroundTask{
name: name,
interval: interval,
stopChan: make(chan struct{}),
taskFunc: taskFunc,
logger: logger,
name: name,
interval: interval,
stopChan: make(chan struct{}),
taskFunc: taskFunc,
logger: logger,
externalWG: externalWG,
}
}
// Start begins the background task execution in a separate goroutine.
// The task runs immediately upon start and then at the specified interval.
func (bt *BackgroundTask) Start() {
bt.mu.Lock()
defer bt.mu.Unlock()
if bt.stopped {
if bt.logger != nil {
bt.logger.Infof("Attempted to start already stopped task: %s", bt.name)
}
return
}
// Add to both internal and external WaitGroups
bt.internalWG.Add(1)
if bt.externalWG != nil {
bt.externalWG.Add(1)
}
go bt.run()
}
// Stop gracefully terminates the background task by closing the stop channel.
// It waits for the goroutine to complete using the internal WaitGroup.
// This method is safe to call multiple times.
func (bt *BackgroundTask) Stop() {
close(bt.stopChan)
bt.stopOnce.Do(func() {
bt.mu.Lock()
bt.stopped = true
bt.mu.Unlock()
close(bt.stopChan)
// Wait only on the internal WaitGroup
bt.internalWG.Wait()
})
}
// run is the main execution loop for the background task.
// It executes the task function immediately and then at regular intervals
// until the stop signal is received.
func (bt *BackgroundTask) run() {
defer func() {
// Always decrement internal WaitGroup
bt.internalWG.Done()
// Decrement external WaitGroup if provided
if bt.externalWG != nil {
bt.externalWG.Done()
}
}()
ticker := time.NewTicker(bt.interval)
defer ticker.Stop()
+3 -3
View File
@@ -14,7 +14,7 @@ import (
// during the callback from Azure AD
func TestAzureOAuthCallbackScenario(t *testing.T) {
t.Run("Azure_OAuth_Complete_Flow", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Step 1: User visits https://app.example.com/protected
@@ -110,7 +110,7 @@ func TestAzureOAuthCallbackScenario(t *testing.T) {
t.Run("Cookie_Not_Sent_With_Wrong_Domain", func(t *testing.T) {
// This test verifies that cookies with wrong domain won't be sent
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Initial request sets cookie for app.example.com
@@ -148,7 +148,7 @@ func TestAzureOAuthCallbackScenario(t *testing.T) {
// This test demonstrates why we can't use SameSite=Strict for OAuth
// With Strict, cookies wouldn't be sent when redirecting from Azure to our app
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// If we had SameSite=Strict (which we don't anymore), the browser would:
+10 -3
View File
@@ -41,9 +41,16 @@ func (m *mockTraefikOidc) VerifyJWTSignatureAndClaims(jwt *JWT, token string) er
}
func TestAzureOIDCRegression(t *testing.T) {
// Create test cleanup helper
tc := newTestCleanup(t)
// Create a mocked TraefikOidc instance configured for Azure AD
mockLogger := NewLogger("debug")
// Create caches with cleanup tracking
tokenCache := tc.addTokenCache(NewTokenCache())
tokenBlacklist := tc.addCache(NewCache())
// Configure for Azure AD provider
baseOidc := &TraefikOidc{
issuerURL: "https://login.microsoftonline.com/tenant-id/v2.0",
@@ -58,8 +65,8 @@ func TestAzureOIDCRegression(t *testing.T) {
logger: mockLogger,
httpClient: createDefaultHTTPClient(), // Add HTTP client
jwkCache: &JWKCache{}, // Add JWK cache
tokenCache: NewTokenCache(),
tokenBlacklist: NewCache(),
tokenCache: tokenCache,
tokenBlacklist: tokenBlacklist,
allowedUserDomains: make(map[string]struct{}),
allowedUsers: make(map[string]struct{}),
allowedRolesAndGroups: make(map[string]struct{}),
@@ -71,7 +78,7 @@ func TestAzureOIDCRegression(t *testing.T) {
tOidc := &mockTraefikOidc{TraefikOidc: baseOidc}
// Initialize session manager
sessionManager, _ := NewSessionManager("test-encryption-key-32-bytes-long", false, mockLogger)
sessionManager, _ := NewSessionManager("test-encryption-key-32-bytes-long", false, "", mockLogger)
tOidc.sessionManager = sessionManager
// Mock the JWT verification to avoid JWKS lookup issues
+12 -6
View File
@@ -53,7 +53,7 @@ func NewCacheWithLogger(logger *Logger) *Cache {
order: list.New(),
elems: make(map[string]*list.Element, DefaultMaxSize),
maxSize: DefaultMaxSize,
autoCleanupInterval: 5 * time.Minute,
autoCleanupInterval: 15 * time.Minute, // Increased from 5 minutes to reduce overhead
logger: logger,
}
c.startAutoCleanup()
@@ -155,16 +155,21 @@ func (c *Cache) Cleanup() {
}
// evictOldest removes the least recently used (oldest) item from the cache.
// It first attempts to find and remove an expired item from the front of the LRU list.
// If no expired items are found at the front, it removes the absolute oldest item (front of the list).
// It first attempts to find and remove expired items, checking up to 5 items
// from the front of the LRU list for efficiency. If no expired items are found,
// it removes the absolute oldest item (front of the list).
// This method is called internally by Set when the cache reaches its maximum size.
// Note: This function assumes the write lock is already held.
func (c *Cache) evictOldest() {
now := time.Now()
elem := c.order.Front()
// First try to find an expired item from the front
for elem != nil {
// Check up to 5 items from the front for expired entries
// This limits the search overhead while still finding expired items efficiently
const maxExpiredCheck = 5
checked := 0
for elem != nil && checked < maxExpiredCheck {
entry := elem.Value.(lruEntry)
if item, exists := c.items[entry.key]; exists {
if now.After(item.ExpiresAt) {
@@ -173,9 +178,10 @@ func (c *Cache) evictOldest() {
}
}
elem = elem.Next()
checked++
}
// If no expired items found, remove the oldest item
// If no expired items found in the first few entries, remove the oldest item
if elem = c.order.Front(); elem != nil {
entry := elem.Value.(lruEntry)
c.removeItem(entry.key)
+336
View File
@@ -0,0 +1,336 @@
package traefikoidc
import (
"container/list"
"fmt"
"runtime"
"sync"
"testing"
"time"
)
// TestCacheMemoryLeaks tests various cache scenarios for memory leaks
func TestCacheMemoryLeaks(t *testing.T) {
t.Run("Cache doesn't release expired items memory", func(t *testing.T) {
runtime.GC()
var m runtime.MemStats
runtime.ReadMemStats(&m)
baselineAlloc := m.Alloc
cache := NewCache()
defer cache.Close()
// Add many large items with short expiration
largeData := make([]byte, 1024*1024) // 1MB
for i := 0; i < 100; i++ {
key := fmt.Sprintf("key-%d", i)
// Items expire in 1 second
cache.Set(key, largeData, 1*time.Second)
}
// Wait for items to expire
time.Sleep(2 * time.Second)
// Force cleanup
cache.Cleanup()
// Check memory after cleanup
runtime.GC()
runtime.ReadMemStats(&m)
afterCleanupAlloc := m.Alloc
allocIncrease := float64(afterCleanupAlloc-baselineAlloc) / 1024 / 1024
t.Logf("Memory after adding and expiring 100MB of data: %.2f MB", allocIncrease)
// Memory should be mostly released after cleanup
if allocIncrease > 10.0 {
t.Errorf("Cache retains too much memory after cleanup: %.2f MB", allocIncrease)
}
})
t.Run("Token blacklist unbounded growth", func(t *testing.T) {
runtime.GC()
var m runtime.MemStats
runtime.ReadMemStats(&m)
baselineAlloc := m.Alloc
blacklist := NewCache()
blacklist.SetMaxSize(1000) // Limit size
defer blacklist.Close()
// Simulate continuous token blacklisting
for i := 0; i < 10000; i++ {
token := fmt.Sprintf("token-%d", i)
// All tokens expire in 24 hours (typical blacklist duration)
blacklist.Set(token, true, 24*time.Hour)
}
runtime.GC()
runtime.ReadMemStats(&m)
currentAlloc := m.Alloc
allocIncrease := float64(currentAlloc-baselineAlloc) / 1024 / 1024
t.Logf("Memory after adding 10000 blacklisted tokens (max 1000): %.2f MB", allocIncrease)
// Should respect max size limit
if len(blacklist.items) > 1000 {
t.Errorf("Blacklist exceeded max size: %d items", len(blacklist.items))
}
// Memory should be bounded
if allocIncrease > 5.0 {
t.Errorf("Blacklist uses too much memory: %.2f MB for max 1000 items", allocIncrease)
}
})
t.Run("Replay cache with high JTI volume", func(t *testing.T) {
initReplayCache()
defer cleanupReplayCache()
runtime.GC()
var m runtime.MemStats
runtime.ReadMemStats(&m)
baselineAlloc := m.Alloc
// Simulate high volume of JTIs
for i := 0; i < 20000; i++ {
jti := fmt.Sprintf("jti-%d", i)
replayCacheMu.Lock()
if replayCache != nil {
// JTIs expire after token expiry (typically 1 hour)
replayCache.Set(jti, true, 1*time.Hour)
}
replayCacheMu.Unlock()
}
runtime.GC()
runtime.ReadMemStats(&m)
currentAlloc := m.Alloc
allocIncrease := float64(currentAlloc-baselineAlloc) / 1024 / 1024
t.Logf("Memory after adding 20000 JTIs (max 10000): %.2f MB", allocIncrease)
// Check size limit is enforced
replayCacheMu.RLock()
cacheSize := 0
if replayCache != nil {
cacheSize = len(replayCache.items)
}
replayCacheMu.RUnlock()
if cacheSize > 10000 {
t.Errorf("Replay cache exceeded max size: %d items", cacheSize)
}
// Memory should be bounded
if allocIncrease > 10.0 {
t.Errorf("Replay cache uses too much memory: %.2f MB for max 10000 items", allocIncrease)
}
})
t.Run("Cache cleanup interval effectiveness", func(t *testing.T) {
// Create a cache with custom settings - don't use NewCache to avoid default cleanup
cache := &Cache{
items: make(map[string]CacheItem, DefaultMaxSize),
order: list.New(),
elems: make(map[string]*list.Element, DefaultMaxSize),
maxSize: DefaultMaxSize,
autoCleanupInterval: 200 * time.Millisecond, // Fast cleanup for test
logger: newNoOpLogger(),
}
// Start cleanup with our custom interval
cache.startAutoCleanup()
defer cache.Close()
// Add expired items
for i := 0; i < 1000; i++ {
key := fmt.Sprintf("key-%d", i)
cache.Set(key, "data", 50*time.Millisecond) // Very short expiry
}
// Wait for items to expire and cleanup to run (at least 2 cleanup cycles)
time.Sleep(600 * time.Millisecond)
// Manually trigger cleanup to ensure it runs
cache.Cleanup()
// Check that expired items are removed
cache.mutex.RLock()
remainingItems := len(cache.items)
cache.mutex.RUnlock()
t.Logf("Remaining items after auto cleanup: %d", remainingItems)
if remainingItems > 100 {
t.Errorf("Auto cleanup not effective: %d items remain", remainingItems)
}
})
t.Run("Concurrent cache operations memory stability", func(t *testing.T) {
cache := NewCache()
defer cache.Close()
runtime.GC()
var m runtime.MemStats
runtime.ReadMemStats(&m)
baselineAlloc := m.Alloc
var wg sync.WaitGroup
stop := make(chan struct{})
// Writers continuously add items
for i := 0; i < 5; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
for j := 0; j < 1000; j++ {
select {
case <-stop:
return
default:
key := fmt.Sprintf("writer-%d-%d", id, j)
cache.Set(key, "data", 1*time.Second)
time.Sleep(1 * time.Millisecond)
}
}
}(i)
}
// Readers continuously read items
for i := 0; i < 5; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
for j := 0; j < 1000; j++ {
select {
case <-stop:
return
default:
key := fmt.Sprintf("writer-%d-%d", id%5, j)
cache.Get(key)
time.Sleep(1 * time.Millisecond)
}
}
}(i)
}
// Let it run for a bit
time.Sleep(5 * time.Second)
close(stop)
wg.Wait()
runtime.GC()
runtime.ReadMemStats(&m)
finalAlloc := m.Alloc
// Handle potential underflow
var allocIncrease float64
if finalAlloc > baselineAlloc {
allocIncrease = float64(finalAlloc-baselineAlloc) / 1024 / 1024
} else {
allocIncrease = -float64(baselineAlloc-finalAlloc) / 1024 / 1024
}
t.Logf("Memory increase under concurrent load: %.2f MB", allocIncrease)
if allocIncrease > 5.0 {
t.Errorf("Memory leak under concurrent operations: %.2f MB", allocIncrease)
}
})
t.Run("LRU eviction memory release", func(t *testing.T) {
cache := NewCache()
cache.SetMaxSize(100) // Small cache
defer cache.Close()
runtime.GC()
var m runtime.MemStats
runtime.ReadMemStats(&m)
baselineAlloc := m.Alloc
// Add many items to trigger eviction
for i := 0; i < 1000; i++ {
key := fmt.Sprintf("key-%d", i)
data := make([]byte, 10240) // 10KB per item
cache.Set(key, data, 1*time.Hour)
}
runtime.GC()
runtime.ReadMemStats(&m)
afterEvictionAlloc := m.Alloc
allocIncrease := float64(afterEvictionAlloc-baselineAlloc) / 1024 / 1024
t.Logf("Memory after LRU eviction (1000 items, max 100): %.2f MB", allocIncrease)
// Should only keep 100 items worth of memory
if allocIncrease > 2.0 { // 100 * 10KB = ~1MB
t.Errorf("LRU eviction doesn't release memory properly: %.2f MB", allocIncrease)
}
// Verify cache size
if len(cache.items) > 100 {
t.Errorf("Cache size exceeded limit: %d items", len(cache.items))
}
})
t.Run("Token cache with claims memory", func(t *testing.T) {
tokenCache := NewTokenCache()
defer tokenCache.Close()
runtime.GC()
var m runtime.MemStats
runtime.ReadMemStats(&m)
baselineAlloc := m.Alloc
// Add tokens with large claims
for i := 0; i < 1000; i++ {
token := fmt.Sprintf("token-%d", i)
claims := map[string]interface{}{
"sub": fmt.Sprintf("user-%d", i),
"email": fmt.Sprintf("user%[email protected]", i),
"groups": make([]string, 100), // Large groups list
"data": make([]byte, 1024), // Extra data
}
tokenCache.Set(token, claims, 1*time.Hour)
}
runtime.GC()
runtime.ReadMemStats(&m)
currentAlloc := m.Alloc
allocIncrease := float64(currentAlloc-baselineAlloc) / 1024 / 1024
t.Logf("Memory after adding 1000 tokens with large claims: %.2f MB", allocIncrease)
// Check if memory is reasonable
if allocIncrease > 20.0 {
t.Errorf("Token cache uses excessive memory: %.2f MB", allocIncrease)
}
})
}
// TestCacheEvictionBug tests the inefficient eviction in evictOldest
func TestCacheEvictionBug(t *testing.T) {
t.Run("evictOldest scans entire list", func(t *testing.T) {
cache := NewCache()
cache.SetMaxSize(100)
defer cache.Close()
// Fill cache with non-expired items
for i := 0; i < 100; i++ {
key := fmt.Sprintf("key-%d", i)
cache.Set(key, "data", 1*time.Hour) // Long expiry
}
// Try to add one more item to trigger eviction
start := time.Now()
cache.Set("trigger", "data", 1*time.Hour)
elapsed := time.Since(start)
t.Logf("Time to evict and add one item: %v", elapsed)
// Should be fast even with full cache
if elapsed > 10*time.Millisecond {
t.Errorf("Eviction too slow, possibly scanning entire list: %v", elapsed)
}
})
}
+27 -28
View File
@@ -7,45 +7,36 @@ import (
func TestCache_Cleanup(t *testing.T) {
c := NewCache()
defer c.Close() // Stop the auto-cleanup goroutine
// 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
// Create test items
c.items["expired"] = CacheItem{
Value: "expired-value",
ExpiresAt: pastTime,
}
c.items["valid"] = CacheItem{
Value: "valid-value",
ExpiresAt: futureTime,
}
// 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"})
// Use the Set method to properly add items with synchronization
c.Set("expired", "expired-value", pastTime.Sub(now))
c.Set("valid", "valid-value", futureTime.Sub(now))
// Call cleanup, which should only remove expired items
c.Cleanup()
// Check that only the expired item was removed
if _, exists := c.items["expired"]; exists {
if _, exists := c.Get("expired"); exists {
t.Error("Expired item was not removed by Cleanup()")
}
if _, exists := c.items["valid"]; !exists {
if _, exists := c.Get("valid"); !exists {
t.Error("Valid item was incorrectly removed by Cleanup()")
}
}
func TestCache_SetMaxSize(t *testing.T) {
c := NewCache()
defer c.Close() // Stop the auto-cleanup goroutine
// Set a lower max size
originalMaxSize := c.maxSize
originalMaxSize := 10
newMaxSize := 3
// Add more items than the new max size
@@ -54,25 +45,33 @@ func TestCache_SetMaxSize(t *testing.T) {
c.Set(key, i, 1*time.Hour)
}
// Verify items were added
if len(c.items) != originalMaxSize {
t.Errorf("Expected %d items before SetMaxSize, got %d", originalMaxSize, len(c.items))
// Verify items were added by checking a few keys
for i := 0; i < originalMaxSize; i++ {
key := "key" + string(rune('A'+i))
if _, exists := c.Get(key); !exists {
t.Errorf("Expected key %s to exist before SetMaxSize", key)
}
}
// Change the max size to a smaller value
c.SetMaxSize(newMaxSize)
// Count remaining items
count := 0
for i := 0; i < originalMaxSize; i++ {
key := "key" + string(rune('A'+i))
if _, exists := c.Get(key); exists {
count++
}
}
// 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)
if count > newMaxSize {
t.Errorf("Cache size %d exceeds new max size %d after SetMaxSize", count, newMaxSize)
}
if c.maxSize != newMaxSize {
t.Errorf("Cache maxSize not updated, expected %d, got %d", newMaxSize, c.maxSize)
}
// Check that the oldest items were evicted (should keep "keyC", "keyD", "keyE", etc.)
if _, exists := c.items["keyA"]; exists {
// Check that the oldest items were evicted (keyA should be evicted)
if _, exists := c.Get("keyA"); exists {
t.Error("Expected oldest item 'keyA' to be evicted, but it still exists")
}
}
+218
View File
@@ -0,0 +1,218 @@
package traefikoidc
import (
"net/http"
"net/http/httptest"
"strings"
"testing"
)
// TestCookieCleanupBehavior tests the cookie cleanup function's behavior
func TestCookieCleanupBehavior(t *testing.T) {
logger := NewLogger("debug")
// Create session manager with a specific domain
sm, err := NewSessionManager("test-encryption-key-32-characters", false, "app.example.com", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
// Create request from app.example.com
req := httptest.NewRequest("GET", "http://app.example.com/test", nil)
// Add some OIDC cookies (browsers don't send domain info)
cookies := []*http.Cookie{
{
Name: "_oidc_raczylo_m",
Value: "test-value-1",
},
{
Name: "_oidc_raczylo_a",
Value: "test-value-2",
},
{
Name: "access_token_chunk_0",
Value: "chunk-value",
},
}
for _, cookie := range cookies {
req.AddCookie(cookie)
}
// Create response recorder to capture Set-Cookie headers
rr := httptest.NewRecorder()
// Run cleanup - it should attempt to delete cookies with various domains
sm.CleanupOldCookies(rr, req)
// Check Set-Cookie headers
setCookies := rr.Result().Cookies()
// The cleanup should have attempted to delete cookies with various domain variations
// We should see deletion attempts for:
// - app.example.com
// - .app.example.com
// - example.com
// - .example.com
var deletionAttempts int
domainsAttempted := make(map[string]bool)
for _, cookie := range setCookies {
if cookie.MaxAge == -1 {
deletionAttempts++
domainsAttempted[cookie.Domain] = true
}
}
// We should see deletion attempts for various domain variations
// Note: The exact number depends on the implementation, but we should see multiple attempts
if deletionAttempts == 0 {
t.Error("Expected cleanup to attempt cookie deletions, but none were found")
}
// Log the domains attempted for debugging
t.Logf("Deletion attempts: %d", deletionAttempts)
for domain := range domainsAttempted {
t.Logf("Attempted deletion for domain: %q", domain)
}
}
// TestConfiguredDomainPersistence tests that configured domain is consistently used
func TestConfiguredDomainPersistence(t *testing.T) {
logger := NewLogger("debug")
// Create session manager with explicit domain configuration
configuredDomain := ".example.com"
sm, err := NewSessionManager("test-encryption-key-32-characters", false, configuredDomain, logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
// Create requests from different subdomains
tests := []struct {
requestHost string
name string
}{
{
name: "Request from main domain",
requestHost: "example.com",
},
{
name: "Request from subdomain",
requestHost: "app.example.com",
},
{
name: "Request from nested subdomain",
requestHost: "api.app.example.com",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
req := httptest.NewRequest("GET", "http://"+tt.requestHost+"/test", nil)
rr := httptest.NewRecorder()
// Get session and set a value
session, err := sm.GetSession(req)
if err != nil {
t.Fatalf("Failed to get session: %v", err)
}
defer session.ReturnToPool()
session.SetNonce("test-nonce")
// Save session
err = session.Save(req, rr)
if err != nil {
t.Fatalf("Failed to save session: %v", err)
}
// Check that all cookies use the configured domain
cookies := rr.Result().Cookies()
for _, cookie := range cookies {
if strings.HasPrefix(cookie.Name, "_oidc_") {
// The domain should match the configured domain (with possible normalization)
expectedDomains := []string{configuredDomain, strings.TrimPrefix(configuredDomain, ".")}
domainMatches := false
for _, expected := range expectedDomains {
if cookie.Domain == expected || cookie.Domain == "" {
domainMatches = true
break
}
}
if !domainMatches {
t.Errorf("Cookie %s has unexpected domain %q, expected one of %v",
cookie.Name, cookie.Domain, expectedDomains)
}
}
}
})
}
}
// TestDomainMigration simulates migrating from no configured domain to explicit domain
func TestDomainMigration(t *testing.T) {
logger := NewLogger("debug")
// Step 1: Create session without configured domain (auto-detection)
sm1, err := NewSessionManager("test-encryption-key-32-characters", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
req1 := httptest.NewRequest("GET", "http://app.example.com/test", nil)
rr1 := httptest.NewRecorder()
session1, err := sm1.GetSession(req1)
if err != nil {
t.Fatalf("Failed to get session: %v", err)
}
session1.SetNonce("test-nonce")
session1.Save(req1, rr1)
session1.ReturnToPool()
// The cookies will have auto-detected domain
oldCookies := rr1.Result().Cookies()
t.Logf("Old cookies count: %d", len(oldCookies))
// Step 2: Create new session manager with explicit domain configuration
sm2, err := NewSessionManager("test-encryption-key-32-characters", false, ".example.com", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
// Step 3: Make request with old cookies
req2 := httptest.NewRequest("GET", "http://app.example.com/test", nil)
// Add the old cookies to the new request
for _, cookie := range oldCookies {
// Simulate browser behavior - don't include domain in Cookie header
simpleCookie := &http.Cookie{
Name: cookie.Name,
Value: cookie.Value,
}
req2.AddCookie(simpleCookie)
}
rr2 := httptest.NewRecorder()
// Run cleanup - should attempt to delete old cookies
sm2.CleanupOldCookies(rr2, req2)
// Check that deletion cookies were sent
newCookies := rr2.Result().Cookies()
deletionCount := 0
for _, cookie := range newCookies {
if cookie.MaxAge == -1 {
deletionCount++
}
}
if deletionCount == 0 {
t.Error("Expected cleanup to send deletion cookies during migration, but none were found")
}
t.Logf("Sent %d deletion cookies during migration", deletionCount)
}
+205
View File
@@ -0,0 +1,205 @@
package traefikoidc
import (
"net/http"
"net/http/httptest"
"strings"
"testing"
)
// TestCookieDomainConfiguration tests that the cookie domain configuration is properly applied
func TestCookieDomainConfiguration(t *testing.T) {
tests := []struct {
configDomain string
requestHost string
forwardedHost string
expectedDomain string
name string
}{
{
name: "Configured domain takes precedence",
configDomain: ".example.com",
requestHost: "app.example.com",
expectedDomain: ".example.com",
},
{
name: "Auto-detection when no domain configured",
configDomain: "",
requestHost: "app.example.com",
expectedDomain: "app.example.com",
},
{
name: "X-Forwarded-Host used for auto-detection",
configDomain: "",
requestHost: "internal.local",
forwardedHost: "public.example.com",
expectedDomain: "public.example.com",
},
{
name: "No domain for localhost",
configDomain: "",
requestHost: "localhost:8080",
expectedDomain: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// Create session manager with configured domain
logger := NewLogger("debug")
sm, err := NewSessionManager("test-encryption-key-32-characters", false, tt.configDomain, logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
// Create request
req := httptest.NewRequest("GET", "http://"+tt.requestHost+"/test", nil)
if tt.forwardedHost != "" {
req.Header.Set("X-Forwarded-Host", tt.forwardedHost)
}
// Create a dummy response writer to test getCookieOptions behavior
// We'll examine the session options domain instead
options := &http.Cookie{
Domain: sm.cookieDomain,
}
// If no configured domain, simulate auto-detection
if sm.cookieDomain == "" && req != nil {
host := req.Host
if forwardedHost := req.Header.Get("X-Forwarded-Host"); forwardedHost != "" {
host = forwardedHost
}
if host != "" && !strings.Contains(host, "localhost") && !strings.Contains(host, "127.0.0.1") {
if colonIndex := strings.Index(host, ":"); colonIndex != -1 {
host = host[:colonIndex]
}
options.Domain = host
}
}
// Check domain
if options.Domain != tt.expectedDomain {
t.Errorf("Expected domain %q, got %q", tt.expectedDomain, options.Domain)
}
})
}
}
// TestCookieDomainConsistency tests that all session cookies use the same domain
func TestCookieDomainConsistency(t *testing.T) {
logger := NewLogger("debug")
// Test with configured domain
sm, err := NewSessionManager("test-encryption-key-32-characters", false, ".example.com", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
req := httptest.NewRequest("GET", "http://app.example.com/test", nil)
rr := httptest.NewRecorder()
// Get session and set some values
session, err := sm.GetSession(req)
if err != nil {
t.Fatalf("Failed to get session: %v", err)
}
defer session.ReturnToPool()
// Set various session values including ID token
session.SetNonce("test-nonce")
session.SetAccessToken("test-access-token")
session.SetRefreshToken("test-refresh-token")
// Set a valid JWT-like ID token (needs 2 dots)
session.SetIDToken("header.payload.signature")
// Save session
err = session.Save(req, rr)
if err != nil {
t.Fatalf("Failed to save session: %v", err)
}
// Check all cookies have the same domain
cookies := rr.Result().Cookies()
var seenDomain string
for _, cookie := range cookies {
// Only check our OIDC cookies
if strings.HasPrefix(cookie.Name, "_oidc_") ||
strings.HasPrefix(cookie.Name, "access_token_chunk_") ||
strings.HasPrefix(cookie.Name, "refresh_token_chunk_") {
// Normalize domain for comparison (handle leading dot and empty domain)
normalizedDomain := cookie.Domain
if normalizedDomain == "" {
// Empty domain means host-only cookie, should match configured domain
normalizedDomain = "example.com"
} else if strings.HasPrefix(normalizedDomain, ".") {
normalizedDomain = normalizedDomain[1:]
}
if seenDomain == "" {
seenDomain = normalizedDomain
} else if normalizedDomain != seenDomain {
t.Errorf("Inconsistent cookie domains: %q vs %q for cookie %s",
seenDomain, normalizedDomain, cookie.Name)
}
// Verify it matches configured domain (browsers may normalize by removing leading dot)
// Empty domain is also acceptable for host-only cookies
if cookie.Domain != "" && cookie.Domain != ".example.com" && cookie.Domain != "example.com" {
t.Errorf("Cookie %s has domain %q, expected %q, %q, or empty",
cookie.Name, cookie.Domain, ".example.com", "example.com")
}
}
}
}
// TestCookieDomainWithReverseProxy simulates a reverse proxy scenario
func TestCookieDomainWithReverseProxy(t *testing.T) {
logger := NewLogger("debug")
// No configured domain, should auto-detect from X-Forwarded-Host
sm, err := NewSessionManager("test-encryption-key-32-characters", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
// Simulate reverse proxy request
req := httptest.NewRequest("GET", "http://internal.local:8080/test", nil)
req.Header.Set("X-Forwarded-Host", "public.example.com")
req.Header.Set("X-Forwarded-Proto", "https")
// Test the domain configuration
// Since getCookieOptions is private, we'll check the configured domain directly
options := &http.Cookie{
Domain: sm.cookieDomain,
}
// If no configured domain, simulate auto-detection
if sm.cookieDomain == "" && req != nil {
host := req.Host
if forwardedHost := req.Header.Get("X-Forwarded-Host"); forwardedHost != "" {
host = forwardedHost
}
if host != "" && !strings.Contains(host, "localhost") && !strings.Contains(host, "127.0.0.1") {
if colonIndex := strings.Index(host, ":"); colonIndex != -1 {
host = host[:colonIndex]
}
options.Domain = host
}
}
// Check secure flag based on X-Forwarded-Proto
isSecure := req.Header.Get("X-Forwarded-Proto") == "https" || req.TLS != nil
options.Secure = isSecure // Note: forceHTTPS is private so we can't access it in test
// Should use the forwarded host
if options.Domain != "public.example.com" {
t.Errorf("Expected domain from X-Forwarded-Host %q, got %q",
"public.example.com", options.Domain)
}
// Should be secure due to X-Forwarded-Proto
if !options.Secure {
t.Error("Expected Secure flag to be true with X-Forwarded-Proto: https")
}
}
+8 -8
View File
@@ -18,7 +18,7 @@ func TestCSRFTokenSessionManagement(t *testing.T) {
// Test that CSRF tokens persist through the authentication flow
t.Run("CSRF_Token_Persists_After_Selective_Clear", func(t *testing.T) {
// Create a session manager
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Create initial request
@@ -90,7 +90,7 @@ func TestCSRFTokenSessionManagement(t *testing.T) {
// Test that marking session as dirty forces save
t.Run("Mark_Dirty_Forces_Session_Save", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
req := httptest.NewRequest("GET", "http://example.com/test", nil)
@@ -126,7 +126,7 @@ func TestCSRFTokenSessionManagement(t *testing.T) {
// Test Azure-specific session handling
t.Run("Azure_Session_Cookie_Configuration", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Simulate Azure callback scenario
@@ -158,7 +158,7 @@ func TestCSRFTokenSessionManagement(t *testing.T) {
// Test session continuity through auth flow
t.Run("Session_Continuity_Through_Auth_Flow", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Step 1: Initial request
@@ -199,7 +199,7 @@ func TestCSRFTokenSessionManagement(t *testing.T) {
// Test large token handling doesn't affect CSRF
t.Run("Large_Tokens_Dont_Affect_CSRF", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
req := httptest.NewRequest("GET", "http://example.com/test", nil)
@@ -262,7 +262,7 @@ func TestAuthFlowWithoutExternalDependencies(t *testing.T) {
// We can't fully initialize TraefikOidc without network access,
// but we can test the session management directly
sessionManager, err := NewSessionManager(plugin.SessionEncryptionKey, plugin.ForceHTTPS, NewLogger(plugin.LogLevel))
sessionManager, err := NewSessionManager(plugin.SessionEncryptionKey, plugin.ForceHTTPS, "", NewLogger(plugin.LogLevel))
require.NoError(t, err)
t.Run("Session_Created_On_Protected_Request", func(t *testing.T) {
@@ -291,7 +291,7 @@ func TestAuthFlowWithoutExternalDependencies(t *testing.T) {
// TestRegressionLoginLoop specifically tests the fix for issue #53
func TestRegressionLoginLoop(t *testing.T) {
// This test verifies that the specific changes made to fix the login loop work correctly
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Simulate the exact flow that was causing the login loop
@@ -392,7 +392,7 @@ func TestRegressionLoginLoop(t *testing.T) {
// TestCSRFValidationTiming tests timing-sensitive CSRF validation scenarios
func TestCSRFValidationTiming(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
t.Run("Rapid_Redirect_Maintains_CSRF", func(t *testing.T) {
+664
View File
@@ -0,0 +1,664 @@
package traefikoidc
import (
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"runtime"
"runtime/debug"
"strings"
"sync"
"testing"
"time"
"golang.org/x/time/rate"
)
// TestBackgroundGoroutineLeaks tests that background goroutines don't leak memory
// even when no requests are made to protected resources
func TestBackgroundGoroutineLeaks(t *testing.T) {
t.Run("Idle middleware memory growth", func(t *testing.T) {
// Force GC to get clean baseline
runtime.GC()
runtime.GC()
time.Sleep(100 * time.Millisecond)
var m runtime.MemStats
runtime.ReadMemStats(&m)
baselineAlloc := m.Alloc
baselineGoroutines := runtime.NumGoroutine()
t.Logf("Baseline: Memory=%d KB, Goroutines=%d", baselineAlloc/1024, baselineGoroutines)
// Create test cleanup
tc := newTestCleanup(t)
// Create middleware instance with mock setup
config := createTestConfig()
config.LogLevel = "error" // Reduce noise
middleware, server := setupTestOIDCMiddleware(t, config)
tc.addServer(server)
tc.addOIDC(middleware)
// Let it sit idle for a while - simulating no requests
// During this time, background goroutines are running
t.Log("Letting middleware sit idle for 3 seconds...")
// Take measurements every 1 second
for i := 0; i < 3; i++ {
time.Sleep(1 * time.Second)
runtime.GC()
runtime.ReadMemStats(&m)
currentAlloc := m.Alloc
currentGoroutines := runtime.NumGoroutine()
allocIncrease := float64(currentAlloc-baselineAlloc) / 1024 / 1024
goroutineIncrease := currentGoroutines - baselineGoroutines
t.Logf("After %d seconds: Memory increase=%.2f MB, Goroutine increase=%d",
(i + 1), allocIncrease, goroutineIncrease)
// Check for significant memory growth (more than 5MB)
if allocIncrease > 5.0 {
t.Errorf("Significant memory increase detected: %.2f MB after %d seconds of idle",
allocIncrease, (i + 1))
}
// Check for goroutine leaks (more than 10 extra goroutines)
if goroutineIncrease > 10 {
t.Errorf("Goroutine leak detected: %d extra goroutines after %d seconds",
goroutineIncrease, (i+1)*5)
}
}
// Clean up
if err := middleware.Close(); err != nil {
t.Errorf("Failed to close middleware: %v", err)
}
// Wait for cleanup
time.Sleep(500 * time.Millisecond)
// Final check
runtime.GC()
runtime.ReadMemStats(&m)
finalAlloc := m.Alloc
finalGoroutines := runtime.NumGoroutine()
finalAllocIncrease := float64(finalAlloc-baselineAlloc) / 1024 / 1024
finalGoroutineIncrease := finalGoroutines - baselineGoroutines
t.Logf("Final: Memory increase=%.2f MB, Goroutine increase=%d",
finalAllocIncrease, finalGoroutineIncrease)
if finalGoroutineIncrease > 2 {
t.Errorf("Goroutines not cleaned up properly: %d extra goroutines remain",
finalGoroutineIncrease)
}
})
}
// TestHTTPClientConnectionLeaks tests that HTTP clients don't leak connections
func TestHTTPClientConnectionLeaks(t *testing.T) {
t.Run("HTTP client connection accumulation", func(t *testing.T) {
// Create test server that simulates OIDC endpoints
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/.well-known/openid-configuration":
w.Header().Set("Content-Type", "application/json")
w.Write([]byte(`{
"issuer": "https://example.com",
"authorization_endpoint": "https://example.com/auth",
"token_endpoint": "https://example.com/token",
"jwks_uri": "https://example.com/jwks",
"userinfo_endpoint": "https://example.com/userinfo"
}`))
case "/jwks":
w.Header().Set("Content-Type", "application/json")
w.Write([]byte(`{"keys": []}`))
default:
w.WriteHeader(http.StatusNotFound)
}
}))
defer server.Close()
// Monitor connection count
getActiveConnections := func(client *http.Client) int {
if transport, ok := client.Transport.(*http.Transport); ok {
// This is a simplified check - in reality we'd need to inspect
// the transport's connection pool
return transport.MaxIdleConnsPerHost
}
return 0
}
// Create multiple HTTP clients like the middleware does
clients := make([]*http.Client, 10)
for i := 0; i < 10; i++ {
clients[i] = createDefaultHTTPClient()
// Make requests
resp, err := clients[i].Get(server.URL + "/.well-known/openid-configuration")
if err != nil {
t.Fatal(err)
}
resp.Body.Close()
// Check connection settings
conns := getActiveConnections(clients[i])
if conns > 1 {
t.Logf("Client %d has %d max idle connections per host", i, conns)
}
}
// Let connections sit idle briefly to test cleanup
time.Sleep(1 * time.Second) // Reduced for faster tests
// Force cleanup
for _, client := range clients {
if transport, ok := client.Transport.(*http.Transport); ok {
transport.CloseIdleConnections()
}
}
t.Log("HTTP clients cleaned up successfully")
})
}
// TestCacheBackgroundTaskLeaks tests that cache background tasks don't leak
func TestCacheBackgroundTaskLeaks(t *testing.T) {
t.Run("Multiple cache instances with cleanup tasks", func(t *testing.T) {
// Reset global state to prevent test interference
resetGlobalState()
initialGoroutines := runtime.NumGoroutine()
// Create many cache instances
caches := make([]*Cache, 50)
// Ensure cleanup even on test failure or panic
defer func() {
for _, cache := range caches {
if cache != nil {
cache.Close()
}
}
// Wait a bit for goroutines to stop
time.Sleep(100 * time.Millisecond)
}()
for i := 0; i < 50; i++ {
caches[i] = NewCache()
// Add some data
for j := 0; j < 100; j++ {
key := fmt.Sprintf("key-%d-%d", i, j)
caches[i].Set(key, "value", 5*time.Minute)
}
}
// Wait for all cleanup goroutines to start
time.Sleep(200 * time.Millisecond)
afterCreateGoroutines := runtime.NumGoroutine()
goroutineIncrease := afterCreateGoroutines - initialGoroutines
t.Logf("Created %d caches, goroutine increase: %d", len(caches), goroutineIncrease)
// Expected: one cleanup goroutine per cache
if goroutineIncrease < len(caches) {
t.Errorf("Expected at least %d goroutines, got %d increase",
len(caches), goroutineIncrease)
}
// Close all caches
for _, cache := range caches {
cache.Close()
}
// Wait for goroutines to stop with timeout
done := make(chan bool)
go func() {
for i := 0; i < 20; i++ { // Try for 2 seconds
time.Sleep(100 * time.Millisecond)
currentGoroutines := runtime.NumGoroutine()
if currentGoroutines-initialGoroutines <= 5 {
done <- true
return
}
}
done <- false
}()
success := <-done
finalGoroutines := runtime.NumGoroutine()
remainingGoroutines := finalGoroutines - initialGoroutines
if !success && remainingGoroutines > 5 { // Allow small tolerance
t.Errorf("Cache cleanup goroutines not stopped properly: %d extra goroutines remain",
remainingGoroutines)
}
})
}
// TestGlobalSingletonMemoryGrowth tests that global singletons don't grow unbounded
func TestGlobalSingletonMemoryGrowth(t *testing.T) {
t.Run("Global cache manager memory growth", func(t *testing.T) {
// Clean up any existing global state
CleanupGlobalCacheManager()
CleanupGlobalMemoryPools()
runtime.GC()
var m runtime.MemStats
runtime.ReadMemStats(&m)
baselineAlloc := m.Alloc
// Get global cache manager
wg := &sync.WaitGroup{}
cm := GetGlobalCacheManager(wg)
// Simulate continuous usage without cleanup
for i := 0; i < 1000; i++ {
// Add to token cache
cm.GetSharedTokenCache().Set(
fmt.Sprintf("token-%d", i),
map[string]interface{}{"claim": fmt.Sprintf("value-%d", i)},
5*time.Minute,
)
// Add to blacklist
cm.GetSharedTokenBlacklist().Set(
fmt.Sprintf("blacklist-%d", i),
true,
5*time.Minute,
)
// Every 100 iterations, check memory
if i%100 == 0 {
runtime.GC()
runtime.ReadMemStats(&m)
currentAlloc := m.Alloc
allocIncrease := float64(currentAlloc-baselineAlloc) / 1024 / 1024
t.Logf("After %d items: Memory increase=%.2f MB", i, allocIncrease)
// The caches should have max size limits
// If memory grows more than 10MB, there's likely a leak
if allocIncrease > 10.0 {
t.Errorf("Excessive memory growth in global caches: %.2f MB after %d items",
allocIncrease, i)
break
}
}
}
// Cleanup
CleanupGlobalCacheManager()
CleanupGlobalMemoryPools()
wg.Wait()
// Force full cleanup of replay cache too
cleanupReplayCache()
runtime.GC()
runtime.ReadMemStats(&m)
finalAlloc := m.Alloc
// Handle potential negative differences (memory freed beyond baseline)
var finalAllocIncrease float64
if finalAlloc >= baselineAlloc {
finalAllocIncrease = float64(finalAlloc-baselineAlloc) / 1024 / 1024
} else {
// Memory decreased below baseline - this is good!
finalAllocIncrease = -float64(baselineAlloc-finalAlloc) / 1024 / 1024
}
t.Logf("Final memory increase after cleanup: %.2f MB", finalAllocIncrease)
if finalAllocIncrease > 2.0 {
t.Errorf("Memory not properly released after cleanup: %.2f MB remains",
finalAllocIncrease)
}
})
}
// TestMetadataCacheRefreshLeak tests for memory leaks in metadata refresh
func TestMetadataCacheRefreshLeak(t *testing.T) {
t.Run("Metadata cache refresh memory leak", func(t *testing.T) {
wg := &sync.WaitGroup{}
cache := NewMetadataCacheWithLogger(wg, NewLogger("error"))
defer cache.Close()
// Mock HTTP client that returns metadata
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
// Return a large metadata response to amplify any leaks
w.Write([]byte(`{
"issuer": "https://example.com",
"authorization_endpoint": "https://example.com/auth",
"token_endpoint": "https://example.com/token",
"jwks_uri": "https://example.com/jwks",
"userinfo_endpoint": "https://example.com/userinfo",
"extra_field_1": "` + strings.Repeat("x", 1000) + `",
"extra_field_2": "` + strings.Repeat("y", 1000) + `"
}`))
}))
defer server.Close()
client := &http.Client{Timeout: 5 * time.Second}
runtime.GC()
var m runtime.MemStats
runtime.ReadMemStats(&m)
baselineAlloc := m.Alloc
// Simulate many metadata refreshes
for i := 0; i < 100; i++ {
_, err := cache.GetMetadata(server.URL, client, NewLogger("error"))
if err != nil {
t.Logf("Metadata fetch error (expected for test server): %v", err)
}
// Force cache expiry to trigger refresh
cache.mutex.Lock()
cache.expiresAt = time.Now().Add(-1 * time.Hour)
cache.mutex.Unlock()
if i%20 == 0 && i > 0 {
runtime.GC()
runtime.ReadMemStats(&m)
currentAlloc := m.Alloc
allocIncrease := float64(currentAlloc-baselineAlloc) / 1024 / 1024
t.Logf("After %d refreshes: Memory increase=%.2f MB", i, allocIncrease)
// Metadata cache should only store one copy
if allocIncrease > 3.0 {
t.Errorf("Metadata cache leak detected: %.2f MB after %d refreshes",
allocIncrease, i)
break
}
}
}
cache.Close()
wg.Wait()
runtime.GC()
runtime.ReadMemStats(&m)
finalAlloc := m.Alloc
finalAllocIncrease := float64(finalAlloc-baselineAlloc) / 1024 / 1024
t.Logf("Final memory after metadata cache closure: %.2f MB", finalAllocIncrease)
if finalAllocIncrease > 1.0 {
t.Errorf("Metadata cache not cleaned properly: %.2f MB remains", finalAllocIncrease)
}
})
}
// TestMemoryPoolLeak tests for leaks in memory pool management
func TestMemoryPoolLeak(t *testing.T) {
t.Run("Memory pool buffer leaks", func(t *testing.T) {
// Clean up any existing pools
CleanupGlobalMemoryPools()
pools := GetGlobalMemoryPools()
runtime.GC()
var m runtime.MemStats
runtime.ReadMemStats(&m)
baselineAlloc := m.Alloc
// Simulate heavy buffer usage
var buffers [][]byte
for i := 0; i < 1000; i++ {
buf := pools.GetHTTPResponseBuffer()
// Simulate using the buffer
copy(buf, []byte("test data"))
// 90% of the time, return the buffer
// 10% of the time, "forget" to return it (simulating a leak)
if i%10 != 0 {
pools.PutHTTPResponseBuffer(buf)
} else {
// Keep reference to simulate leak
buffers = append(buffers, buf)
}
if i%100 == 0 && i > 0 {
runtime.GC()
runtime.ReadMemStats(&m)
currentAlloc := m.Alloc
allocIncrease := float64(currentAlloc-baselineAlloc) / 1024 / 1024
t.Logf("After %d buffer operations: Memory increase=%.2f MB, Leaked buffers=%d",
i, allocIncrease, len(buffers))
// With proper pooling, memory should be bounded
if allocIncrease > 5.0 {
t.Errorf("Memory pool leak detected: %.2f MB after %d operations",
allocIncrease, i)
break
}
}
}
// Return the "leaked" buffers
for _, buf := range buffers {
pools.PutHTTPResponseBuffer(buf)
}
CleanupGlobalMemoryPools()
runtime.GC()
runtime.ReadMemStats(&m)
finalAlloc := m.Alloc
finalAllocIncrease := float64(finalAlloc-baselineAlloc) / 1024 / 1024
t.Logf("Final memory after pool cleanup: %.2f MB", finalAllocIncrease)
if finalAllocIncrease > 1.0 {
t.Errorf("Memory pools not cleaned properly: %.2f MB remains", finalAllocIncrease)
}
})
}
// TestConcurrentMemoryLeaks tests for memory leaks under concurrent load
func TestConcurrentMemoryLeaks(t *testing.T) {
t.Run("Concurrent operations memory stability", func(t *testing.T) {
// Reset global state
resetGlobalState()
// Set lower GC percentage to trigger GC more frequently
debug.SetGCPercent(50)
defer debug.SetGCPercent(100)
runtime.GC()
var m runtime.MemStats
runtime.ReadMemStats(&m)
baselineAlloc := m.Alloc
baselineGoroutines := runtime.NumGoroutine()
// Create a mock OIDC server
var mockServerURL string
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/.well-known/openid-configuration":
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]interface{}{
"issuer": mockServerURL,
"authorization_endpoint": mockServerURL + "/auth",
"token_endpoint": mockServerURL + "/token",
"userinfo_endpoint": mockServerURL + "/userinfo",
"jwks_uri": mockServerURL + "/keys",
})
case "/keys":
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]interface{}{
"keys": []map[string]interface{}{
{
"kty": "RSA",
"use": "sig",
"kid": "test-key",
"n": "test-modulus",
"e": "AQAB",
},
},
})
default:
w.WriteHeader(http.StatusNotFound)
}
}))
defer mockServer.Close()
mockServerURL = mockServer.URL
// Create middleware config with mock server
config := createTestConfig()
config.ProviderURL = mockServerURL
config.LogLevel = "error"
// Create middleware directly without using New() to avoid automatic metadata fetch
middleware := &TraefikOidc{
next: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) }),
providerURL: config.ProviderURL,
clientID: config.ClientID,
clientSecret: config.ClientSecret,
redirURLPath: "/callback",
scopes: []string{"openid", "email", "profile"},
logger: NewLogger(config.LogLevel),
excludedURLs: make(map[string]struct{}),
httpClient: &http.Client{Timeout: 5 * time.Second},
sessionManager: nil, // Will be created below
tokenCache: NewTokenCache(),
tokenBlacklist: NewCache(),
limiter: rate.NewLimiter(rate.Every(time.Second), 10),
goroutineWG: &sync.WaitGroup{},
firstRequestMutex: sync.Mutex{},
firstRequestReceived: false,
}
// Create session manager
var err error
middleware.sessionManager, err = NewSessionManager(
config.SessionEncryptionKey,
config.ForceHTTPS,
config.CookieDomain,
middleware.logger,
)
if err != nil {
t.Fatal(err)
}
// Simulate concurrent operations
var wg sync.WaitGroup
stopChan := make(chan struct{})
// Worker that continuously performs operations
worker := func(id int) {
defer wg.Done()
for {
select {
case <-stopChan:
return
default:
// Simulate various operations
switch id % 4 {
case 0:
// Cache operations
cache := NewCache()
cache.Set(fmt.Sprintf("key-%d", id), "value", time.Minute)
cache.Get(fmt.Sprintf("key-%d", id))
cache.Close()
case 1:
// Session operations
req := httptest.NewRequest("GET", "/", nil)
session, _ := middleware.sessionManager.GetSession(req)
if session != nil {
session.SetAccessToken("token")
session.returnToPoolSafely()
}
case 2:
// Token cache operations
cm := GetGlobalCacheManager(nil)
if cm != nil {
tc := cm.GetSharedTokenCache()
tc.Set("test-token", map[string]interface{}{"test": "data"}, time.Minute)
tc.Get("test-token")
}
case 3:
// Memory pool operations
pools := GetGlobalMemoryPools()
buf := pools.GetHTTPResponseBuffer()
pools.PutHTTPResponseBuffer(buf)
}
// Small delay to prevent tight loop
time.Sleep(10 * time.Millisecond)
}
}
}
// Start workers
numWorkers := 20
wg.Add(numWorkers)
for i := 0; i < numWorkers; i++ {
go worker(i)
}
// Let it run and measure periodically
for i := 0; i < 3; i++ {
time.Sleep(1 * time.Second)
runtime.GC()
runtime.ReadMemStats(&m)
currentAlloc := m.Alloc
currentGoroutines := runtime.NumGoroutine()
allocIncrease := float64(currentAlloc-baselineAlloc) / 1024 / 1024
goroutineIncrease := currentGoroutines - baselineGoroutines
t.Logf("After %d seconds of concurrent load: Memory increase=%.2f MB, Goroutine increase=%d",
(i + 1), allocIncrease, goroutineIncrease)
// Under sustained concurrent load, memory should stabilize
if i > 1 && allocIncrease > 20.0 {
t.Errorf("Memory leak under concurrent load: %.2f MB after %d seconds",
allocIncrease, (i + 1))
}
}
// Stop workers
close(stopChan)
wg.Wait()
// Clean up
middleware.Close()
// Final measurements
time.Sleep(500 * time.Millisecond)
runtime.GC()
runtime.ReadMemStats(&m)
finalAlloc := m.Alloc
finalGoroutines := runtime.NumGoroutine()
finalAllocIncrease := float64(finalAlloc-baselineAlloc) / 1024 / 1024
finalGoroutineIncrease := finalGoroutines - baselineGoroutines
t.Logf("Final after cleanup: Memory increase=%.2f MB, Goroutine increase=%d",
finalAllocIncrease, finalGoroutineIncrease)
if finalGoroutineIncrease > 5 {
t.Errorf("Goroutines not cleaned up after concurrent operations: %d extra remain",
finalGoroutineIncrease)
}
if finalAllocIncrease > 5.0 {
t.Errorf("Memory not released after concurrent operations: %.2f MB remains",
finalAllocIncrease)
}
})
}
+7 -4
View File
@@ -43,7 +43,7 @@ func TestGoogleOIDCRefreshTokenHandling(t *testing.T) {
}
// Create a session manager
sessionManager, _ := NewSessionManager("0123456789abcdef0123456789abcdef", true, mockLogger)
sessionManager, _ := NewSessionManager("0123456789abcdef0123456789abcdef", true, "", mockLogger)
tOidc.sessionManager = sessionManager
t.Run("Google provider detection adds required parameters", func(t *testing.T) {
@@ -349,6 +349,9 @@ func TestGoogleOIDCRefreshTokenHandling(t *testing.T) {
Err: nil,
}
// Create test cleanup helper
tc := newTestCleanup(t)
// Create a complete test instance with all required fields
mockLogger := NewLogger("debug")
googleTOidc := &TraefikOidc{
@@ -358,8 +361,8 @@ func TestGoogleOIDCRefreshTokenHandling(t *testing.T) {
logger: mockLogger,
scopes: []string{"openid", "profile", "email"},
refreshGracePeriod: 60,
tokenCache: NewTokenCache(), // Initialize tokenCache
tokenBlacklist: NewCache(), // Initialize tokenBlacklist
tokenCache: tc.addTokenCache(NewTokenCache()), // Initialize tokenCache with cleanup
tokenBlacklist: tc.addCache(NewCache()), // Initialize tokenBlacklist with cleanup
enablePKCE: false,
limiter: rate.NewLimiter(rate.Inf, 0), // No rate limiting for tests
jwkCache: mockJWKCache,
@@ -367,7 +370,7 @@ func TestGoogleOIDCRefreshTokenHandling(t *testing.T) {
}
// Create a session manager
sessionManager, _ := NewSessionManager("0123456789abcdef0123456789abcdef", true, mockLogger)
sessionManager, _ := NewSessionManager("0123456789abcdef0123456789abcdef", true, "", mockLogger)
googleTOidc.sessionManager = sessionManager
// Create a mock token verifier
+24 -5
View File
@@ -64,15 +64,37 @@ func (r *ProviderRegistry) ClearCache() {
// DetectProvider determines the most appropriate provider for a given issuer URL.
// It iterates through the registered providers and returns the first one that matches.
// Detection is based on URL patterns and other provider-specific criteria.
// Uses double-checked locking pattern to avoid race conditions while caching results.
func (r *ProviderRegistry) DetectProvider(issuerURL string) OIDCProvider {
// First check: read lock for cache lookup
r.mu.RLock()
defer r.mu.RUnlock()
if provider, found := r.cache[issuerURL]; found {
r.mu.RUnlock()
return provider
}
r.mu.RUnlock()
// Check cache first for performance
// Cache miss - acquire write lock for detection and caching
r.mu.Lock()
defer r.mu.Unlock()
// Second check: another goroutine might have cached the result while we waited for write lock
if provider, found := r.cache[issuerURL]; found {
return provider
}
// Perform detection under write lock
detectedProvider := r.detectProviderUnsafe(issuerURL)
// Cache the result (even if nil to avoid repeated expensive operations)
r.cache[issuerURL] = detectedProvider
return detectedProvider
}
// detectProviderUnsafe performs the actual provider detection logic.
// This method assumes the caller holds the appropriate lock and should not be called directly.
func (r *ProviderRegistry) detectProviderUnsafe(issuerURL string) OIDCProvider {
// Normalize issuer URL for consistent matching
normalizedURL, err := url.Parse(issuerURL)
if err != nil {
@@ -86,12 +108,10 @@ func (r *ProviderRegistry) DetectProvider(issuerURL string) OIDCProvider {
switch p.GetType() {
case ProviderTypeGoogle:
if strings.Contains(host, "accounts.google.com") {
r.cache[issuerURL] = p
return p
}
case ProviderTypeAzure:
if strings.Contains(host, "login.microsoftonline.com") || strings.Contains(host, "sts.windows.net") {
r.cache[issuerURL] = p
return p
}
}
@@ -100,7 +120,6 @@ func (r *ProviderRegistry) DetectProvider(issuerURL string) OIDCProvider {
// Fallback to the generic provider if no specific provider is detected
for _, p := range r.providers {
if p.GetType() == ProviderTypeGeneric {
r.cache[issuerURL] = p
return p
}
}
+5 -5
View File
@@ -20,7 +20,7 @@ func TestIssue53Regression(t *testing.T) {
// 3. Session cookies must be properly configured for HTTPS
// 4. CSRF token must persist through the OAuth flow
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Step 1: Initial request to protected resource
@@ -108,7 +108,7 @@ func TestIssue53Regression(t *testing.T) {
// This test ensures that once the session is properly maintained,
// token signature verification works correctly for Azure tokens
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Create authenticated session with Azure tokens
@@ -154,7 +154,7 @@ func TestIssue53Regression(t *testing.T) {
// This test verifies the redirect loop prevention mechanism
// that was added to handle authentication failures gracefully
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
req := httptest.NewRequest("GET", "http://internal/protected", nil)
@@ -183,7 +183,7 @@ func TestIssue53Regression(t *testing.T) {
// in different reverse proxy scenarios
func TestReverseProxySameSiteHandling(t *testing.T) {
t.Run("SameSite_Lax_For_HTTPS_OAuth", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// HTTPS request via proxy
@@ -211,7 +211,7 @@ func TestReverseProxySameSiteHandling(t *testing.T) {
})
t.Run("SameSite_Lax_For_HTTP", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Plain HTTP request (no proxy headers)
+39 -12
View File
@@ -17,9 +17,13 @@ import (
)
var (
replayCacheMu sync.RWMutex // Use RWMutex for better read performance
replayCache *Cache // Replace unbounded map with bounded Cache
replayCacheOnce sync.Once
replayCacheMu sync.RWMutex // Use RWMutex for better read performance
replayCache *Cache // Replace unbounded map with bounded Cache
replayCacheOnce sync.Once
replayCacheCleanupWG sync.WaitGroup
replayCacheCtx context.Context
replayCacheCancel context.CancelFunc
replayCacheCleanupMu sync.Mutex
)
// initReplayCache initializes the global replay cache for JWT ID tracking.
@@ -36,6 +40,18 @@ func initReplayCache() {
// It acquires a write lock, closes the cache, and sets it to nil
// to ensure proper cleanup during shutdown.
func cleanupReplayCache() {
// Cancel cleanup goroutine first
replayCacheCleanupMu.Lock()
if replayCacheCancel != nil {
replayCacheCancel()
replayCacheCancel = nil
replayCacheCtx = nil
}
replayCacheCleanupMu.Unlock()
// Wait for cleanup goroutine to finish
replayCacheCleanupWG.Wait()
replayCacheMu.Lock()
defer replayCacheMu.Unlock()
@@ -73,7 +89,21 @@ func getReplayCacheStats() (size int, maxSize int) {
// - ctx: Context for cancellation.
// - logger: Logger for debug output (can be nil).
func startReplayCacheCleanup(ctx context.Context, logger *Logger) {
replayCacheCleanupMu.Lock()
defer replayCacheCleanupMu.Unlock()
// Cancel any existing cleanup goroutine
if replayCacheCancel != nil {
replayCacheCancel()
replayCacheCleanupWG.Wait()
}
// Create new context for this cleanup goroutine
replayCacheCtx, replayCacheCancel = context.WithCancel(ctx)
replayCacheCleanupWG.Add(1)
go func() {
defer replayCacheCleanupWG.Done()
ticker := time.NewTicker(5 * time.Minute)
defer ticker.Stop()
@@ -91,8 +121,7 @@ func startReplayCacheCleanup(ctx context.Context, logger *Logger) {
}
replayCacheMu.RUnlock()
case <-ctx.Done():
cleanupReplayCache()
case <-replayCacheCtx.Done():
if logger != nil {
logger.Debug("Replay cache cleanup goroutine stopped due to context cancellation")
}
@@ -286,19 +315,17 @@ func (j *JWT) Verify(issuerURL, clientID string, skipReplayCheck ...bool) error
shouldSkipReplay := len(skipReplayCheck) > 0 && skipReplayCheck[0]
if jti, ok := claims["jti"].(string); ok && !shouldSkipReplay {
if j.Token == "" {
return nil
}
jtiValue, jtiOk := claims["jti"].(string)
if jtiOk && !shouldSkipReplay && jtiValue != "" {
initReplayCache()
replayCacheMu.RLock()
_, exists := replayCache.Get(jti)
_, exists := replayCache.Get(jtiValue)
replayCacheMu.RUnlock()
if exists {
return fmt.Errorf("token replay detected (jti: %s)", jti)
return fmt.Errorf("token replay detected (jti: %s)", jtiValue)
}
expFloat, ok := claims["exp"].(float64)
@@ -313,7 +340,7 @@ func (j *JWT) Verify(issuerURL, clientID string, skipReplayCheck ...bool) error
if duration > 0 {
replayCacheMu.Lock()
if replayCache != nil {
replayCache.Set(jti, true, duration)
replayCache.Set(jtiValue, true, duration)
}
replayCacheMu.Unlock()
}
+87 -38
View File
@@ -36,27 +36,27 @@ func createDefaultHTTPClient() *http.Client {
Proxy: http.ProxyFromEnvironment,
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
dialer := &net.Dialer{
Timeout: 10 * time.Second, // Connection timeout for faster failure detection
KeepAlive: 30 * time.Second, // Keep-alive interval for connection reuse
Timeout: 5 * time.Second, // Reduced connection timeout for faster failure detection
KeepAlive: 15 * time.Second, // Reduced keep-alive interval for connection reuse
}
return dialer.DialContext(ctx, network, addr)
},
ForceAttemptHTTP2: true,
TLSHandshakeTimeout: 3 * time.Second, // TLS handshake timeout
ExpectContinueTimeout: 1 * time.Second, // Timeout for 100-continue responses
MaxIdleConns: 10, // Reduced from 20 to prevent memory buildup
MaxIdleConnsPerHost: 2, // Reduced from 5 to limit per-host connections
IdleConnTimeout: 30 * time.Second, // Reduced from 60 to close idle connections faster
DisableKeepAlives: false, // Enable connection reuse
MaxConnsPerHost: 10, // Reduced from 20 to limit concurrent connections
ResponseHeaderTimeout: 5 * time.Second, // Timeout for reading response headers
DisableCompression: false, // Enable compression for bandwidth efficiency
WriteBufferSize: 4096, // Write buffer size for connections
ReadBufferSize: 4096, // Read buffer size for connections
TLSHandshakeTimeout: 2 * time.Second, // Reduced TLS handshake timeout
ExpectContinueTimeout: 1 * time.Second, // Timeout for 100-continue responses
MaxIdleConns: 2, // Reduced from 5 to minimize idle connections
MaxIdleConnsPerHost: 1, // Keep minimal idle connections per host
IdleConnTimeout: 5 * time.Second, // Reduced from 15s to close idle connections faster
DisableKeepAlives: false, // Enable connection reuse
MaxConnsPerHost: 2, // Reduced from 5 to limit concurrent connections
ResponseHeaderTimeout: 3 * time.Second, // Reduced timeout for reading response headers
DisableCompression: false, // Enable compression for bandwidth efficiency
WriteBufferSize: 4096, // Write buffer size for connections
ReadBufferSize: 4096, // Read buffer size for connections
}
return &http.Client{
Timeout: time.Second * 10, // HTTP client timeout
Timeout: time.Second * 5, // Reduced HTTP client timeout to prevent hanging connections
Transport: transport,
CheckRedirect: func(req *http.Request, via []*http.Request) error {
// Limit redirects to prevent redirect loops
@@ -124,7 +124,7 @@ type CacheManager struct {
// It initializes all cache types on first call with appropriate default settings.
// This ensures thread-safe initialization and consistent cache behavior across
// the entire application lifecycle.
func GetGlobalCacheManager() *CacheManager {
func GetGlobalCacheManager(wg *sync.WaitGroup) *CacheManager {
cacheManagerOnce.Do(func() {
globalCacheManager = &CacheManager{
tokenBlacklist: func() *Cache {
@@ -133,7 +133,7 @@ func GetGlobalCacheManager() *CacheManager {
return c
}(),
tokenCache: NewTokenCache(),
metadataCache: NewMetadataCache(),
metadataCache: NewMetadataCache(wg),
jwkCache: &JWKCache{},
}
})
@@ -283,13 +283,17 @@ type TraefikOidc struct {
issuerURL string
revocationURL string
scopes []string
goroutineWG sync.WaitGroup
goroutineWG *sync.WaitGroup
refreshGracePeriod time.Duration
shutdownOnce sync.Once
forceHTTPS bool
enablePKCE bool
overrideScopes bool
suppressDiagnosticLogs bool
firstRequestReceived bool
firstRequestMutex sync.Mutex
metadataRefreshStarted bool
providerURL string
}
// ProviderMetadata represents the OpenID Connect provider's discovery metadata.
@@ -690,13 +694,15 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
} else {
httpClient = createDefaultHTTPClient()
}
cacheManager := GetGlobalCacheManager()
goroutineWG := &sync.WaitGroup{}
cacheManager := GetGlobalCacheManager(goroutineWG)
pluginCtx, cancelFunc := context.WithCancel(context.Background())
t := &TraefikOidc{
next: next,
name: name,
goroutineWG: goroutineWG,
redirURLPath: config.CallbackURL,
logoutURLPath: func() string {
if config.LogoutURL == "" {
@@ -756,7 +762,7 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
suppressDiagnosticLogs: isTestMode(), // Suppress diagnostic logs during tests
}
t.sessionManager, _ = NewSessionManager(config.SessionEncryptionKey, config.ForceHTTPS, t.logger)
t.sessionManager, _ = NewSessionManager(config.SessionEncryptionKey, config.ForceHTTPS, config.CookieDomain, t.logger)
t.errorRecoveryManager = NewErrorRecoveryManager(t.logger)
t.extractClaimsFunc = extractClaims
// t.exchangeCodeForTokenFunc = t.exchangeCodeForToken // Removed, using interface now
@@ -771,7 +777,7 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
t.tokenVerifier = t
t.jwtVerifier = t
t.startTokenCleanup()
// Delay startTokenCleanup() until first request
t.tokenExchanger = t // Initialize the interface field to self
// Initialize and parse header templates with safe field access
@@ -815,6 +821,9 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
startReplayCacheCleanup(pluginCtx, logger)
logger.Debugf("TraefikOidc.New: Final t.scopes initialized to: %v", t.scopes)
// Store provider URL for later use
t.providerURL = config.ProviderURL
go t.initializeMetadata(config.ProviderURL)
return t, nil
@@ -851,8 +860,8 @@ func (t *TraefikOidc) initializeMetadata(providerURL string) {
t.logger.Debug("Successfully initialized provider metadata")
t.updateMetadataEndpoints(metadata)
// Start metadata refresh goroutine
go t.startMetadataRefresh(providerURL)
// Delay metadata refresh goroutine until first request
// It will be started in ServeHTTP when needed
// Only close channel on success
close(t.initComplete)
@@ -888,11 +897,19 @@ func (t *TraefikOidc) updateMetadataEndpoints(metadata *ProviderMetadata) {
func (t *TraefikOidc) startMetadataRefresh(providerURL string) {
// Use longer interval to reduce memory pressure from refresh attempts
ticker := time.NewTicker(2 * time.Hour) // Increased from 1 hour
t.goroutineWG.Add(1) // Track this goroutine
// Check if WaitGroup is initialized (it might be nil in tests)
if t.goroutineWG != nil {
t.goroutineWG.Add(1) // Track this goroutine
}
go func() {
defer t.goroutineWG.Done() // Signal completion when goroutine exits
defer ticker.Stop() // Ensure ticker is always stopped
defer func() {
if t.goroutineWG != nil {
t.goroutineWG.Done() // Signal completion when goroutine exits
}
}()
defer ticker.Stop() // Ensure ticker is always stopped
consecutiveFailures := 0
for {
@@ -1056,6 +1073,23 @@ func fetchMetadata(wellKnownURL string, httpClient *http.Client) (*ProviderMetad
// - Authentication state management
// - Header injection for authenticated requests
func (t *TraefikOidc) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
// Start background tasks on first request (except for health checks)
if !strings.HasPrefix(req.URL.Path, "/health") {
t.firstRequestMutex.Lock()
if !t.firstRequestReceived {
t.firstRequestReceived = true
t.logger.Debug("Starting background tasks on first request")
t.startTokenCleanup()
// Also start metadata refresh if not already started
if !t.metadataRefreshStarted && t.providerURL != "" {
t.metadataRefreshStarted = true
go t.startMetadataRefresh(t.providerURL)
}
}
t.firstRequestMutex.Unlock()
}
// Wait for provider metadata initialization to complete
select {
case <-t.initComplete:
@@ -1087,6 +1121,10 @@ func (t *TraefikOidc) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
return
}
// Clean up any old cookies with different domains before processing
// This prevents CSRF token mismatch errors when cookie domain configuration changes
t.sessionManager.CleanupOldCookies(rw, req)
// Retrieve or create session for request
session, err := t.sessionManager.GetSession(req)
if err != nil {
@@ -2044,11 +2082,17 @@ func (t *TraefikOidc) validateHost(host string) error {
// panic recovery to ensure stability.
func (t *TraefikOidc) startTokenCleanup() {
ticker := time.NewTicker(1 * time.Minute) // Run cleanup every minute
t.goroutineWG.Add(1) // Track this goroutine
// Check if WaitGroup is initialized (it might be nil in tests)
if t.goroutineWG != nil {
t.goroutineWG.Add(1) // Track this goroutine
}
go func() {
defer func() {
t.goroutineWG.Done() // Signal completion when goroutine exits
ticker.Stop() // Ensure ticker is always stopped
if t.goroutineWG != nil {
t.goroutineWG.Done() // Signal completion when goroutine exits
}
ticker.Stop() // Ensure ticker is always stopped
// CRITICAL: Recover from panics to prevent middleware crashes
if r := recover(); r != nil {
@@ -2834,17 +2878,22 @@ func (t *TraefikOidc) Close() error {
t.logger.Debug("metadataRefreshStopChan closed")
}
done := make(chan struct{})
go func() {
t.goroutineWG.Wait()
close(done)
}()
// Only wait for goroutines if WaitGroup is initialized
if t.goroutineWG != nil {
done := make(chan struct{})
go func() {
t.goroutineWG.Wait()
close(done)
}()
select {
case <-done:
t.logger.Debug("All background goroutines stopped gracefully")
case <-time.After(10 * time.Second):
t.logger.Errorf("Timeout waiting for background goroutines to stop")
select {
case <-done:
t.logger.Debug("All background goroutines stopped gracefully")
case <-time.After(10 * time.Second):
t.logger.Errorf("Timeout waiting for background goroutines to stop")
}
} else {
t.logger.Debug("No goroutineWG to wait for (likely in test)")
}
if t.httpClient != nil {
+3 -1
View File
@@ -10,7 +10,9 @@ import (
func BenchmarkOIDCMiddleware(b *testing.B) {
// Setup test environment
ts := &TestSuite{}
// Create a testing.T wrapper for benchmarks
t := &testing.T{}
ts := NewTestSuite(t)
ts.Setup()
ts.token = "valid.jwt.token"
+99 -99
View File
@@ -32,10 +32,25 @@ type TestSuite struct {
mockJWKCache *MockJWKCache
sessionManager *SessionManager
token string
utf *UnifiedTestFramework // Add unified test framework
}
// NewTestSuite creates a new test suite with automatic cleanup
func NewTestSuite(t *testing.T) *TestSuite {
ts := &TestSuite{
t: t,
utf: NewUnifiedTestFramework(t),
}
return ts
}
// Setup initializes the test suite
func (ts *TestSuite) Setup() {
// Initialize unified test framework if not already done
if ts.utf == nil {
ts.utf = NewUnifiedTestFramework(ts.t)
}
var err error
ts.rsaPrivateKey, err = rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
@@ -90,7 +105,10 @@ func (ts *TestSuite) Setup() {
}
logger := NewLogger("info")
ts.sessionManager, _ = NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, logger)
ts.sessionManager, _ = NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, "", logger)
// Create WaitGroup for the OIDC instance
goroutineWG := ts.utf.NewWaitGroup()
// Common TraefikOidc instance
ts.tOidc = &TraefikOidc{
@@ -101,19 +119,23 @@ func (ts *TestSuite) Setup() {
jwksURL: "https://test-jwks-url.com",
revocationURL: "https://revocation-endpoint.com",
limiter: rate.NewLimiter(rate.Every(time.Second), 10),
tokenBlacklist: NewCache(), // Use generic cache for blacklist
tokenCache: NewTokenCache(),
tokenBlacklist: ts.utf.NewCache(), // Use framework's tracked cache
tokenCache: ts.utf.NewTokenCache(), // Use framework's tracked token cache
logger: logger,
allowedUserDomains: map[string]struct{}{"example.com": {}},
excludedURLs: map[string]struct{}{"/favicon": {}},
httpClient: &http.Client{},
httpClient: ts.utf.NewHTTPClient(), // Use framework's tracked HTTP client
// Explicitly set paths as New() is bypassed
redirURLPath: "/callback", // Assume default callback path for tests
logoutURLPath: "/callback/logout", // Assume default logout path for tests
tokenURL: "https://test-issuer.com/token", // Explicitly set for refresh tests
extractClaimsFunc: extractClaims,
initComplete: make(chan struct{}),
sessionManager: ts.sessionManager,
redirURLPath: "/callback", // Assume default callback path for tests
logoutURLPath: "/callback/logout", // Assume default logout path for tests
tokenURL: "https://test-issuer.com/token", // Explicitly set for refresh tests
extractClaimsFunc: extractClaims,
initComplete: make(chan struct{}),
sessionManager: ts.sessionManager,
goroutineWG: goroutineWG,
ctx: ts.utf.NewContext(), // Use framework's tracked context
tokenCleanupStopChan: make(chan struct{}),
metadataRefreshStopChan: make(chan struct{}),
}
close(ts.tOidc.initComplete)
// ts.tOidc.exchangeCodeForTokenFunc = ts.exchangeCodeForTokenFunc // Removed
@@ -139,6 +161,9 @@ func (ts *TestSuite) Setup() {
return nil
},
}
// Track the OIDC instance for cleanup
ts.utf.AddOIDCInstance(ts.tOidc)
}
// Helper function exchangeCodeForTokenFunc removed as it's unused after refactoring to TokenExchanger interface.
@@ -272,7 +297,7 @@ func bigIntToBytes(i *big.Int) []byte {
// TestVerifyToken tests the VerifyToken method
func TestVerifyToken(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
tests := []struct {
@@ -361,7 +386,7 @@ func TestVerifyToken(t *testing.T) {
// TestServeHTTP tests the ServeHTTP method
func TestServeHTTP(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -739,10 +764,8 @@ func TestServeHTTP(t *testing.T) {
ts.tOidc.tokenCache = NewTokenCache()
// Reset the global replayCache to prevent "token replay detected" errors
replayCacheMu.Lock()
replayCache = NewCache()
replayCache.SetMaxSize(10000)
replayCacheMu.Unlock()
cleanupReplayCache()
initReplayCache()
// Store original tokenVerifier to restore later
origTokenVerifier := ts.tOidc.tokenVerifier
@@ -752,10 +775,8 @@ func TestServeHTTP(t *testing.T) {
mockTokenVerifier := &MockTokenVerifier{
VerifyFunc: func(token string) error {
// Clear replay cache before token verification
replayCacheMu.Lock()
replayCache = NewCache()
replayCache.SetMaxSize(10000)
replayCacheMu.Unlock()
cleanupReplayCache()
initReplayCache()
// Call the original verifier's VerifyToken method
// Ensure origTokenVerifier is not nil and is the correct type if necessary,
@@ -869,7 +890,7 @@ func TestServeHTTP(t *testing.T) {
}
func TestJWKToPEM(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
tests := []struct {
@@ -922,7 +943,7 @@ func TestJWKToPEM(t *testing.T) {
}
func TestParseJWT(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
tests := []struct {
@@ -963,7 +984,7 @@ func TestParseJWT(t *testing.T) {
}
func TestJWTVerify_MissingClaims(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
jwt := &JWT{
@@ -983,7 +1004,7 @@ func TestJWTVerify_MissingClaims(t *testing.T) {
}
func TestHandleCallback(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
redirectURL := "http://example.com/"
@@ -1162,10 +1183,8 @@ func TestHandleCallback(t *testing.T) {
// Capture range variable
t.Run(tc.name, func(t *testing.T) {
// Clear the global replay cache before each test run
replayCacheMu.Lock()
replayCache = NewCache()
replayCache.SetMaxSize(10000)
replayCacheMu.Unlock()
cleanupReplayCache()
initReplayCache()
// Explicitly clear the shared blacklist at the start of each sub-test
// to ensure no state leaks, even though we expect the local one to be used.
@@ -1173,7 +1192,7 @@ func TestHandleCallback(t *testing.T) {
ts.tOidc.tokenBlacklist = NewCache() // Use generic cache for blacklist
logger := NewLogger("info")
sessionManager, _ := NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, logger)
sessionManager, _ := NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, "", logger)
// Create a new instance for each test to avoid state carryover
instanceExtractClaimsFunc := tc.extractClaimsFunc
@@ -1252,7 +1271,7 @@ func TestHandleCallback(t *testing.T) {
}
func TestIsAllowedDomain(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
tests := []struct {
@@ -1337,7 +1356,7 @@ func TestIsAllowedDomain(t *testing.T) {
}
func TestOIDCHandler(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
ts.token = "valid.jwt.token"
@@ -1481,7 +1500,7 @@ func TestOIDCHandler(t *testing.T) {
// TestHandleLogout tests the logout functionality
func TestHandleLogout(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create mock revocation endpoint server
@@ -1561,7 +1580,7 @@ func TestHandleLogout(t *testing.T) {
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
logger := NewLogger("info")
sessionManager, _ := NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, logger)
sessionManager, _ := NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, "", logger)
tOidc := &TraefikOidc{
revocationURL: mockRevocationServer.URL,
endSessionURL: tc.endSessionURL,
@@ -1652,7 +1671,7 @@ func TestHandleLogout(t *testing.T) {
// TestRevokeTokenWithProvider tests the token revocation with provider
func TestRevokeTokenWithProvider(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
tests := []struct {
@@ -1728,7 +1747,7 @@ func TestRevokeTokenWithProvider(t *testing.T) {
// TestRevokeToken tests the token revocation functionality
func TestRevokeToken(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
token := "test.token.with.claims"
@@ -1820,7 +1839,7 @@ func TestBuildLogoutURL(t *testing.T) {
// Add this new test function
func TestHandleExpiredToken(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
tests := []struct {
@@ -1863,7 +1882,7 @@ func TestHandleExpiredToken(t *testing.T) {
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
logger := NewLogger("info")
sessionManager, _ := NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, logger)
sessionManager, _ := NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, "", logger)
tOidc := &TraefikOidc{
sessionManager: sessionManager,
@@ -1928,7 +1947,7 @@ func TestHandleExpiredToken(t *testing.T) {
// Add this new test function
func TestExtractGroupsAndRoles(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
tests := []struct {
@@ -2094,7 +2113,7 @@ func TestMultipleMiddlewareInstances(t *testing.T) {
}
func TestServeHTTPRolesAndGroups(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create consistent timestamps for all test cases
@@ -2309,7 +2328,7 @@ func stringSliceEqual(a, b []string) bool {
// TestExchangeTokensWithRedirects tests the token exchange process with redirects
func TestExchangeTokensWithRedirects(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
tests := []struct {
@@ -2414,7 +2433,7 @@ func TestExchangeTokensWithRedirects(t *testing.T) {
// TestBuildAuthURL tests the buildAuthURL function with various URL scenarios
func TestBuildAuthURL(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
tests := []struct {
@@ -2570,7 +2589,7 @@ func TestBuildAuthURL(t *testing.T) {
// TestExchangeCodeForToken tests the exchangeCodeForToken function with PKCE support
func TestExchangeCodeForToken(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
tests := []struct {
@@ -2688,7 +2707,7 @@ func TestExchangeCodeForToken(t *testing.T) {
// TestDefaultInitiateAuthentication_PreservesQueryParameters tests that defaultInitiateAuthentication preserves query parameters in the incoming path.
func TestDefaultInitiateAuthentication_PreservesQueryParameters(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create a request with query parameters
@@ -2845,14 +2864,12 @@ func TestVerifyTimeConstraint(t *testing.T) {
// TestJWTVerifyWithSkipReplayCheck tests the new skipReplayCheck parameter functionality
func TestJWTVerifyWithSkipReplayCheck(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Clear the global replay cache before test
replayCacheMu.Lock()
replayCache = NewCache()
replayCache.SetMaxSize(10000)
replayCacheMu.Unlock()
cleanupReplayCache()
initReplayCache()
// Create a test JWT with unique JTI
jti := generateRandomString(16)
@@ -2913,10 +2930,8 @@ func TestJWTVerifyWithSkipReplayCheck(t *testing.T) {
t.Run(tc.name, func(t *testing.T) {
if tc.firstCall {
// Clear replay cache for first call tests
replayCacheMu.Lock()
replayCache = NewCache()
replayCache.SetMaxSize(10000)
replayCacheMu.Unlock()
cleanupReplayCache()
initReplayCache()
}
err := jwt.Verify("https://test-issuer.com", "test-client-id", tc.skipReplayCheck)
@@ -2938,14 +2953,12 @@ func TestJWTVerifyWithSkipReplayCheck(t *testing.T) {
// TestJWTVerifyBackwardCompatibility tests that calls without the skipReplayCheck parameter default to replay checking
func TestJWTVerifyBackwardCompatibility(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Clear the global replay cache
replayCacheMu.Lock()
replayCache = NewCache()
replayCache.SetMaxSize(10000)
replayCacheMu.Unlock()
cleanupReplayCache()
initReplayCache()
// Create a test JWT with unique JTI
jti := generateRandomString(16)
@@ -2991,14 +3004,12 @@ func TestJWTVerifyBackwardCompatibility(t *testing.T) {
// TestTokenReplayDetectionFalsePositiveFix tests the specific scenario that was causing false positives
func TestTokenReplayDetectionFalsePositiveFix(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Clear the global replay cache
replayCacheMu.Lock()
replayCache = NewCache()
replayCache.SetMaxSize(10000)
replayCacheMu.Unlock()
cleanupReplayCache()
initReplayCache()
// Create a test JWT with unique JTI
jti := generateRandomString(16)
@@ -3064,14 +3075,12 @@ func TestTokenReplayDetectionFalsePositiveFix(t *testing.T) {
// TestAuthenticationFlowReplayDetection tests the complete authentication flow
func TestAuthenticationFlowReplayDetection(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Clear the global replay cache
replayCacheMu.Lock()
replayCache = NewCache()
replayCache.SetMaxSize(10000)
replayCacheMu.Unlock()
cleanupReplayCache()
initReplayCache()
// Create a test JWT with unique JTI
jti := generateRandomString(16)
@@ -3139,14 +3148,12 @@ func TestAuthenticationFlowReplayDetection(t *testing.T) {
// TestActualReplayAttackDetection ensures real replay attacks are still properly detected
func TestActualReplayAttackDetection(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Clear the global replay cache
replayCacheMu.Lock()
replayCache = NewCache()
replayCache.SetMaxSize(10000)
replayCacheMu.Unlock()
cleanupReplayCache()
initReplayCache()
// Create a test JWT with unique JTI
jti := generateRandomString(16)
@@ -3217,17 +3224,15 @@ func TestActualReplayAttackDetection(t *testing.T) {
// TestConcurrentTokenValidation tests thread safety of replay detection
func TestConcurrentTokenValidation(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Configure rate limiter to allow more requests for concurrent testing
ts.tOidc.limiter = rate.NewLimiter(rate.Limit(1000), 1000) // Allow 1000 requests per second with burst of 1000
// Clear the global replay cache
replayCacheMu.Lock()
replayCache = NewCache()
replayCache.SetMaxSize(10000)
replayCacheMu.Unlock()
cleanupReplayCache()
initReplayCache()
// Create multiple tokens with unique JTIs
var tokens []string
@@ -3306,17 +3311,16 @@ func TestConcurrentTokenValidation(t *testing.T) {
// TestJTIBlacklistBehavior tests the JTI blacklist cache management
func TestJTIBlacklistBehavior(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Clear the global replay cache
replayCacheMu.Lock()
replayCache = NewCache()
replayCache.SetMaxSize(10000)
replayCacheMu.Unlock()
// Properly reinitialize the global replay cache
cleanupReplayCache() // Clean up any existing cache and reset sync.Once
initReplayCache() // Initialize new cache through proper channel
// Create a test JWT with unique JTI
jti := generateRandomString(16)
t.Logf("TestJTIBlacklistBehavior - JTI: %s", jti)
now := time.Now()
exp := now.Add(1 * time.Hour).Unix()
iat := now.Unix()
@@ -3355,8 +3359,8 @@ func TestJTIBlacklistBehavior(t *testing.T) {
{
name: "JTI exists in blacklist after verification",
action: func() error {
replayCacheMu.Lock()
defer replayCacheMu.Unlock()
replayCacheMu.RLock()
defer replayCacheMu.RUnlock()
if _, exists := replayCache.Get(jti); !exists {
return fmt.Errorf("JTI not found in blacklist cache")
}
@@ -3406,14 +3410,12 @@ func TestJTIBlacklistBehavior(t *testing.T) {
// TestSessionBasedTokenRevalidation tests token revalidation in session-based scenarios
func TestSessionBasedTokenRevalidation(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Clear the global replay cache
replayCacheMu.Lock()
replayCache = NewCache()
replayCache.SetMaxSize(10000)
replayCacheMu.Unlock()
cleanupReplayCache()
initReplayCache()
// Create a test JWT with unique JTI
jti := generateRandomString(16)
@@ -3480,14 +3482,12 @@ func TestSessionBasedTokenRevalidation(t *testing.T) {
// TestEdgeCasesWithDifferentTokenTypes tests replay detection with different token types
func TestEdgeCasesWithDifferentTokenTypes(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Clear the global replay cache
replayCacheMu.Lock()
replayCache = NewCache()
replayCache.SetMaxSize(10000)
replayCacheMu.Unlock()
// Properly reinitialize the global replay cache
cleanupReplayCache() // Clean up any existing cache and reset sync.Once
initReplayCache() // Initialize new cache through proper channel
now := time.Now()
exp := now.Add(1 * time.Hour).Unix()
@@ -3969,7 +3969,7 @@ func TestNewWithScopeAppending(t *testing.T) {
// TestBuildAuthURLWithMergedScopes tests that the auth URL includes the properly merged scopes
func TestBuildAuthURLWithMergedScopes(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
tests := []struct {
@@ -4026,7 +4026,7 @@ func TestBuildAuthURLWithMergedScopes(t *testing.T) {
// TestBuildAuthURL_OverrideScopes_And_OfflineAccess tests the offline_access logic in buildAuthURL
// considering the overrideScopes flag.
func TestBuildAuthURL_OverrideScopes_And_OfflineAccess(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup() // Sets up ts.tOidc
tests := []struct {
@@ -4164,7 +4164,7 @@ func TestBuildAuthURL_OverrideScopes_And_OfflineAccess(t *testing.T) {
// TestBuildAuthURL_SpecificUserCase tests the buildAuthURL function with the specific user-reported scenario.
func TestBuildAuthURL_SpecificUserCase(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup() // Basic setup for tOidc
// Configure the TraefikOidc instance for the specific scenario
+111
View File
@@ -0,0 +1,111 @@
package traefikoidc
import (
"container/list"
"net/http"
"sync"
"time"
)
// LazyBackgroundTask wraps BackgroundTask to start only when needed
type LazyBackgroundTask struct {
*BackgroundTask
started bool
startOnce sync.Once
}
// NewLazyBackgroundTask creates a background task that doesn't start immediately
func NewLazyBackgroundTask(name string, interval time.Duration, taskFunc func(), logger *Logger, wg ...*sync.WaitGroup) *LazyBackgroundTask {
return &LazyBackgroundTask{
BackgroundTask: NewBackgroundTask(name, interval, taskFunc, logger, wg...),
started: false,
}
}
// StartIfNeeded starts the task only if it hasn't been started yet
func (lt *LazyBackgroundTask) StartIfNeeded() {
lt.startOnce.Do(func() {
if !lt.started {
lt.BackgroundTask.Start()
lt.started = true
}
})
}
// Stop stops the task if it was started
func (lt *LazyBackgroundTask) Stop() {
if lt.started {
lt.BackgroundTask.Stop()
lt.started = false
lt.startOnce = sync.Once{} // Reset for potential restart
}
}
// NewLazyCacheWithLogger creates a cache that doesn't start cleanup until first use
func NewLazyCacheWithLogger(logger *Logger) *Cache {
if logger == nil {
logger = newNoOpLogger()
}
c := &Cache{
items: make(map[string]CacheItem, DefaultMaxSize),
order: list.New(),
elems: make(map[string]*list.Element, DefaultMaxSize),
maxSize: DefaultMaxSize,
autoCleanupInterval: 10 * time.Minute, // Increased from 5 minutes
logger: logger,
}
// Don't start cleanup immediately - it will be started on first use
return c
}
// NewLazyCache creates a cache that doesn't start cleanup immediately
func NewLazyCache() *Cache {
return NewLazyCacheWithLogger(nil)
}
// CleanupIdleConnections periodically closes idle HTTP connections
func CleanupIdleConnections(client *http.Client, interval time.Duration, stopChan <-chan struct{}) {
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
select {
case <-ticker.C:
if transport, ok := client.Transport.(*http.Transport); ok {
transport.CloseIdleConnections()
}
case <-stopChan:
// Final cleanup
if transport, ok := client.Transport.(*http.Transport); ok {
transport.CloseIdleConnections()
}
return
}
}
}
// OptimizedMiddlewareConfig provides configuration for memory-optimized middleware
type OptimizedMiddlewareConfig struct {
// DelayBackgroundTasks delays starting background tasks until first request
DelayBackgroundTasks bool
// ReducedCleanupIntervals uses longer intervals for cleanup tasks
ReducedCleanupIntervals bool
// AggressiveConnectionCleanup closes idle connections more aggressively
AggressiveConnectionCleanup bool
// MinimalCacheSize uses smaller default cache sizes
MinimalCacheSize bool
}
// DefaultOptimizedConfig returns a configuration optimized for low memory usage
func DefaultOptimizedConfig() *OptimizedMiddlewareConfig {
return &OptimizedMiddlewareConfig{
DelayBackgroundTasks: true,
ReducedCleanupIntervals: true,
AggressiveConnectionCleanup: true,
MinimalCacheSize: true,
}
}
+3 -3
View File
@@ -48,7 +48,7 @@ func TestMemoryLeakFixes(t *testing.T) {
t.Run("Global cache manager cleanup", func(t *testing.T) {
// Get the global cache manager
cm := GetGlobalCacheManager()
cm := GetGlobalCacheManager(nil)
if cm == nil {
t.Fatal("Failed to get global cache manager")
}
@@ -64,7 +64,7 @@ func TestMemoryLeakFixes(t *testing.T) {
}
// Verify it can be re-initialized
cm2 := GetGlobalCacheManager()
cm2 := GetGlobalCacheManager(nil)
if cm2 == nil {
t.Fatal("Failed to re-initialize global cache manager")
}
@@ -72,7 +72,7 @@ func TestMemoryLeakFixes(t *testing.T) {
t.Run("Session pool returns properly", func(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("test-encryption-key-that-is-long-enough-32bytes", false, logger)
sm, err := NewSessionManager("test-encryption-key-that-is-long-enough-32bytes", false, "", logger)
if err != nil {
t.Fatal(err)
}
+9 -5
View File
@@ -19,23 +19,27 @@ type MetadataCache struct {
logger *Logger
autoCleanupInterval time.Duration
mutex sync.RWMutex
wg *sync.WaitGroup
stopChan chan struct{}
}
// NewMetadataCache creates a new MetadataCache instance.
// It initializes the cache structure and starts the background cleanup task.
func NewMetadataCache() *MetadataCache {
return NewMetadataCacheWithLogger(nil)
func NewMetadataCache(wg *sync.WaitGroup) *MetadataCache {
return NewMetadataCacheWithLogger(wg, nil)
}
// NewMetadataCacheWithLogger creates a new MetadataCache with a specified logger.
func NewMetadataCacheWithLogger(logger *Logger) *MetadataCache {
func NewMetadataCacheWithLogger(wg *sync.WaitGroup, logger *Logger) *MetadataCache {
if logger == nil {
logger = newNoOpLogger()
}
c := &MetadataCache{
autoCleanupInterval: 5 * time.Minute,
autoCleanupInterval: 30 * time.Minute, // Increased from 5 minutes since metadata changes rarely
logger: logger,
wg: wg,
stopChan: make(chan struct{}),
}
c.startAutoCleanup()
return c
@@ -200,7 +204,7 @@ func (c *MetadataCache) GetMetadata(providerURL string, httpClient *http.Client,
// startAutoCleanup starts the background task that periodically calls Cleanup
// to remove expired metadata from the cache.
func (c *MetadataCache) startAutoCleanup() {
c.cleanupTask = NewBackgroundTask("metadata-cache-cleanup", c.autoCleanupInterval, c.Cleanup, c.logger)
c.cleanupTask = NewBackgroundTask("metadata-cache-cleanup", c.autoCleanupInterval, c.Cleanup, c.logger, c.wg)
c.cleanupTask.Start()
}
+852
View File
@@ -0,0 +1,852 @@
package traefikoidc
import (
"bytes"
"fmt"
"net/http"
"runtime"
"runtime/pprof"
"sync"
"time"
)
// MemoryProfiler defines the interface for memory profiling operations
type MemoryProfiler interface {
// TakeSnapshot captures current memory statistics
TakeSnapshot() (*MemorySnapshot, error)
// StartProfiling begins memory profiling with specified configuration
StartProfiling(config ProfilingConfig) error
// StopProfiling ends memory profiling and returns final snapshot
StopProfiling() (*MemorySnapshot, error)
// GetCurrentStats returns current runtime memory statistics
GetCurrentStats() *runtime.MemStats
// AnalyzeLeaks performs leak detection analysis
AnalyzeLeaks(baseline, current *MemorySnapshot) *LeakAnalysis
}
// MemorySnapshot represents a point-in-time capture of memory statistics
type MemorySnapshot struct {
Timestamp time.Time
RuntimeStats runtime.MemStats
HeapProfile []byte
GoroutineProfile []byte
CustomMetrics map[string]interface{}
}
// LeakAnalysis contains the results of memory leak detection
type LeakAnalysis struct {
HasLeak bool
LeakDescription string
MemoryIncrease uint64
GoroutineIncrease int
SuspectedLeaks []string
Recommendations []string
}
// ProfilingManager coordinates memory profiling operations
type ProfilingManager struct {
mu sync.RWMutex
isProfiling bool
startTime time.Time
baselineSnapshot *MemorySnapshot
config ProfilingConfig
logger *Logger
profilers map[string]MemoryProfiler
}
// ProfilingConfig contains configuration for profiling operations
type ProfilingConfig struct {
EnableHeapProfiling bool
EnableGoroutineProfiling bool
SnapshotInterval time.Duration
LeakThresholdMB uint64
MaxSnapshots int
EnableContinuousMonitoring bool
MonitoringInterval time.Duration
}
// LeakDetectionConfig contains configuration for leak detection
type LeakDetectionConfig struct {
EnableLeakDetection bool
LeakThresholdMB uint64
GoroutineLeakThreshold int
SessionPoolThreshold int
CacheMemoryThreshold uint64
HTTPClientThreshold int
TokenCompressionThreshold uint64
}
// NewProfilingManager creates a new profiling manager instance
func NewProfilingManager(logger *Logger) *ProfilingManager {
if logger == nil {
logger = newNoOpLogger()
}
return &ProfilingManager{
profilers: make(map[string]MemoryProfiler),
config: ProfilingConfig{
EnableHeapProfiling: true,
EnableGoroutineProfiling: true,
SnapshotInterval: 30 * time.Second,
LeakThresholdMB: 50, // 50MB
MaxSnapshots: 100,
EnableContinuousMonitoring: true,
MonitoringInterval: 60 * time.Second,
},
logger: logger,
}
}
// TakeSnapshot captures current memory statistics
func (pm *ProfilingManager) TakeSnapshot() (*MemorySnapshot, error) {
var buf bytes.Buffer
snapshot := &MemorySnapshot{
Timestamp: time.Now(),
CustomMetrics: make(map[string]interface{}),
}
// Capture runtime memory statistics
runtime.ReadMemStats(&snapshot.RuntimeStats)
// Capture heap profile if enabled
if pm.config.EnableHeapProfiling {
if err := pprof.WriteHeapProfile(&buf); err != nil {
pm.logger.Errorf("Failed to capture heap profile: %v", err)
} else {
snapshot.HeapProfile = make([]byte, buf.Len())
copy(snapshot.HeapProfile, buf.Bytes())
buf.Reset()
}
}
// Capture goroutine profile if enabled
if pm.config.EnableGoroutineProfiling {
if err := pprof.Lookup("goroutine").WriteTo(&buf, 0); err != nil {
pm.logger.Errorf("Failed to capture goroutine profile: %v", err)
} else {
snapshot.GoroutineProfile = make([]byte, buf.Len())
copy(snapshot.GoroutineProfile, buf.Bytes())
buf.Reset()
}
}
// Capture custom metrics from registered profilers
pm.mu.RLock()
for name, profiler := range pm.profilers {
if customStats := profiler.GetCurrentStats(); customStats != nil {
snapshot.CustomMetrics[name] = customStats
}
}
pm.mu.RUnlock()
return snapshot, nil
}
// StartProfiling begins memory profiling with specified configuration
func (pm *ProfilingManager) StartProfiling(config ProfilingConfig) error {
pm.mu.Lock()
defer pm.mu.Unlock()
if pm.isProfiling {
return fmt.Errorf("profiling already in progress")
}
pm.config = config
pm.isProfiling = true
pm.startTime = time.Now()
// Take baseline snapshot
baseline, err := pm.TakeSnapshot()
if err != nil {
pm.isProfiling = false
return fmt.Errorf("failed to take baseline snapshot: %w", err)
}
pm.baselineSnapshot = baseline
pm.logger.Infof("Memory profiling started at %v", pm.startTime)
return nil
}
// StopProfiling ends memory profiling and returns final snapshot
func (pm *ProfilingManager) StopProfiling() (*MemorySnapshot, error) {
pm.mu.Lock()
defer pm.mu.Unlock()
if !pm.isProfiling {
return nil, fmt.Errorf("profiling not in progress")
}
// Take final snapshot
finalSnapshot, err := pm.TakeSnapshot()
if err != nil {
pm.logger.Errorf("Failed to take final snapshot: %v", err)
}
pm.isProfiling = false
duration := time.Since(pm.startTime)
pm.logger.Infof("Memory profiling stopped after %v", duration)
return finalSnapshot, err
}
// GetCurrentStats returns current runtime memory statistics
func (pm *ProfilingManager) GetCurrentStats() *runtime.MemStats {
stats := &runtime.MemStats{}
runtime.ReadMemStats(stats)
return stats
}
// AnalyzeLeaks performs leak detection analysis
func (pm *ProfilingManager) AnalyzeLeaks(baseline, current *MemorySnapshot) *LeakAnalysis {
analysis := &LeakAnalysis{
SuspectedLeaks: make([]string, 0),
Recommendations: make([]string, 0),
}
if baseline == nil || current == nil {
analysis.HasLeak = false
analysis.LeakDescription = "Insufficient data for leak analysis"
return analysis
}
// Calculate memory increase
memoryIncrease := current.RuntimeStats.Alloc - baseline.RuntimeStats.Alloc
analysis.MemoryIncrease = memoryIncrease
// Calculate goroutine increase
currentGoroutines := runtime.NumGoroutine()
baselineGoroutines := runtime.NumGoroutine() // Note: This is not accurate for baseline, but we don't have historical data
goroutineIncrease := currentGoroutines - baselineGoroutines
analysis.GoroutineIncrease = goroutineIncrease
// Check for memory leaks
memoryThreshold := pm.config.LeakThresholdMB * 1024 * 1024 // Convert MB to bytes
if memoryIncrease > memoryThreshold {
analysis.HasLeak = true
analysis.SuspectedLeaks = append(analysis.SuspectedLeaks,
fmt.Sprintf("Memory usage increased by %.2f MB", float64(memoryIncrease)/(1024*1024)))
analysis.Recommendations = append(analysis.Recommendations,
"Consider checking for unreleased memory pools or growing caches")
}
// Check for goroutine leaks
if goroutineIncrease > 10 { // Arbitrary threshold
analysis.HasLeak = true
analysis.SuspectedLeaks = append(analysis.SuspectedLeaks,
fmt.Sprintf("Goroutine count increased by %d", goroutineIncrease))
analysis.Recommendations = append(analysis.Recommendations,
"Check for goroutines that are not being properly cleaned up")
}
if analysis.HasLeak {
analysis.LeakDescription = fmt.Sprintf("Potential memory leak detected: %s",
fmt.Sprintf("%.2f MB increase, %d goroutines", float64(memoryIncrease)/(1024*1024), goroutineIncrease))
} else {
analysis.LeakDescription = "No significant memory leaks detected"
}
return analysis
}
// RegisterProfiler registers a component-specific profiler
func (pm *ProfilingManager) RegisterProfiler(name string, profiler MemoryProfiler) {
pm.mu.Lock()
defer pm.mu.Unlock()
pm.profilers[name] = profiler
pm.logger.Debugf("Registered profiler: %s", name)
}
// UnregisterProfiler removes a component-specific profiler
func (pm *ProfilingManager) UnregisterProfiler(name string) {
pm.mu.Lock()
defer pm.mu.Unlock()
delete(pm.profilers, name)
pm.logger.Debugf("Unregistered profiler: %s", name)
}
// GetRegisteredProfilers returns list of registered profiler names
func (pm *ProfilingManager) GetRegisteredProfilers() []string {
pm.mu.RLock()
defer pm.mu.RUnlock()
names := make([]string, 0, len(pm.profilers))
for name := range pm.profilers {
names = append(names, name)
}
return names
}
// MemoryTestOrchestrator coordinates memory leak testing across components
type MemoryTestOrchestrator struct {
mu sync.RWMutex
profilers map[string]MemoryProfiler
config LeakDetectionConfig
logger *Logger
isRunning bool
stopChan chan struct{}
testResults map[string]*LeakAnalysis
}
// NewMemoryTestOrchestrator creates a new test orchestrator
func NewMemoryTestOrchestrator(config LeakDetectionConfig, logger *Logger) *MemoryTestOrchestrator {
if logger == nil {
logger = newNoOpLogger()
}
return &MemoryTestOrchestrator{
profilers: make(map[string]MemoryProfiler),
config: config,
logger: logger,
stopChan: make(chan struct{}),
testResults: make(map[string]*LeakAnalysis),
}
}
// RegisterComponent registers a component for memory leak testing
func (mto *MemoryTestOrchestrator) RegisterComponent(name string, profiler MemoryProfiler) {
mto.mu.Lock()
defer mto.mu.Unlock()
mto.profilers[name] = profiler
mto.logger.Debugf("Registered component for leak testing: %s", name)
}
// UnregisterComponent removes a component from leak testing
func (mto *MemoryTestOrchestrator) UnregisterComponent(name string) {
mto.mu.Lock()
defer mto.mu.Unlock()
delete(mto.profilers, name)
delete(mto.testResults, name)
mto.logger.Debugf("Unregistered component from leak testing: %s", name)
}
// StartLeakDetection begins continuous leak detection monitoring
func (mto *MemoryTestOrchestrator) StartLeakDetection() error {
mto.mu.Lock()
defer mto.mu.Unlock()
if mto.isRunning {
return fmt.Errorf("leak detection already running")
}
if !mto.config.EnableLeakDetection {
return fmt.Errorf("leak detection is disabled in configuration")
}
mto.isRunning = true
go mto.runLeakDetection()
mto.logger.Infof("Memory leak detection started")
return nil
}
// StopLeakDetection stops continuous leak detection monitoring
func (mto *MemoryTestOrchestrator) StopLeakDetection() error {
mto.mu.Lock()
defer mto.mu.Unlock()
if !mto.isRunning {
return fmt.Errorf("leak detection not running")
}
mto.isRunning = false
close(mto.stopChan)
mto.stopChan = make(chan struct{}) // Reset for potential restart
mto.logger.Infof("Memory leak detection stopped")
return nil
}
// runLeakDetection performs continuous leak detection monitoring
func (mto *MemoryTestOrchestrator) runLeakDetection() {
ticker := time.NewTicker(5 * time.Minute) // Check every 5 minutes
defer ticker.Stop()
baselineSnapshots := make(map[string]*MemorySnapshot)
// Take initial baseline snapshots
mto.mu.RLock()
for name, profiler := range mto.profilers {
if snapshot, err := profiler.TakeSnapshot(); err == nil {
baselineSnapshots[name] = snapshot
}
}
mto.mu.RUnlock()
for {
select {
case <-ticker.C:
mto.performLeakCheck(baselineSnapshots)
case <-mto.stopChan:
return
}
}
}
// performLeakCheck performs leak detection for all registered components
func (mto *MemoryTestOrchestrator) performLeakCheck(baselineSnapshots map[string]*MemorySnapshot) {
mto.mu.RLock()
defer mto.mu.RUnlock()
for name, profiler := range mto.profilers {
baseline, exists := baselineSnapshots[name]
if !exists {
continue
}
current, err := profiler.TakeSnapshot()
if err != nil {
mto.logger.Errorf("Failed to take snapshot for component %s: %v", name, err)
continue
}
analysis := profiler.AnalyzeLeaks(baseline, current)
if analysis.HasLeak {
mto.logger.Errorf("Memory leak detected in component %s: %s", name, analysis.LeakDescription)
for _, rec := range analysis.Recommendations {
mto.logger.Errorf("Recommendation for %s: %s", name, rec)
}
}
mto.testResults[name] = analysis
}
}
// GetLeakAnalysis returns leak analysis for a specific component
func (mto *MemoryTestOrchestrator) GetLeakAnalysis(componentName string) (*LeakAnalysis, bool) {
mto.mu.RLock()
defer mto.mu.RUnlock()
analysis, exists := mto.testResults[componentName]
return analysis, exists
}
// GetAllLeakAnalyses returns leak analyses for all components
func (mto *MemoryTestOrchestrator) GetAllLeakAnalyses() map[string]*LeakAnalysis {
mto.mu.RLock()
defer mto.mu.RUnlock()
results := make(map[string]*LeakAnalysis)
for name, analysis := range mto.testResults {
results[name] = analysis
}
return results
}
// Component-specific profiler implementations
// SessionPoolProfiler monitors session pool memory usage
type SessionPoolProfiler struct {
sessionManager *SessionManager
logger *Logger
}
// NewSessionPoolProfiler creates a new session pool profiler
func NewSessionPoolProfiler(sm *SessionManager, logger *Logger) *SessionPoolProfiler {
if logger == nil {
logger = newNoOpLogger()
}
return &SessionPoolProfiler{
sessionManager: sm,
logger: logger,
}
}
// TakeSnapshot captures session pool memory statistics
func (spp *SessionPoolProfiler) TakeSnapshot() (*MemorySnapshot, error) {
snapshot := &MemorySnapshot{
Timestamp: time.Now(),
CustomMetrics: make(map[string]interface{}),
}
// Capture runtime stats
runtime.ReadMemStats(&snapshot.RuntimeStats)
// Add session pool specific metrics
snapshot.CustomMetrics["session_pool_metrics"] = spp.sessionManager.GetSessionMetrics()
return snapshot, nil
}
// StartProfiling begins profiling (no-op for session pools)
func (spp *SessionPoolProfiler) StartProfiling(config ProfilingConfig) error {
return nil
}
// StopProfiling ends profiling (no-op for session pools)
func (spp *SessionPoolProfiler) StopProfiling() (*MemorySnapshot, error) {
return spp.TakeSnapshot()
}
// GetCurrentStats returns current memory statistics
func (spp *SessionPoolProfiler) GetCurrentStats() *runtime.MemStats {
stats := &runtime.MemStats{}
runtime.ReadMemStats(stats)
return stats
}
// AnalyzeLeaks analyzes session pool for leaks
func (spp *SessionPoolProfiler) AnalyzeLeaks(baseline, current *MemorySnapshot) *LeakAnalysis {
analysis := &LeakAnalysis{
SuspectedLeaks: make([]string, 0),
Recommendations: make([]string, 0),
}
if baseline == nil || current == nil {
analysis.LeakDescription = "Insufficient session pool data"
return analysis
}
// Check for session pool specific leaks
memoryIncrease := current.RuntimeStats.Alloc - baseline.RuntimeStats.Alloc
if memoryIncrease > 10*1024*1024 { // 10MB threshold
analysis.HasLeak = true
analysis.SuspectedLeaks = append(analysis.SuspectedLeaks,
"Session pool memory usage increased significantly")
analysis.Recommendations = append(analysis.Recommendations,
"Check for sessions not being returned to pool properly")
}
return analysis
}
// CacheMemoryProfiler monitors cache memory usage
type CacheMemoryProfiler struct {
cache *Cache
logger *Logger
}
// NewCacheMemoryProfiler creates a new cache memory profiler
func NewCacheMemoryProfiler(cache *Cache, logger *Logger) *CacheMemoryProfiler {
if logger == nil {
logger = newNoOpLogger()
}
return &CacheMemoryProfiler{
cache: cache,
logger: logger,
}
}
// TakeSnapshot captures cache memory statistics
func (cmp *CacheMemoryProfiler) TakeSnapshot() (*MemorySnapshot, error) {
snapshot := &MemorySnapshot{
Timestamp: time.Now(),
CustomMetrics: make(map[string]interface{}),
}
runtime.ReadMemStats(&snapshot.RuntimeStats)
// Add cache-specific metrics (would need to be added to Cache struct)
snapshot.CustomMetrics["cache_size"] = "unknown" // Placeholder
return snapshot, nil
}
// StartProfiling begins profiling (no-op for cache)
func (cmp *CacheMemoryProfiler) StartProfiling(config ProfilingConfig) error {
return nil
}
// StopProfiling ends profiling
func (cmp *CacheMemoryProfiler) StopProfiling() (*MemorySnapshot, error) {
return cmp.TakeSnapshot()
}
// GetCurrentStats returns current memory statistics
func (cmp *CacheMemoryProfiler) GetCurrentStats() *runtime.MemStats {
stats := &runtime.MemStats{}
runtime.ReadMemStats(stats)
return stats
}
// AnalyzeLeaks analyzes cache for memory leaks
func (cmp *CacheMemoryProfiler) AnalyzeLeaks(baseline, current *MemorySnapshot) *LeakAnalysis {
analysis := &LeakAnalysis{
SuspectedLeaks: make([]string, 0),
Recommendations: make([]string, 0),
}
if baseline == nil || current == nil {
analysis.LeakDescription = "Insufficient cache data"
return analysis
}
memoryIncrease := current.RuntimeStats.Alloc - baseline.RuntimeStats.Alloc
if memoryIncrease > 20*1024*1024 { // 20MB threshold for cache
analysis.HasLeak = true
analysis.SuspectedLeaks = append(analysis.SuspectedLeaks,
"Cache memory usage increased significantly")
analysis.Recommendations = append(analysis.Recommendations,
"Check cache size limits and cleanup intervals")
}
return analysis
}
// HTTPClientProfiler monitors HTTP client connection pools
type HTTPClientProfiler struct {
httpClient *http.Client
logger *Logger
}
// NewHTTPClientProfiler creates a new HTTP client profiler
func NewHTTPClientProfiler(client *http.Client, logger *Logger) *HTTPClientProfiler {
if logger == nil {
logger = newNoOpLogger()
}
return &HTTPClientProfiler{
httpClient: client,
logger: logger,
}
}
// TakeSnapshot captures HTTP client memory statistics
func (hcp *HTTPClientProfiler) TakeSnapshot() (*MemorySnapshot, error) {
snapshot := &MemorySnapshot{
Timestamp: time.Now(),
CustomMetrics: make(map[string]interface{}),
}
runtime.ReadMemStats(&snapshot.RuntimeStats)
// Add HTTP client specific metrics
if transport, ok := hcp.httpClient.Transport.(*http.Transport); ok {
snapshot.CustomMetrics["idle_connections"] = transport.IdleConnTimeout.String()
snapshot.CustomMetrics["max_idle_conns"] = transport.MaxIdleConns
}
return snapshot, nil
}
// StartProfiling begins profiling (no-op for HTTP client)
func (hcp *HTTPClientProfiler) StartProfiling(config ProfilingConfig) error {
return nil
}
// StopProfiling ends profiling
func (hcp *HTTPClientProfiler) StopProfiling() (*MemorySnapshot, error) {
return hcp.TakeSnapshot()
}
// GetCurrentStats returns current memory statistics
func (hcp *HTTPClientProfiler) GetCurrentStats() *runtime.MemStats {
stats := &runtime.MemStats{}
runtime.ReadMemStats(stats)
return stats
}
// AnalyzeLeaks analyzes HTTP client for connection leaks
func (hcp *HTTPClientProfiler) AnalyzeLeaks(baseline, current *MemorySnapshot) *LeakAnalysis {
analysis := &LeakAnalysis{
SuspectedLeaks: make([]string, 0),
Recommendations: make([]string, 0),
}
if baseline == nil || current == nil {
analysis.LeakDescription = "Insufficient HTTP client data"
return analysis
}
memoryIncrease := current.RuntimeStats.Alloc - baseline.RuntimeStats.Alloc
if memoryIncrease > 5*1024*1024 { // 5MB threshold for HTTP client
analysis.HasLeak = true
analysis.SuspectedLeaks = append(analysis.SuspectedLeaks,
"HTTP client memory usage increased significantly")
analysis.Recommendations = append(analysis.Recommendations,
"Check for HTTP response bodies not being drained properly")
}
return analysis
}
// TokenCompressionProfiler monitors token compression memory usage
type TokenCompressionProfiler struct {
compressionPool *TokenCompressionPool
logger *Logger
}
// NewTokenCompressionProfiler creates a new token compression profiler
func NewTokenCompressionProfiler(pool *TokenCompressionPool, logger *Logger) *TokenCompressionProfiler {
if logger == nil {
logger = newNoOpLogger()
}
return &TokenCompressionProfiler{
compressionPool: pool,
logger: logger,
}
}
// TakeSnapshot captures token compression memory statistics
func (tcp *TokenCompressionProfiler) TakeSnapshot() (*MemorySnapshot, error) {
snapshot := &MemorySnapshot{
Timestamp: time.Now(),
CustomMetrics: make(map[string]interface{}),
}
runtime.ReadMemStats(&snapshot.RuntimeStats)
// Add compression pool specific metrics
snapshot.CustomMetrics["compression_pool_active"] = true
return snapshot, nil
}
// StartProfiling begins profiling (no-op for compression)
func (tcp *TokenCompressionProfiler) StartProfiling(config ProfilingConfig) error {
return nil
}
// StopProfiling ends profiling
func (tcp *TokenCompressionProfiler) StopProfiling() (*MemorySnapshot, error) {
return tcp.TakeSnapshot()
}
// GetCurrentStats returns current memory statistics
func (tcp *TokenCompressionProfiler) GetCurrentStats() *runtime.MemStats {
stats := &runtime.MemStats{}
runtime.ReadMemStats(stats)
return stats
}
// AnalyzeLeaks analyzes token compression for memory leaks
func (tcp *TokenCompressionProfiler) AnalyzeLeaks(baseline, current *MemorySnapshot) *LeakAnalysis {
analysis := &LeakAnalysis{
SuspectedLeaks: make([]string, 0),
Recommendations: make([]string, 0),
}
if baseline == nil || current == nil {
analysis.LeakDescription = "Insufficient compression data"
return analysis
}
memoryIncrease := current.RuntimeStats.Alloc - baseline.RuntimeStats.Alloc
if memoryIncrease > 2*1024*1024 { // 2MB threshold for compression
analysis.HasLeak = true
analysis.SuspectedLeaks = append(analysis.SuspectedLeaks,
"Token compression memory usage increased significantly")
analysis.Recommendations = append(analysis.Recommendations,
"Check for compression buffers not being returned to pool")
}
return analysis
}
// MemoryPoolProfiler monitors memory pool usage and detects leaks
type MemoryPoolProfiler struct {
memoryPoolManager *MemoryPoolManager
tokenCompressionPool *TokenCompressionPool
logger *Logger
}
// NewMemoryPoolProfiler creates a new memory pool profiler
func NewMemoryPoolProfiler(memoryPoolManager *MemoryPoolManager, tokenCompressionPool *TokenCompressionPool, logger *Logger) *MemoryPoolProfiler {
if logger == nil {
logger = newNoOpLogger()
}
return &MemoryPoolProfiler{
memoryPoolManager: memoryPoolManager,
tokenCompressionPool: tokenCompressionPool,
logger: logger,
}
}
// TakeSnapshot captures memory pool statistics
func (mpp *MemoryPoolProfiler) TakeSnapshot() (*MemorySnapshot, error) {
snapshot := &MemorySnapshot{
Timestamp: time.Now(),
CustomMetrics: make(map[string]interface{}),
}
// Capture runtime stats
runtime.ReadMemStats(&snapshot.RuntimeStats)
// Add memory pool metrics
if mpp.memoryPoolManager != nil {
snapshot.CustomMetrics["memory_pool_active"] = true
// Note: sync.Pool doesn't expose internal statistics, so we track usage patterns
}
if mpp.tokenCompressionPool != nil {
snapshot.CustomMetrics["token_compression_pool_active"] = true
}
return snapshot, nil
}
// StartProfiling begins profiling (no-op for memory pools)
func (mpp *MemoryPoolProfiler) StartProfiling(config ProfilingConfig) error {
return nil
}
// StopProfiling ends profiling
func (mpp *MemoryPoolProfiler) StopProfiling() (*MemorySnapshot, error) {
return mpp.TakeSnapshot()
}
// GetCurrentStats returns current memory statistics
func (mpp *MemoryPoolProfiler) GetCurrentStats() *runtime.MemStats {
stats := &runtime.MemStats{}
runtime.ReadMemStats(stats)
return stats
}
// AnalyzeLeaks analyzes memory pools for leaks
func (mpp *MemoryPoolProfiler) AnalyzeLeaks(baseline, current *MemorySnapshot) *LeakAnalysis {
analysis := &LeakAnalysis{
SuspectedLeaks: make([]string, 0),
Recommendations: make([]string, 0),
}
if baseline == nil || current == nil {
analysis.LeakDescription = "Insufficient memory pool data"
return analysis
}
// Check for memory leaks in pool operations
memoryIncrease := current.RuntimeStats.Alloc - baseline.RuntimeStats.Alloc
if memoryIncrease > 5*1024*1024 { // 5MB threshold for pool operations
analysis.HasLeak = true
analysis.SuspectedLeaks = append(analysis.SuspectedLeaks,
"Memory pool operations caused significant memory increase")
analysis.Recommendations = append(analysis.Recommendations,
"Check for objects not being returned to memory pools properly")
}
return analysis
}
// Global profiling manager instance
var globalProfilingManager *ProfilingManager
var profilingManagerOnce sync.Once
// GetGlobalProfilingManager returns the singleton profiling manager
func GetGlobalProfilingManager() *ProfilingManager {
profilingManagerOnce.Do(func() {
globalProfilingManager = NewProfilingManager(nil)
})
return globalProfilingManager
}
// Global test orchestrator instance
var globalTestOrchestrator *MemoryTestOrchestrator
var testOrchestratorOnce sync.Once
// GetGlobalTestOrchestrator returns the singleton test orchestrator
func GetGlobalTestOrchestrator() *MemoryTestOrchestrator {
testOrchestratorOnce.Do(func() {
config := LeakDetectionConfig{
EnableLeakDetection: true,
LeakThresholdMB: 50,
GoroutineLeakThreshold: 10,
SessionPoolThreshold: 100,
CacheMemoryThreshold: 20 * 1024 * 1024, // 20MB
HTTPClientThreshold: 50,
TokenCompressionThreshold: 2 * 1024 * 1024, // 2MB
}
globalTestOrchestrator = NewMemoryTestOrchestrator(config, nil)
})
return globalTestOrchestrator
}
+788
View File
@@ -0,0 +1,788 @@
package traefikoidc
import (
"encoding/json"
"fmt"
"net"
"net/http"
"os"
"runtime"
"sync"
"testing"
"time"
)
func TestProfilingManager(t *testing.T) {
logger := NewLogger("debug")
pm := NewProfilingManager(logger)
// Test taking a snapshot
snapshot, err := pm.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to take snapshot: %v", err)
}
if snapshot == nil {
t.Fatal("Snapshot is nil")
}
if snapshot.RuntimeStats.Alloc == 0 {
t.Error("Runtime stats Alloc should not be zero")
}
if snapshot.Timestamp.IsZero() {
t.Error("Snapshot timestamp should not be zero")
}
}
func TestMemoryTestOrchestrator(t *testing.T) {
logger := NewLogger("debug")
config := LeakDetectionConfig{
EnableLeakDetection: true,
LeakThresholdMB: 10,
}
mto := NewMemoryTestOrchestrator(config, logger)
// Test registering a component
sessionManager, err := NewSessionManager("test-key-32-chars-long-for-testing", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
profiler := NewSessionPoolProfiler(sessionManager, logger)
mto.RegisterComponent("session_pool", profiler)
// Test getting leak analysis (should return false initially since no checks have been performed)
_, exists := mto.GetLeakAnalysis("session_pool")
if exists {
t.Error("Should not have leak analysis before any checks are performed")
}
// Perform a manual leak check
baseline, err := profiler.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to take baseline snapshot: %v", err)
}
time.Sleep(10 * time.Millisecond) // Small delay
// Manually trigger leak check with baseline
baselineSnapshots := make(map[string]*MemorySnapshot)
baselineSnapshots["session_pool"] = baseline
mto.performLeakCheck(baselineSnapshots)
// Now test getting leak analysis
analysis, exists := mto.GetLeakAnalysis("session_pool")
if !exists {
t.Error("Should have leak analysis after performing checks")
}
if analysis == nil {
t.Error("Leak analysis should not be nil after checks")
}
}
func TestComponentProfilers(t *testing.T) {
logger := NewLogger("debug")
// Test Session Pool Profiler
sessionManager, err := NewSessionManager("test-key-32-chars-long-for-testing", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
spp := NewSessionPoolProfiler(sessionManager, logger)
snapshot, err := spp.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to take session pool snapshot: %v", err)
}
if snapshot == nil {
t.Fatal("Session pool snapshot is nil")
}
// Test Cache Memory Profiler
cache := NewCache()
cmp := NewCacheMemoryProfiler(cache, logger)
snapshot, err = cmp.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to take cache snapshot: %v", err)
}
if snapshot == nil {
t.Fatal("Cache snapshot is nil")
}
// Test HTTP Client Profiler
httpClient := createDefaultHTTPClient()
hcp := NewHTTPClientProfiler(httpClient, logger)
snapshot, err = hcp.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to take HTTP client snapshot: %v", err)
}
if snapshot == nil {
t.Fatal("HTTP client snapshot is nil")
}
// Test Token Compression Profiler
compressionPool := NewTokenCompressionPool()
tcp := NewTokenCompressionProfiler(compressionPool, logger)
snapshot, err = tcp.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to take compression snapshot: %v", err)
}
if snapshot == nil {
t.Fatal("Compression snapshot is nil")
}
}
func TestLeakAnalysis(t *testing.T) {
logger := NewLogger("debug")
pm := NewProfilingManager(logger)
// Create baseline snapshot
baseline, err := pm.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to create baseline: %v", err)
}
// Wait a bit and create current snapshot
time.Sleep(10 * time.Millisecond)
current, err := pm.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to create current snapshot: %v", err)
}
// Test leak analysis
analysis := pm.AnalyzeLeaks(baseline, current)
if analysis == nil {
t.Fatal("Leak analysis is nil")
}
// Analysis should not have leaks for normal operation
if analysis.HasLeak {
t.Logf("Leak detected: %s", analysis.LeakDescription)
// This is acceptable as the test environment may have varying memory usage
}
}
func TestGlobalInstances(t *testing.T) {
// Test global profiling manager
gpm := GetGlobalProfilingManager()
if gpm == nil {
t.Fatal("Global profiling manager is nil")
}
// Test global test orchestrator
gto := GetGlobalTestOrchestrator()
if gto == nil {
t.Fatal("Global test orchestrator is nil")
}
// Test that they're singletons
gpm2 := GetGlobalProfilingManager()
if gpm != gpm2 {
t.Error("Global profiling manager should be singleton")
}
gto2 := GetGlobalTestOrchestrator()
if gto != gto2 {
t.Error("Global test orchestrator should be singleton")
}
}
func TestProfilingConfig(t *testing.T) {
config := ProfilingConfig{
EnableHeapProfiling: true,
EnableGoroutineProfiling: true,
SnapshotInterval: 30 * time.Second,
LeakThresholdMB: 50,
MaxSnapshots: 100,
EnableContinuousMonitoring: true,
MonitoringInterval: 60 * time.Second,
}
if !config.EnableHeapProfiling {
t.Error("Heap profiling should be enabled")
}
if !config.EnableGoroutineProfiling {
t.Error("Goroutine profiling should be enabled")
}
if config.LeakThresholdMB != 50 {
t.Errorf("Expected leak threshold 50, got %d", config.LeakThresholdMB)
}
}
func TestLeakDetectionConfig(t *testing.T) {
config := LeakDetectionConfig{
EnableLeakDetection: true,
LeakThresholdMB: 50,
GoroutineLeakThreshold: 10,
SessionPoolThreshold: 100,
CacheMemoryThreshold: 20 * 1024 * 1024,
HTTPClientThreshold: 50,
TokenCompressionThreshold: 2 * 1024 * 1024,
}
if !config.EnableLeakDetection {
t.Error("Leak detection should be enabled")
}
if config.LeakThresholdMB != 50 {
t.Errorf("Expected leak threshold 50, got %d", config.LeakThresholdMB)
}
if config.CacheMemoryThreshold != 20*1024*1024 {
t.Errorf("Expected cache threshold 20MB, got %d", config.CacheMemoryThreshold)
}
}
// ProviderMetadataProfiler monitors provider metadata fetching and caching operations
type ProviderMetadataProfiler struct {
metadataCache *MetadataCache
httpClient *http.Client
logger *Logger
providerURL string
}
// NewProviderMetadataProfiler creates a new provider metadata profiler
func NewProviderMetadataProfiler(metadataCache *MetadataCache, httpClient *http.Client, providerURL string, logger *Logger) *ProviderMetadataProfiler {
if logger == nil {
logger = newNoOpLogger()
}
return &ProviderMetadataProfiler{
metadataCache: metadataCache,
httpClient: httpClient,
providerURL: providerURL,
logger: logger,
}
}
// TakeSnapshot captures current memory statistics for metadata operations
func (pmp *ProviderMetadataProfiler) TakeSnapshot() (*MemorySnapshot, error) {
snapshot := &MemorySnapshot{
Timestamp: time.Now(),
CustomMetrics: make(map[string]interface{}),
}
// Capture runtime memory statistics
runtime.ReadMemStats(&snapshot.RuntimeStats)
// Add metadata-specific metrics
snapshot.CustomMetrics["metadata_cache_size"] = 1 // Placeholder for cache size
snapshot.CustomMetrics["metadata_fetch_count"] = 0 // Placeholder for fetch count
snapshot.CustomMetrics["background_goroutines"] = runtime.NumGoroutine()
return snapshot, nil
}
// StartProfiling begins profiling (no-op for metadata profiler)
func (pmp *ProviderMetadataProfiler) StartProfiling(config ProfilingConfig) error {
return nil
}
// StopProfiling ends profiling
func (pmp *ProviderMetadataProfiler) StopProfiling() (*MemorySnapshot, error) {
return pmp.TakeSnapshot()
}
// GetCurrentStats returns current memory statistics
func (pmp *ProviderMetadataProfiler) GetCurrentStats() *runtime.MemStats {
stats := &runtime.MemStats{}
runtime.ReadMemStats(stats)
return stats
}
// AnalyzeLeaks analyzes metadata operations for memory leaks
func (pmp *ProviderMetadataProfiler) AnalyzeLeaks(baseline, current *MemorySnapshot) *LeakAnalysis {
analysis := &LeakAnalysis{
SuspectedLeaks: make([]string, 0),
Recommendations: make([]string, 0),
}
if baseline == nil || current == nil {
analysis.LeakDescription = "Insufficient metadata data"
return analysis
}
// Check for memory leaks
memoryIncrease := current.RuntimeStats.Alloc - baseline.RuntimeStats.Alloc
if memoryIncrease > 5*1024*1024 { // 5MB threshold for metadata operations
analysis.HasLeak = true
analysis.SuspectedLeaks = append(analysis.SuspectedLeaks,
"Metadata operations memory usage increased significantly")
analysis.Recommendations = append(analysis.Recommendations,
"Check for metadata cache not being cleaned up properly")
}
// Check for goroutine leaks
goroutineIncrease := current.CustomMetrics["background_goroutines"].(int) - baseline.CustomMetrics["background_goroutines"].(int)
if goroutineIncrease > 2 { // Allow some variance
analysis.HasLeak = true
analysis.SuspectedLeaks = append(analysis.SuspectedLeaks,
fmt.Sprintf("Goroutine count increased by %d during metadata operations", goroutineIncrease))
analysis.Recommendations = append(analysis.Recommendations,
"Check for background goroutines not being cleaned up")
}
return analysis
}
// TestProviderMetadataMemoryLeakDetection tests for memory leaks in provider metadata operations
func TestProviderMetadataMemoryLeakDetection(t *testing.T) {
logger := NewLogger("debug")
strictMode := os.Getenv("STRICT_MEMORY_TEST") == "true"
if strictMode {
t.Log("Running in strict memory test mode - will fail on detected leaks")
} else {
t.Log("Running in lenient memory test mode - will log warnings instead of failing")
}
config := LeakDetectionConfig{
EnableLeakDetection: true,
LeakThresholdMB: 10,
}
mto := NewMemoryTestOrchestrator(config, logger)
// Create mock HTTP server for metadata endpoint with failure simulation
requestCount := 0
serverFailures := 0
mockServer := &http.Server{
Addr: "localhost:0", // Let system assign port
Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requestCount++
if r.URL.Path == "/.well-known/openid-configuration" {
// Simulate occasional failures to test cache extension
if requestCount%4 == 0 { // Fail every 4th request
serverFailures++
w.WriteHeader(http.StatusInternalServerError)
return
}
metadata := ProviderMetadata{
Issuer: "https://mock-provider.com",
AuthURL: "https://mock-provider.com/auth",
TokenURL: "https://mock-provider.com/token",
JWKSURL: "https://mock-provider.com/jwks",
RevokeURL: "https://mock-provider.com/revoke",
EndSessionURL: "https://mock-provider.com/logout",
}
w.Header().Set("Content-Type", "application/json")
w.Header().Set("Cache-Control", "max-age=3600") // 1 hour cache hint
json.NewEncoder(w).Encode(metadata)
} else {
http.NotFound(w, r)
}
}),
}
// Start mock server
listener, err := net.Listen("tcp", "localhost:0")
if err != nil {
t.Fatalf("Failed to create listener: %v", err)
}
go mockServer.Serve(listener)
defer mockServer.Close()
providerURL := fmt.Sprintf("http://%s", listener.Addr().String())
httpClient := createDefaultHTTPClient()
// Create metadata cache with WaitGroup for proper goroutine synchronization
var wg sync.WaitGroup
metadataCache := NewMetadataCacheWithLogger(&wg, logger)
// Create profiler
profiler := NewProviderMetadataProfiler(metadataCache, httpClient, providerURL, logger)
mto.RegisterComponent("provider_metadata", profiler)
// Take initial baseline
baseline, err := profiler.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to take baseline snapshot: %v", err)
}
initialGoroutines := runtime.NumGoroutine()
// Phase 1: Simulate periodic metadata fetching with some failures
t.Log("Phase 1: Testing periodic fetching with occasional failures...")
for i := 0; i < 20; i++ {
_, err := metadataCache.GetMetadata(providerURL, httpClient, logger)
if err != nil {
t.Logf("Metadata fetch %d failed (expected for cache extension testing): %v", i+1, err)
} else {
t.Logf("Metadata fetch %d succeeded", i+1)
}
time.Sleep(100 * time.Millisecond)
}
// Wait for background cleanup (normally every 5 minutes)
time.Sleep(300 * time.Millisecond)
// Take intermediate snapshot
intermediate, err := profiler.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to take intermediate snapshot: %v", err)
}
// Phase 2: Continue with more fetches to test sustained operation
t.Log("Phase 2: Testing sustained operation with 1000 iterations...")
for i := 20; i < 1020; i++ {
_, err := metadataCache.GetMetadata(providerURL, httpClient, logger)
if err != nil {
t.Logf("Metadata fetch %d failed: %v", i+1, err)
}
time.Sleep(50 * time.Millisecond) // Reduced sleep for faster execution
}
// Take final snapshot
current, err := profiler.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to take current snapshot: %v", err)
}
finalGoroutines := runtime.NumGoroutine()
// Analyze for leaks
analysis := profiler.AnalyzeLeaks(baseline, current)
// Assertions for memory leaks
if analysis.HasLeak {
if strictMode {
t.Errorf("Memory leak detected in provider metadata operations: %s", analysis.LeakDescription)
for _, leak := range analysis.SuspectedLeaks {
t.Errorf("Suspected leak: %s", leak)
}
} else {
t.Logf("Memory leak warning in provider metadata operations: %s", analysis.LeakDescription)
for _, leak := range analysis.SuspectedLeaks {
t.Logf("Suspected leak: %s", leak)
}
}
for _, rec := range analysis.Recommendations {
t.Logf("Recommendation: %s", rec)
}
}
// Check total memory growth
totalMemoryIncrease := current.RuntimeStats.Alloc - baseline.RuntimeStats.Alloc
if totalMemoryIncrease > 20*1024*1024 { // 20MB threshold for entire test
if strictMode {
t.Errorf("Total memory usage increased by %.2f MB during metadata operations", float64(totalMemoryIncrease)/(1024*1024))
} else {
t.Logf("Total memory usage increased by %.2f MB during metadata operations", float64(totalMemoryIncrease)/(1024*1024))
}
}
// Check for gradual memory growth patterns
intermediateMemoryIncrease := intermediate.RuntimeStats.Alloc - baseline.RuntimeStats.Alloc
if intermediateMemoryIncrease > 10*1024*1024 { // 10MB threshold for first phase
if strictMode {
t.Errorf("Memory usage increased by %.2f MB during first phase of metadata operations", float64(intermediateMemoryIncrease)/(1024*1024))
} else {
t.Logf("Memory usage increased by %.2f MB during first phase of metadata operations", float64(intermediateMemoryIncrease)/(1024*1024))
}
}
// Check goroutine count stability
goroutineIncrease := finalGoroutines - initialGoroutines
if goroutineIncrease > 5 { // Allow some variance for test environment
if strictMode {
t.Errorf("Goroutine count increased by %d during metadata operations (initial: %d, final: %d)",
goroutineIncrease, initialGoroutines, finalGoroutines)
} else {
t.Logf("Goroutine count increased by %d during metadata operations (initial: %d, final: %d)",
goroutineIncrease, initialGoroutines, finalGoroutines)
}
}
// Phase 3: Test cache extension behavior on persistent failures
t.Log("Phase 3: Testing cache extension on persistent failures...")
// Stop mock server to simulate provider unavailability
mockServer.Close()
// Try multiple fetches after server shutdown
postShutdownFailures := 0
for i := 0; i < 5; i++ {
_, err = metadataCache.GetMetadata(providerURL, httpClient, logger)
if err != nil {
postShutdownFailures++
t.Logf("Expected failure %d after server shutdown: %v", i+1, err)
} else {
t.Logf("Unexpected success %d after server shutdown - cache extension working", i+1)
}
time.Sleep(200 * time.Millisecond)
}
if postShutdownFailures == 0 {
if strictMode {
t.Error("Expected some metadata fetches to fail after server shutdown")
} else {
t.Log("Warning: No metadata fetches failed after server shutdown - cache extension may not be working as expected")
}
}
// Phase 4: Test background goroutine lifecycle and cleanup
t.Log("Phase 4: Testing background goroutine lifecycle...")
// Wait longer to allow background cleanup to run
time.Sleep(1 * time.Second)
// Take final snapshot after cleanup
finalAfterCleanup, err := profiler.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to take final snapshot after cleanup: %v", err)
}
// Check if memory decreased after cleanup
if finalAfterCleanup.RuntimeStats.Alloc < current.RuntimeStats.Alloc {
memoryDecrease := current.RuntimeStats.Alloc - finalAfterCleanup.RuntimeStats.Alloc
t.Logf("Memory decreased by %.2f MB after cleanup phase", float64(memoryDecrease)/(1024*1024))
}
// Clean up resources
metadataCache.Close()
wg.Wait() // Ensure all background goroutines complete
t.Logf("Test completed: %d total requests, %d server failures, %d post-shutdown failures",
requestCount, serverFailures, postShutdownFailures)
t.Logf("Memory usage: baseline=%.2f MB, intermediate=%.2f MB, final=%.2f MB",
float64(baseline.RuntimeStats.Alloc)/(1024*1024),
float64(intermediate.RuntimeStats.Alloc)/(1024*1024),
float64(current.RuntimeStats.Alloc)/(1024*1024))
}
// TestMemoryPoolLeakDetection tests for memory leaks in memory pool operations
func TestMemoryPoolLeakDetection(t *testing.T) {
logger := NewLogger("debug")
strictMode := os.Getenv("STRICT_MEMORY_TEST") == "true"
if strictMode {
t.Log("Running in strict memory test mode - will fail on detected leaks")
} else {
t.Log("Running in lenient memory test mode - will log warnings instead of failing")
}
config := LeakDetectionConfig{
EnableLeakDetection: true,
LeakThresholdMB: 10,
}
mto := NewMemoryTestOrchestrator(config, logger)
// Create memory pool manager and token compression pool
memoryPoolManager := NewMemoryPoolManager()
tokenCompressionPool := NewTokenCompressionPool()
// Create profiler for memory pools
profiler := NewMemoryPoolProfiler(memoryPoolManager, tokenCompressionPool, logger)
mto.RegisterComponent("memory_pools", profiler)
// Take initial baseline
baseline, err := profiler.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to take baseline snapshot: %v", err)
}
initialGoroutines := runtime.NumGoroutine()
// Phase 1: Simulate various memory pool operations
t.Log("Phase 1: Testing memory pool operations with various patterns...")
// Test compression buffer pool
for i := 0; i < 100; i++ {
buf := memoryPoolManager.GetCompressionBuffer()
// Simulate some work with the buffer
buf.WriteString(fmt.Sprintf("test data %d", i))
// Properly return buffer to pool
memoryPoolManager.PutCompressionBuffer(buf)
}
// Test JWT parsing buffer pool
for i := 0; i < 50; i++ {
jwtBuf := memoryPoolManager.GetJWTParsingBuffer()
// Simulate JWT parsing operations
jwtBuf.HeaderBuf = append(jwtBuf.HeaderBuf, []byte("header")...)
jwtBuf.PayloadBuf = append(jwtBuf.PayloadBuf, []byte("payload")...)
jwtBuf.SignatureBuf = append(jwtBuf.SignatureBuf, []byte("signature")...)
// Properly return buffer to pool
memoryPoolManager.PutJWTParsingBuffer(jwtBuf)
}
// Test HTTP response buffer pool
for i := 0; i < 75; i++ {
httpBuf := memoryPoolManager.GetHTTPResponseBuffer()
// Simulate HTTP response processing
copy(httpBuf[:min(len(httpBuf), 100)], []byte("http response data"))
// Properly return buffer to pool
memoryPoolManager.PutHTTPResponseBuffer(httpBuf)
}
// Test string builder pool
for i := 0; i < 60; i++ {
sb := memoryPoolManager.GetStringBuilder()
// Simulate string building operations
sb.WriteString(fmt.Sprintf("built string %d", i))
_ = sb.String() // Use the result
// Properly return string builder to pool
memoryPoolManager.PutStringBuilder(sb)
}
// Test token compression pool
for i := 0; i < 40; i++ {
compBuf := tokenCompressionPool.GetCompressionBuffer()
// Simulate compression operations
compBuf.WriteString(fmt.Sprintf("compress data %d", i))
// Properly return buffer to pool
tokenCompressionPool.PutCompressionBuffer(compBuf)
decompBuf := tokenCompressionPool.GetDecompressionBuffer()
// Simulate decompression operations
decompBuf.WriteString(fmt.Sprintf("decompress data %d", i))
// Properly return buffer to pool
tokenCompressionPool.PutDecompressionBuffer(decompBuf)
sb := tokenCompressionPool.GetStringBuilder()
// Simulate string operations
sb.WriteString(fmt.Sprintf("token string %d", i))
_ = sb.String()
// Properly return string builder to pool
tokenCompressionPool.PutStringBuilder(sb)
}
// Take intermediate snapshot
intermediate, err := profiler.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to take intermediate snapshot: %v", err)
}
// Phase 2: Continue with more intensive operations to test sustained usage
t.Log("Phase 2: Testing sustained memory pool usage...")
// Simulate mixed operations with varying patterns
for i := 0; i < 200; i++ {
// Mix different pool operations
switch i % 4 {
case 0:
buf := memoryPoolManager.GetCompressionBuffer()
buf.WriteString("mixed operation data")
memoryPoolManager.PutCompressionBuffer(buf)
case 1:
jwtBuf := memoryPoolManager.GetJWTParsingBuffer()
jwtBuf.HeaderBuf = append(jwtBuf.HeaderBuf, []byte("mixed")...)
memoryPoolManager.PutJWTParsingBuffer(jwtBuf)
case 2:
httpBuf := memoryPoolManager.GetHTTPResponseBuffer()
copy(httpBuf[:min(len(httpBuf), 50)], []byte("mixed http"))
memoryPoolManager.PutHTTPResponseBuffer(httpBuf)
case 3:
sb := memoryPoolManager.GetStringBuilder()
sb.WriteString("mixed string building")
_ = sb.String()
memoryPoolManager.PutStringBuilder(sb)
}
}
// Take final snapshot
current, err := profiler.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to take current snapshot: %v", err)
}
finalGoroutines := runtime.NumGoroutine()
// Analyze for leaks
analysis := profiler.AnalyzeLeaks(baseline, current)
// Assertions for memory leaks
if analysis.HasLeak {
if strictMode {
t.Errorf("Memory leak detected in memory pool operations: %s", analysis.LeakDescription)
for _, leak := range analysis.SuspectedLeaks {
t.Errorf("Suspected leak: %s", leak)
}
} else {
t.Logf("Memory leak warning in memory pool operations: %s", analysis.LeakDescription)
for _, leak := range analysis.SuspectedLeaks {
t.Logf("Suspected leak: %s", leak)
}
}
for _, rec := range analysis.Recommendations {
t.Logf("Recommendation: %s", rec)
}
}
// Check total memory growth
totalMemoryIncrease := current.RuntimeStats.Alloc - baseline.RuntimeStats.Alloc
if totalMemoryIncrease > 15*1024*1024 { // 15MB threshold for entire test
if strictMode {
t.Errorf("Total memory usage increased by %.2f MB during memory pool operations", float64(totalMemoryIncrease)/(1024*1024))
} else {
t.Logf("Total memory usage increased by %.2f MB during memory pool operations", float64(totalMemoryIncrease)/(1024*1024))
}
}
// Check for gradual memory growth patterns
intermediateMemoryIncrease := intermediate.RuntimeStats.Alloc - baseline.RuntimeStats.Alloc
if intermediateMemoryIncrease > 8*1024*1024 { // 8MB threshold for first phase
if strictMode {
t.Errorf("Memory usage increased by %.2f MB during first phase of memory pool operations", float64(intermediateMemoryIncrease)/(1024*1024))
} else {
t.Logf("Memory usage increased by %.2f MB during first phase of memory pool operations", float64(intermediateMemoryIncrease)/(1024*1024))
}
}
// Check goroutine count stability
goroutineIncrease := finalGoroutines - initialGoroutines
if goroutineIncrease > 3 { // Allow small variance for test environment
if strictMode {
t.Errorf("Goroutine count increased by %d during memory pool operations (initial: %d, final: %d)",
goroutineIncrease, initialGoroutines, finalGoroutines)
} else {
t.Logf("Goroutine count increased by %d during memory pool operations (initial: %d, final: %d)",
goroutineIncrease, initialGoroutines, finalGoroutines)
}
}
// Phase 3: Test cleanup verification
t.Log("Phase 3: Testing cleanup verification...")
// Force garbage collection to see if pools are properly managed
runtime.GC()
runtime.GC() // Run twice to ensure cleanup
time.Sleep(10 * time.Millisecond) // Allow cleanup to complete
// Take post-cleanup snapshot
postCleanup, err := profiler.TakeSnapshot()
if err != nil {
t.Fatalf("Failed to take post-cleanup snapshot: %v", err)
}
// Check if memory decreased after cleanup
if postCleanup.RuntimeStats.Alloc < current.RuntimeStats.Alloc {
memoryDecrease := current.RuntimeStats.Alloc - postCleanup.RuntimeStats.Alloc
t.Logf("Memory decreased by %.2f MB after cleanup phase", float64(memoryDecrease)/(1024*1024))
} else if postCleanup.RuntimeStats.Alloc > current.RuntimeStats.Alloc {
memoryIncrease := postCleanup.RuntimeStats.Alloc - current.RuntimeStats.Alloc
if strictMode {
t.Errorf("Memory increased by %.2f MB after cleanup phase - possible cleanup issues", float64(memoryIncrease)/(1024*1024))
} else {
t.Logf("Memory increased by %.2f MB after cleanup phase - possible cleanup issues", float64(memoryIncrease)/(1024*1024))
}
}
t.Logf("Memory pool leak detection test completed")
t.Logf("Memory usage: baseline=%.2f MB, intermediate=%.2f MB, final=%.2f MB, post-cleanup=%.2f MB",
float64(baseline.RuntimeStats.Alloc)/(1024*1024),
float64(intermediate.RuntimeStats.Alloc)/(1024*1024),
float64(current.RuntimeStats.Alloc)/(1024*1024),
float64(postCleanup.RuntimeStats.Alloc)/(1024*1024))
}
+6
View File
@@ -74,6 +74,9 @@ func TestRefreshGracePeriodConfiguration(t *testing.T) {
}
func TestTokenRefreshWithinGracePeriod(t *testing.T) {
// Reset global state to prevent test interference
resetGlobalState()
refreshCount := int32(0)
tokenVersion := int32(1)
@@ -286,6 +289,9 @@ func TestGracePeriodWithProviderSpecificBehavior(t *testing.T) {
}
func TestRefreshGracePeriodConcurrency(t *testing.T) {
// Reset global state to prevent test interference
resetGlobalState()
var refreshMutex sync.Mutex
refreshCount := 0
blockedRequests := int32(0)
+9 -9
View File
@@ -12,7 +12,7 @@ import (
// TestReverseProxyHTTPSDetection tests that HTTPS is properly detected in reverse proxy environments
func TestReverseProxyHTTPSDetection(t *testing.T) {
t.Run("HTTPS_Detection_With_X_Forwarded_Proto", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Simulate request from reverse proxy (Traefik/nginx)
@@ -48,7 +48,7 @@ func TestReverseProxyHTTPSDetection(t *testing.T) {
})
t.Run("HTTPS_Detection_Without_Headers", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Request without reverse proxy headers (direct HTTP)
@@ -79,7 +79,7 @@ func TestReverseProxyHTTPSDetection(t *testing.T) {
t.Run("HTTPS_Detection_With_ForceHTTPS", func(t *testing.T) {
// Test with forceHTTPS enabled
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", true, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", true, "", NewLogger("debug"))
require.NoError(t, err)
// Even without headers, forceHTTPS should make cookies secure
@@ -110,7 +110,7 @@ func TestReverseProxyHTTPSDetection(t *testing.T) {
// TestCSRFPersistenceInReverseProxy tests CSRF token persistence in reverse proxy setups
func TestCSRFPersistenceInReverseProxy(t *testing.T) {
t.Run("CSRF_Persists_Through_OAuth_Flow_With_Proxy", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Step 1: Initial request to protected resource (HTTPS via proxy)
@@ -156,7 +156,7 @@ func TestCSRFPersistenceInReverseProxy(t *testing.T) {
})
t.Run("Session_Cookie_Domain_With_Proxy_Headers", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Test with X-Forwarded-Host header
@@ -196,7 +196,7 @@ func TestAzureOIDCWithReverseProxy(t *testing.T) {
// This test verifies Azure-specific provider detection and configuration
// without making actual network calls
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Step 1: Test session setup for Azure OAuth flow
@@ -254,7 +254,7 @@ func TestAzureOIDCWithReverseProxy(t *testing.T) {
})
t.Run("Mixed_HTTP_HTTPS_Requests", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Simulate a scenario where some requests come via HTTPS proxy and some don't
@@ -299,7 +299,7 @@ func TestAzureOIDCWithReverseProxy(t *testing.T) {
// TestEnhanceSessionSecurity verifies the security enhancement function
func TestEnhanceSessionSecurity(t *testing.T) {
t.Run("Security_Enhancement_For_AJAX_Requests", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// AJAX request via HTTPS proxy
@@ -327,7 +327,7 @@ func TestEnhanceSessionSecurity(t *testing.T) {
})
t.Run("Security_Enhancement_Missing_User_Agent", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Request without User-Agent (potential bot/attack)
+37 -34
View File
@@ -17,9 +17,12 @@ import (
// TestConcurrentTokenVerification tests race conditions in token verification
func TestConcurrentTokenVerification(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create cleanup helper
tc := newTestCleanup(t)
// Create multiple valid tokens to avoid replay detection
tokens := make([]string, 10)
for i := 0; i < 10; i++ {
@@ -40,27 +43,25 @@ func TestConcurrentTokenVerification(t *testing.T) {
}
// Create a fresh instance for this test
var goroutineWG sync.WaitGroup
tokenBlacklist := tc.addCache(NewCache())
tokenCache := tc.addTokenCache(NewTokenCache())
tOidc := &TraefikOidc{
issuerURL: "https://test-issuer.com",
clientID: "test-client-id",
jwkCache: ts.mockJWKCache,
tokenBlacklist: NewCache(),
tokenCache: NewTokenCache(),
tokenBlacklist: tokenBlacklist,
tokenCache: tokenCache,
limiter: rate.NewLimiter(rate.Every(time.Microsecond), 10000), // Very high rate limit
logger: NewLogger("debug"),
allowedUserDomains: map[string]struct{}{"example.com": {}},
httpClient: &http.Client{},
extractClaimsFunc: extractClaims,
goroutineWG: &goroutineWG,
}
tOidc.tokenVerifier = tOidc
tOidc.jwtVerifier = tOidc
// Ensure cleanup when test finishes
defer func() {
if err := tOidc.Close(); err != nil {
t.Logf("Error closing TraefikOidc instance: %v", err)
}
}()
tc.addOIDC(tOidc)
// Test concurrent verification
const numGoroutines = 50
@@ -130,7 +131,7 @@ func TestConcurrentTokenVerification(t *testing.T) {
// TestCacheMemoryExhaustion tests cache behavior under memory pressure
func TestCacheMemoryExhaustion(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create a cache with limited size
@@ -197,7 +198,7 @@ func TestCacheMemoryExhaustion(t *testing.T) {
// TestSessionConcurrencyProtection tests session safety under concurrent access
func TestSessionConcurrencyProtection(t *testing.T) {
logger := NewLogger("debug")
sessionManager, err := NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, logger)
sessionManager, err := NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -372,7 +373,7 @@ func TestProviderFailureRecovery(t *testing.T) {
// TestOversizedTokenHandling tests boundary value handling
func TestOversizedTokenHandling(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create an oversized token with large claims
@@ -441,7 +442,7 @@ func TestOversizedTokenHandling(t *testing.T) {
// TestMaliciousInputValidation tests security input validation
func TestMaliciousInputValidation(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
maliciousInputs := []struct {
@@ -496,28 +497,29 @@ func TestMaliciousInputValidation(t *testing.T) {
for _, test := range maliciousInputs {
t.Run(test.name, func(t *testing.T) {
// Create cleanup helper for subtest
tc := newTestCleanup(t)
// Create a fresh instance for each test to avoid rate limiting issues
var goroutineWG sync.WaitGroup
tokenBlacklist := tc.addCache(NewCache())
tokenCache := tc.addTokenCache(NewTokenCache())
freshOidc := &TraefikOidc{
issuerURL: "https://test-issuer.com",
clientID: "test-client-id",
jwkCache: ts.mockJWKCache,
tokenBlacklist: NewCache(),
tokenCache: NewTokenCache(),
tokenBlacklist: tokenBlacklist,
tokenCache: tokenCache,
limiter: rate.NewLimiter(rate.Every(time.Microsecond), 10000), // Very high rate limit
logger: NewLogger("debug"),
allowedUserDomains: map[string]struct{}{"example.com": {}},
httpClient: &http.Client{},
extractClaimsFunc: extractClaims,
goroutineWG: &goroutineWG,
}
freshOidc.tokenVerifier = freshOidc
freshOidc.jwtVerifier = freshOidc
// Ensure cleanup when test finishes
defer func() {
if err := freshOidc.Close(); err != nil {
t.Logf("Error closing TraefikOidc instance: %v", err)
}
}()
tc.addOIDC(freshOidc)
// All malicious inputs should be safely rejected
err := freshOidc.VerifyToken(test.token)
@@ -646,7 +648,7 @@ func TestResourceLimits(t *testing.T) {
// TestErrorRecoveryPatterns tests various error recovery scenarios
func TestErrorRecoveryPatterns(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Test recovery from cache corruption
@@ -707,9 +709,12 @@ func TestPerformanceUnderLoad(t *testing.T) {
t.Skip("Skipping performance test in short mode")
}
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create cleanup helper
tc := newTestCleanup(t)
// Create multiple valid tokens
const numTokens = 100
tokens := make([]string, numTokens)
@@ -731,27 +736,25 @@ func TestPerformanceUnderLoad(t *testing.T) {
}
// Create fresh instance with high rate limit
var goroutineWG sync.WaitGroup
tokenBlacklist := tc.addCache(NewCache())
tokenCache := tc.addTokenCache(NewTokenCache())
tOidc := &TraefikOidc{
issuerURL: "https://test-issuer.com",
clientID: "test-client-id",
jwkCache: ts.mockJWKCache,
tokenBlacklist: NewCache(),
tokenCache: NewTokenCache(),
tokenBlacklist: tokenBlacklist,
tokenCache: tokenCache,
limiter: rate.NewLimiter(rate.Every(time.Microsecond), 10000), // Very high limit
logger: NewLogger("info"), // Reduce logging for performance
allowedUserDomains: map[string]struct{}{"example.com": {}},
httpClient: &http.Client{},
extractClaimsFunc: extractClaims,
goroutineWG: &goroutineWG,
}
tOidc.tokenVerifier = tOidc
tOidc.jwtVerifier = tOidc
// Ensure cleanup when test finishes
defer func() {
if err := tOidc.Close(); err != nil {
t.Logf("Error closing TraefikOidc instance: %v", err)
}
}()
tc.addOIDC(tOidc)
// Performance test
const iterations = 1000
+1 -1
View File
@@ -192,7 +192,7 @@ func TestScopesConfiguration(t *testing.T) {
}
func TestBuildAuthURLScopeHandling(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup() // Basic setup for TraefikOidc instance
// Default scopes expected if not overridden and no user scopes provided
+53 -27
View File
@@ -21,7 +21,7 @@ import (
// TestJWTAlgorithmConfusionAttack tests if the plugin is vulnerable to JWT algorithm confusion attacks
// where an attacker might try to switch from an asymmetric algorithm (RS256) to a symmetric one (HS256)
func TestJWTAlgorithmConfusionAttack(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create a standard JWT with RS256 algorithm
@@ -86,7 +86,7 @@ func TestJWTAlgorithmConfusionAttack(t *testing.T) {
// TestJWTNoneAlgorithmAttack tests the plugin's resistance to the "none" algorithm attack
// where an attacker removes the signature and sets the algorithm to "none"
func TestJWTNoneAlgorithmAttack(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create a standard JWT
@@ -150,7 +150,7 @@ func TestJWTNoneAlgorithmAttack(t *testing.T) {
// TestJWTTokenTampering tests the plugin's ability to detect modifications to the JWT payload
func TestJWTTokenTampering(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create a standard JWT
@@ -215,7 +215,7 @@ func TestJWTTokenTampering(t *testing.T) {
// TestJWTExpiredToken tests the plugin's handling of expired tokens
func TestJWTExpiredToken(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create a JWT that is already expired
@@ -248,7 +248,7 @@ func TestJWTExpiredToken(t *testing.T) {
// TestJWTFutureToken tests the plugin's handling of tokens issued in the future
func TestJWTFutureToken(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create a JWT with a future issuance time
@@ -281,10 +281,13 @@ func TestJWTFutureToken(t *testing.T) {
// TestJWTReplayAttack tests the plugin's protection against token replay attacks
func TestJWTReplayAttack(t *testing.T) {
// Create cleanup helper
tc := newTestCleanup(t)
// Create a new instance for this test to avoid interference from global state
logger := NewLogger("debug")
tokenBlacklist := NewCache()
tokenCache := NewTokenCache()
tokenBlacklist := tc.addCache(NewCache())
tokenCache := tc.addTokenCache(NewTokenCache())
// Create keys
rsaPrivateKey, err := rsa.GenerateKey(rand.Reader, 2048)
@@ -383,7 +386,7 @@ func TestJWTReplayAttack(t *testing.T) {
// TestMissingClaims tests validation of tokens with missing required claims
func TestMissingClaims(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Test cases for missing claims
@@ -460,8 +463,11 @@ func TestMissingClaims(t *testing.T) {
// TestSessionFixationAttack tests the plugin's resistance to session fixation attacks
func TestSessionFixationAttack(t *testing.T) {
// Create cleanup helper
tc := newTestCleanup(t)
logger := NewLogger("debug")
sm, err := NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, logger)
sm, err := NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -536,6 +542,8 @@ func TestSessionFixationAttack(t *testing.T) {
}
// Create the TraefikOidc middleware
tokenBlacklist := tc.addCache(NewCache())
tokenCache := tc.addTokenCache(NewTokenCache())
tOidc := &TraefikOidc{
next: nextHandler,
name: "test",
@@ -546,8 +554,8 @@ func TestSessionFixationAttack(t *testing.T) {
clientSecret: "test-client-secret",
jwkCache: mockJWKCache,
jwksURL: "https://test-jwks-url.com",
tokenBlacklist: NewCache(),
tokenCache: NewTokenCache(),
tokenBlacklist: tokenBlacklist,
tokenCache: tokenCache,
limiter: rate.NewLimiter(rate.Every(time.Second), 10),
logger: logger,
allowedUserDomains: map[string]struct{}{"example.com": {}},
@@ -615,7 +623,7 @@ func TestSessionFixationAttack(t *testing.T) {
// TestCSRFProtection tests CSRF protection in POST requests
func TestCSRFProtection(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, logger)
sm, err := NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -798,10 +806,13 @@ func TestCSRFProtection(t *testing.T) {
// TestTokenBlacklisting tests the token blacklisting mechanism
func TestTokenBlacklisting(t *testing.T) {
// Create cleanup helper
tc := newTestCleanup(t)
// Create a new instance for this test to avoid interference from global state
logger := NewLogger("debug")
tokenBlacklist := NewCache()
tokenCache := NewTokenCache()
tokenBlacklist := tc.addCache(NewCache())
tokenCache := tc.addTokenCache(NewTokenCache())
// Create keys
rsaPrivateKey, err := rsa.GenerateKey(rand.Reader, 2048)
@@ -886,7 +897,7 @@ func TestTokenBlacklisting(t *testing.T) {
// TestDifferentSigningAlgorithms tests that the plugin properly handles different signing algorithms
func TestDifferentSigningAlgorithms(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Test cases for different algorithms - the implementation actually supports multiple algorithms
@@ -1187,7 +1198,7 @@ func createECJWK(privateKey *ecdsa.PrivateKey, alg, kid string) JWK {
// TestMalformedTokens tests the plugin's handling of malformed tokens
func TestMalformedTokens(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
testCases := []struct {
@@ -1250,23 +1261,28 @@ func TestMalformedTokens(t *testing.T) {
// TestRateLimiting tests the rate limiting functionality to prevent brute force attacks
func TestRateLimiting(t *testing.T) {
// Create cleanup helper
tc := newTestCleanup(t)
// Create a fresh instance for this test to avoid affecting other tests with rate limiting
logger := NewLogger("debug")
// Create a new test suite for this test only
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create a separate TraefikOidc instance with a very restrictive rate limiter
// This prevents the global instance from being rate-limited
tokenBlacklist := tc.addCache(NewCache())
tokenCache := tc.addTokenCache(NewTokenCache())
tOidc := &TraefikOidc{
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(),
tokenBlacklist: tokenBlacklist,
tokenCache: tokenCache,
// Allow only 2 requests per 10 seconds
limiter: rate.NewLimiter(rate.Every(10*time.Second), 2),
logger: logger,
@@ -1347,7 +1363,10 @@ func TestRateLimiting(t *testing.T) {
// TestAuthorizationHeaderBypass tests that the plugin correctly handles attempts to bypass
// authorization by directly providing an Authorization header
func TestAuthorizationHeaderBypass(t *testing.T) {
ts := &TestSuite{t: t}
// Create cleanup helper
tc := newTestCleanup(t)
ts := NewTestSuite(t)
ts.Setup()
// Create a test next handler that would indicate successful authentication
@@ -1357,6 +1376,8 @@ func TestAuthorizationHeaderBypass(t *testing.T) {
})
// Create the TraefikOidc instance
tokenBlacklist := tc.addCache(NewCache())
tokenCache := tc.addTokenCache(NewTokenCache())
tOidc := &TraefikOidc{
next: nextHandler,
name: "test",
@@ -1367,8 +1388,8 @@ func TestAuthorizationHeaderBypass(t *testing.T) {
clientSecret: "test-client-secret",
jwkCache: ts.mockJWKCache,
jwksURL: "https://test-jwks-url.com",
tokenBlacklist: NewCache(),
tokenCache: NewTokenCache(),
tokenBlacklist: tokenBlacklist,
tokenCache: tokenCache,
limiter: rate.NewLimiter(rate.Every(time.Second), 10),
logger: NewLogger("debug"),
allowedUserDomains: map[string]struct{}{"example.com": {}},
@@ -1419,7 +1440,7 @@ func TestAuthorizationHeaderBypass(t *testing.T) {
// TestEmptyAudience tests tokens with empty audience claim
func TestEmptyAudience(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create a JWT with empty audience
@@ -1452,7 +1473,7 @@ func TestEmptyAudience(t *testing.T) {
// TestEmptyIssuer tests tokens with empty issuer claim
func TestEmptyIssuer(t *testing.T) {
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
// Create a JWT with empty issuer
@@ -1485,7 +1506,10 @@ func TestEmptyIssuer(t *testing.T) {
// TestInvalidRedirectURI tests the plugin's handling of invalid redirect URIs
func TestInvalidRedirectURI(t *testing.T) {
ts := &TestSuite{t: t}
// Create cleanup helper
tc := newTestCleanup(t)
ts := NewTestSuite(t)
ts.Setup()
// Create a test request with an invalid redirect URI
@@ -1527,6 +1551,8 @@ func TestInvalidRedirectURI(t *testing.T) {
})
// Create the TraefikOidc instance
tokenBlacklist := tc.addCache(NewCache())
tokenCache := tc.addTokenCache(NewTokenCache())
tOidc := &TraefikOidc{
next: nextHandler,
name: "test",
@@ -1537,8 +1563,8 @@ func TestInvalidRedirectURI(t *testing.T) {
clientSecret: "test-client-secret",
jwkCache: ts.mockJWKCache,
jwksURL: "https://test-jwks-url.com",
tokenBlacklist: NewCache(),
tokenCache: NewTokenCache(),
tokenBlacklist: tokenBlacklist,
tokenCache: tokenCache,
limiter: rate.NewLimiter(rate.Every(time.Second), 10),
logger: NewLogger("debug"),
allowedUserDomains: map[string]struct{}{"example.com": {}},
+6 -8
View File
@@ -81,6 +81,7 @@ type SecurityMonitor struct {
eventHandlers []SecurityEventHandler
config SecurityMonitorConfig
ipMutex sync.RWMutex
cleanupTask *BackgroundTask // Store cleanup task per instance
}
// IPFailureTracker maintains failure statistics for a specific IP address.
@@ -141,9 +142,6 @@ func DefaultSecurityMonitorConfig() SecurityMonitorConfig {
}
}
// cleanupTask holds the BackgroundTask for security cleanup
var cleanupTask *BackgroundTask
// NewSecurityMonitor creates a new security monitor instance
func NewSecurityMonitor(config SecurityMonitorConfig, logger *Logger) *SecurityMonitor {
sm := &SecurityMonitor{
@@ -489,19 +487,19 @@ func (spd *SuspiciousPatternDetector) DetectSuspiciousPatterns() []string {
// startCleanupRoutine starts the background cleanup routine
func (sm *SecurityMonitor) startCleanupRoutine() {
// Use BackgroundTask abstraction for consistent management
cleanupTask = NewBackgroundTask(
sm.cleanupTask = NewBackgroundTask(
"security-monitor-cleanup",
time.Duration(sm.config.CleanupIntervalMinutes)*time.Minute,
sm.cleanup,
sm.logger)
cleanupTask.Start()
sm.cleanupTask.Start()
}
// StopCleanupRoutine stops the background cleanup routine
func (sm *SecurityMonitor) StopCleanupRoutine() {
if cleanupTask != nil {
cleanupTask.Stop()
cleanupTask = nil
if sm.cleanupTask != nil {
sm.cleanupTask.Stop()
sm.cleanupTask = nil
}
}
+120 -7
View File
@@ -243,17 +243,21 @@ type SessionManager struct {
store sessions.Store
logger *Logger
forceHTTPS bool
cookieDomain string
chunkManager *ChunkManager
cleanupMutex sync.RWMutex // Protects cleanup operations
cleanupDone bool // Track if we've attempted cookie cleanup
}
// NewSessionManager creates a new session manager with the specified configuration.
// Parameters:
// - encryptionKey: Key used to encrypt session data (must be at least 32 bytes)
// - forceHTTPS: When true, forces secure cookie attributes regardless of request scheme
// - cookieDomain: Explicit domain for cookies (e.g., ".example.com" for subdomain access)
// - logger: Logger instance for recording session-related events
//
// Returns an error if the encryption key does not meet minimum length requirements.
func NewSessionManager(encryptionKey string, forceHTTPS bool, logger *Logger) (*SessionManager, error) {
func NewSessionManager(encryptionKey string, forceHTTPS bool, cookieDomain string, logger *Logger) (*SessionManager, error) {
// Validate encryption key length.
if len(encryptionKey) < minEncryptionKeyLength {
return nil, fmt.Errorf("encryption key must be at least %d bytes long", minEncryptionKeyLength)
@@ -262,6 +266,7 @@ func NewSessionManager(encryptionKey string, forceHTTPS bool, logger *Logger) (*
sm := &SessionManager{
store: sessions.NewCookieStore([]byte(encryptionKey)),
forceHTTPS: forceHTTPS,
cookieDomain: cookieDomain,
logger: logger,
chunkManager: NewChunkManager(logger),
}
@@ -512,8 +517,13 @@ func (sm *SessionManager) EnhanceSessionSecurity(options *sessions.Options, r *h
// Always enforce security best practices
options.HttpOnly = true
// Set secure Domain attribute for production
if options.Domain == "" && r != nil {
// Set cookie domain from configuration if specified
if sm.cookieDomain != "" {
// Use the explicitly configured domain
options.Domain = sm.cookieDomain
sm.logger.Debugf("Using configured cookie domain: %s", sm.cookieDomain)
} else if options.Domain == "" && r != nil {
// Fall back to auto-detection only if no domain is configured
// In reverse proxy setups, use the original host if available
// This ensures cookies work correctly with the public domain
host := r.Host
@@ -529,10 +539,10 @@ func (sm *SessionManager) EnhanceSessionSecurity(options *sessions.Options, r *h
if colonIndex := strings.Index(host, ":"); colonIndex != -1 {
host = host[:colonIndex] // Remove port
}
// For cookie domain, we might want to use the base domain
// to allow cookies to work across subdomains
// But for now, use the exact host for security
// Use the exact host (without leading dot) for host-only cookies
// This prevents subdomain access but ensures consistency
options.Domain = host
sm.logger.Debugf("Auto-detected cookie domain: %s", host)
}
}
@@ -555,10 +565,109 @@ func (sm *SessionManager) getSessionOptions(isSecure bool) *sessions.Options {
SameSite: http.SameSiteLaxMode,
MaxAge: int(absoluteSessionTimeout.Seconds()),
Path: "/",
Domain: sm.cookieDomain, // Use configured domain if set
}
return baseOptions
}
// CleanupOldCookies attempts to remove cookies with potentially mismatched domains.
// Since browsers don't send domain info with cookies, this function tries to delete
// cookies with various domain variations to ensure cleanup after configuration changes.
func (sm *SessionManager) CleanupOldCookies(w http.ResponseWriter, r *http.Request) {
// Get all cookies from the request
cookies := r.Cookies()
// Get the current configured domain
currentDomain := sm.cookieDomain
host := r.Host
if forwardedHost := r.Header.Get("X-Forwarded-Host"); forwardedHost != "" {
host = forwardedHost
}
// Remove port if present
if colonIndex := strings.Index(host, ":"); colonIndex != -1 {
host = host[:colonIndex]
}
// Build list of domains to attempt deletion for
// This ensures we clean up cookies from various possible domains
var domainsToClean []string
// Add variations of the current host
if host != "" && !strings.Contains(host, "localhost") && !strings.Contains(host, "127.0.0.1") {
domainsToClean = append(domainsToClean,
host, // exact host
"."+host, // with leading dot
)
// Add parent domain variations if host has subdomains
parts := strings.Split(host, ".")
if len(parts) > 2 {
// For "sub.example.com", also try "example.com" and ".example.com"
parentDomain := strings.Join(parts[len(parts)-2:], ".")
domainsToClean = append(domainsToClean,
parentDomain,
"."+parentDomain,
)
}
}
// Track which cookies we're processing
processedCookies := make(map[string]bool)
// Check each cookie
for _, cookie := range cookies {
// Only process our OIDC session cookies
if strings.HasPrefix(cookie.Name, mainCookieName) ||
strings.HasPrefix(cookie.Name, accessTokenCookie) ||
strings.HasPrefix(cookie.Name, refreshTokenCookie) ||
strings.HasPrefix(cookie.Name, "_oidc_raczylo_id") ||
strings.HasPrefix(cookie.Name, "access_token_chunk_") ||
strings.HasPrefix(cookie.Name, "refresh_token_chunk_") {
// Mark this cookie as seen
processedCookies[cookie.Name] = true
// If we have a configured domain and this is the first request after config change,
// attempt to delete the cookie with various domain variations to ensure cleanup
sm.cleanupMutex.RLock()
shouldCleanup := currentDomain != "" && !sm.cleanupDone
sm.cleanupMutex.RUnlock()
if shouldCleanup {
for _, domain := range domainsToClean {
// Skip the current configured domain
if domain == currentDomain || domain == "."+currentDomain || "."+domain == currentDomain {
continue
}
// Create a deletion cookie for this domain variation
deleteCookie := &http.Cookie{
Name: cookie.Name,
Value: "",
Path: "/",
Domain: domain,
MaxAge: -1,
HttpOnly: true,
Secure: r.Header.Get("X-Forwarded-Proto") == "https" || r.TLS != nil || sm.forceHTTPS,
SameSite: http.SameSiteLaxMode,
}
http.SetCookie(w, deleteCookie)
sm.logger.Debugf("Attempting to clean up cookie %s with domain %s", cookie.Name, domain)
}
}
}
}
// Mark cleanup as done for this session manager instance
if len(processedCookies) > 0 {
sm.cleanupMutex.Lock()
if !sm.cleanupDone {
sm.cleanupDone = true
}
sm.cleanupMutex.Unlock()
}
}
// GetSession retrieves all session data for the current request.
// It loads the main session and token sessions, including any chunked token data,
// and combines them into a single SessionData structure for easy access.
@@ -571,6 +680,8 @@ func (sm *SessionManager) GetSession(r *http.Request) (*SessionData, error) {
sessionData.dirty = false
var sessionReturned bool
// CRITICAL FIX: Defer logic to return session to pool only on panic, not on successful return
// This prevents memory leaks by ensuring sessions are only pooled when not in use by caller
defer func() {
if !sessionReturned && sessionData != nil {
if r := recover(); r != nil {
@@ -636,7 +747,9 @@ func (sm *SessionManager) GetSession(r *http.Request) (*SessionData, error) {
sm.getTokenChunkSessions(r, refreshTokenCookie, sessionData.refreshTokenChunks)
sm.getTokenChunkSessions(r, idTokenCookie, sessionData.idTokenChunks)
sessionReturned = false
// CRITICAL FIX: Mark session as successfully returned to caller - prevents premature pool return
// This fixes memory leak where sessions were incorrectly returned to pool immediately after success
sessionReturned = true
return sessionData, nil
}
+5 -5
View File
@@ -14,7 +14,7 @@ import (
func TestSessionCompatibilityAfterChanges(t *testing.T) {
t.Run("Plain_HTTP_Without_Proxy_Headers", func(t *testing.T) {
// Test that plain HTTP requests without proxy headers still work
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
req := httptest.NewRequest("GET", "http://localhost:8080/test", nil)
@@ -67,7 +67,7 @@ func TestSessionCompatibilityAfterChanges(t *testing.T) {
t.Run("HTTPS_With_TLS_Field", func(t *testing.T) {
// Test direct HTTPS connection (not proxied)
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
req := httptest.NewRequest("GET", "https://example.com/test", nil)
@@ -101,7 +101,7 @@ func TestSessionCompatibilityAfterChanges(t *testing.T) {
t.Run("ForceHTTPS_Setting", func(t *testing.T) {
// Test forceHTTPS setting works regardless of request
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", true, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", true, "", NewLogger("debug"))
require.NoError(t, err)
// Plain HTTP request
@@ -132,7 +132,7 @@ func TestSessionCompatibilityAfterChanges(t *testing.T) {
})
t.Run("AJAX_Request_Gets_Strict_SameSite", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// AJAX request
@@ -165,7 +165,7 @@ func TestSessionCompatibilityAfterChanges(t *testing.T) {
})
t.Run("Missing_UserAgent_Gets_Reduced_Timeout", func(t *testing.T) {
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, NewLogger("debug"))
sessionManager, err := NewSessionManager("test-encryption-key-32-characters", false, "", NewLogger("debug"))
require.NoError(t, err)
// Request without User-Agent (suspicious)
+11 -11
View File
@@ -16,7 +16,7 @@ import (
func TestSessionPoolMemoryLeak(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -44,7 +44,7 @@ func TestSessionPoolMemoryLeak(t *testing.T) {
// Test 3: Error path in GetSession
// Modify the session store to force an error - use a different encryption key
badSM, _ := NewSessionManager("different0123456789abcdef0123456789abcdef0123456789", false, logger)
badSM, _ := NewSessionManager("different0123456789abcdef0123456789abcdef0123456789", false, "", logger)
// Get session using mismatched manager/request to force error
_, err = badSM.GetSession(req)
@@ -68,7 +68,7 @@ func TestSessionPoolMemoryLeak(t *testing.T) {
func TestSessionErrorHandling(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -95,7 +95,7 @@ func TestSessionErrorHandling(t *testing.T) {
func TestSessionClearAlwaysReturnsToPool(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -181,7 +181,7 @@ func getPooledObjects(sm *SessionManager) int {
// 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)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -344,7 +344,7 @@ func TestTokenCompressionCorruptionDetection(t *testing.T) {
// TestTokenChunkingIntegrity tests that large tokens are properly chunked and reassembled
func TestTokenChunkingIntegrity(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -419,7 +419,7 @@ func TestTokenChunkingIntegrity(t *testing.T) {
// TestTokenChunkingCorruptionResistance tests handling of corrupted chunks
func TestTokenChunkingCorruptionResistance(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -525,7 +525,7 @@ func TestTokenChunkingCorruptionResistance(t *testing.T) {
// TestTokenSizeLimits tests that token size limits are enforced
func TestTokenSizeLimits(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -592,7 +592,7 @@ func TestTokenSizeLimits(t *testing.T) {
// TestConcurrentTokenOperations tests thread safety of token operations
func TestConcurrentTokenOperations(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -646,7 +646,7 @@ func TestConcurrentTokenOperations(t *testing.T) {
// TestSessionValidationAndCleanup tests session validation and orphan cleanup
func TestSessionValidationAndCleanup(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -694,7 +694,7 @@ func TestSessionValidationAndCleanup(t *testing.T) {
// TestLargeIDTokenChunking tests that large ID tokens are properly chunked across multiple cookies
func TestLargeIDTokenChunking(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
+1
View File
@@ -43,6 +43,7 @@ type Config struct {
AllowedUsers []string `json:"allowedUsers"`
Scopes []string `json:"scopes"`
Headers []TemplatedHeader `json:"headers"`
CookieDomain string `json:"cookieDomain"`
RateLimit int `json:"rateLimit"`
RefreshGracePeriodSeconds int `json:"refreshGracePeriodSeconds"`
ForceHTTPS bool `json:"forceHTTPS"`
+7
View File
@@ -252,6 +252,13 @@ func TestTemplateIntegrationWithPlugin(t *testing.T) {
traefikOidc, ok := handler.(*TraefikOidc)
require.True(t, ok)
// Clean up the middleware when test finishes
defer func() {
if err := traefikOidc.Close(); err != nil {
t.Errorf("Failed to close middleware: %v", err)
}
}()
// Create a mock session with claims
req := httptest.NewRequest("GET", "/protected", nil)
+2 -2
View File
@@ -15,7 +15,7 @@ import (
// in the actual middleware flow
func TestTemplatedHeadersIntegration(t *testing.T) {
// Create a TestSuite to use its helper methods and fields
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
tests := []struct {
@@ -423,7 +423,7 @@ func TestTemplatedHeadersIntegration(t *testing.T) {
// TestEdgeCaseTemplatedHeaders tests edge cases for templated headers
func TestEdgeCaseTemplatedHeaders(t *testing.T) {
// Create a TestSuite to use its helper methods and fields
ts := &TestSuite{t: t}
ts := NewTestSuite(t)
ts.Setup()
tests := []struct {
@@ -102,6 +102,13 @@ func TestTraefikConfigurationParsing(t *testing.T) {
t.Fatalf("Handler is not of type *TraefikOidc")
}
// Clean up the middleware when test finishes
defer func() {
if err := middleware.Close(); err != nil {
t.Errorf("Failed to close middleware: %v", err)
}
}()
// Check that templates were parsed correctly
if len(tc.config.Headers) > 0 {
if len(middleware.headerTemplates) != len(tc.config.Headers) {
@@ -232,6 +239,13 @@ func TestIssue55ReproductionAttempt(t *testing.T) {
t.Fatalf("Handler is not of type *TraefikOidc")
}
// Clean up the middleware when test finishes
defer func() {
if err := middleware.Close(); err != nil {
t.Errorf("Failed to close middleware: %v", err)
}
}()
// Check that the header template was parsed
if len(middleware.headerTemplates) != 1 {
t.Errorf("Expected 1 header template, got %d", len(middleware.headerTemplates))
@@ -6,6 +6,7 @@ import (
"log"
"net/http"
"net/http/httptest"
"sync"
"testing"
"time"
@@ -25,6 +26,125 @@ func (w *testWriter) Write(p []byte) (n int, err error) {
// Test helper adapters for the new test files
// resetGlobalState resets all global singletons to prevent test interference
func resetGlobalState() {
// Reset global cache manager
cacheManagerMutex.Lock()
if globalCacheManager != nil {
globalCacheManager.Close()
globalCacheManager = nil
}
cacheManagerOnce = sync.Once{}
cacheManagerMutex.Unlock()
// Reset and cleanup replay cache
cleanupReplayCache()
// Reset memory pools
memoryPoolMutex.Lock()
globalMemoryPools = nil
memoryPoolOnce = sync.Once{}
memoryPoolMutex.Unlock()
}
// testCleanup provides comprehensive cleanup for tests to prevent goroutine leaks
type testCleanup struct {
t *testing.T
caches []*Cache
servers []*httptest.Server
oidcs []*TraefikOidc
mu sync.Mutex
}
// newTestCleanup creates a new test cleanup helper that automatically registers cleanup
func newTestCleanup(t *testing.T) *testCleanup {
tc := &testCleanup{
t: t,
caches: make([]*Cache, 0),
servers: make([]*httptest.Server, 0),
oidcs: make([]*TraefikOidc, 0),
}
// Register cleanup to run even if test panics
t.Cleanup(func() {
tc.cleanupAll()
})
return tc
}
// addCache registers a cache for cleanup
func (tc *testCleanup) addCache(c *Cache) *Cache {
tc.mu.Lock()
defer tc.mu.Unlock()
tc.caches = append(tc.caches, c)
return c
}
// addTokenCache registers a token cache for cleanup
func (tc *testCleanup) addTokenCache(c *TokenCache) *TokenCache {
tc.mu.Lock()
defer tc.mu.Unlock()
if c != nil && c.cache != nil {
tc.caches = append(tc.caches, c.cache)
}
return c
}
// addServer registers an httptest server for cleanup
func (tc *testCleanup) addServer(s *httptest.Server) *httptest.Server {
tc.mu.Lock()
defer tc.mu.Unlock()
tc.servers = append(tc.servers, s)
return s
}
// addOIDC registers a TraefikOidc instance for cleanup
func (tc *testCleanup) addOIDC(o *TraefikOidc) *TraefikOidc {
tc.mu.Lock()
defer tc.mu.Unlock()
tc.oidcs = append(tc.oidcs, o)
return o
}
// cleanupAll cleans up all registered resources
func (tc *testCleanup) cleanupAll() {
tc.mu.Lock()
defer tc.mu.Unlock()
// Close all caches
for _, c := range tc.caches {
if c != nil {
c.Close()
}
}
// Close all servers
for _, s := range tc.servers {
if s != nil {
s.Close()
}
}
// Close all OIDC instances
for _, o := range tc.oidcs {
if o != nil {
// Close caches within the OIDC instance
if o.tokenCache != nil && o.tokenCache.cache != nil {
o.tokenCache.cache.Close()
}
if o.tokenBlacklist != nil {
o.tokenBlacklist.Close()
}
// Call Close if it exists
o.Close()
}
}
// Reset global state
resetGlobalState()
}
// createTestConfig creates a config with all required fields populated for testing
func createTestConfig() *Config {
config := CreateConfig()
@@ -38,6 +158,9 @@ func createTestConfig() *Config {
// setupTestOIDCMiddleware creates a test OIDC middleware instance with mock servers
func setupTestOIDCMiddleware(t *testing.T, config *Config) (*TraefikOidc, *httptest.Server) {
// Reset global state to ensure test isolation
resetGlobalState()
// Create mock OIDC server
var serverURL string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -85,7 +208,7 @@ func setupTestOIDCMiddleware(t *testing.T, config *Config) (*TraefikOidc, *httpt
logInfo: log.New(&testWriter{t}, "INFO: ", 0),
logDebug: log.New(&testWriter{t}, "DEBUG: ", 0),
}
sessionManager, _ := NewSessionManager(config.SessionEncryptionKey, false, logger)
sessionManager, _ := NewSessionManager(config.SessionEncryptionKey, false, "", logger)
// Create next handler
nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -114,6 +237,9 @@ func setupTestOIDCMiddleware(t *testing.T, config *Config) (*TraefikOidc, *httpt
testTokenURL := testIssuerURL + "/token"
testJWKSURL := testIssuerURL + "/keys"
// Create WaitGroup for background goroutines
var wg sync.WaitGroup
// Create TraefikOidc instance directly
oidc := &TraefikOidc{
next: nextHandler,
@@ -143,8 +269,10 @@ func setupTestOIDCMiddleware(t *testing.T, config *Config) (*TraefikOidc, *httpt
forceHTTPS: config.ForceHTTPS,
allowedUserDomains: make(map[string]struct{}),
jwkCache: &JWKCache{},
metadataCache: NewMetadataCache(),
metadataCache: NewMetadataCache(nil),
ctx: context.Background(),
goroutineWG: &wg,
providerURL: serverURL,
}
// Process excluded URLs
@@ -217,7 +345,7 @@ func createMockJWT(t *testing.T, sub, email string) string {
func createTestSession() *SessionData {
// Create a minimal session manager for testing
logger := newNoOpLogger()
sessionManager, _ := NewSessionManager("test-encryption-key-32-characters", false, logger)
sessionManager, _ := NewSessionManager("test-encryption-key-32-characters", false, "", logger)
// Create a test request
req := httptest.NewRequest("GET", "/", nil)
+1 -1
View File
@@ -14,4 +14,4 @@ func generateRandomString(length int) string {
return "random-string-fallback"
}
return hex.EncodeToString(bytes)
}
}
+4 -4
View File
@@ -17,7 +17,7 @@ import (
// Token verified successfully multiple times, then fails with "signature verification failed"
func TestTokenCorruptionScenario(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -180,7 +180,7 @@ func TestCompressionIntegrityFailure(t *testing.T) {
// TestChunkReassemblyEdgeCases tests edge cases in chunk reassembly that could cause corruption
func TestChunkReassemblyEdgeCases(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -296,7 +296,7 @@ func TestChunkReassemblyEdgeCases(t *testing.T) {
// TestRaceConditionProtection tests that concurrent access doesn't cause corruption
func TestRaceConditionProtection(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -432,7 +432,7 @@ func TestMemoryExhaustionProtection(t *testing.T) {
// TestBackwardCompatibility ensures that sessions created before the fixes still work
func TestBackwardCompatibility(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
+8 -8
View File
@@ -88,7 +88,7 @@ func TestTokenTypeDistinction(t *testing.T) {
// 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 := NewTestSuite(t)
ts.Setup()
// Create different tokens for ID and access tokens
@@ -157,8 +157,8 @@ func TestTokenTypeIntegration(t *testing.T) {
clientSecret: "test-client-secret",
jwkCache: ts.mockJWKCache,
jwksURL: "https://test-jwks-url.com",
tokenBlacklist: NewCache(),
tokenCache: NewTokenCache(),
tokenBlacklist: ts.tOidc.tokenBlacklist,
tokenCache: ts.tOidc.tokenCache,
limiter: rate.NewLimiter(rate.Every(time.Second), 10),
logger: NewLogger("debug"),
allowedUserDomains: map[string]struct{}{"example.com": {}},
@@ -247,7 +247,7 @@ func TestSessionIDTokenAccessToken(t *testing.T) {
logger := NewLogger("debug")
// Create a session manager
sessionManager, err := NewSessionManager("test-session-encryption-key-at-least-32-bytes", false, logger)
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)
}
@@ -318,7 +318,7 @@ func TestSessionIDTokenAccessToken(t *testing.T) {
// TestTokenCorruptionIntegrationFlows tests the complete token handling flow with corruption scenarios
func TestTokenCorruptionIntegrationFlows(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -442,7 +442,7 @@ func TestTokenCorruptionIntegrationFlows(t *testing.T) {
// TestSessionPersistenceWithCorruption tests that session corruption is handled across requests
func TestSessionPersistenceWithCorruption(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -506,7 +506,7 @@ func TestSessionPersistenceWithCorruption(t *testing.T) {
// TestConcurrentTokenOperationsWithCorruption tests concurrent access with intentional corruption
func TestConcurrentTokenOperationsWithCorruption(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
@@ -592,7 +592,7 @@ func TestConcurrentTokenOperationsWithCorruption(t *testing.T) {
// TestTokenValidationEdgeCases tests edge cases in token validation
func TestTokenValidationEdgeCases(t *testing.T) {
logger := NewLogger("debug")
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger)
if err != nil {
t.Fatalf("Failed to create session manager: %v", err)
}
+437
View File
@@ -0,0 +1,437 @@
package traefikoidc
import (
"context"
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/rsa"
"encoding/base64"
"fmt"
"math/big"
"net/http"
"net/http/httptest"
"sync"
"testing"
"time"
"golang.org/x/time/rate"
)
// UnifiedTestFramework provides comprehensive test management with automatic cleanup
type UnifiedTestFramework struct {
t *testing.T
mu sync.Mutex
cleanupFuncs []func()
caches []*Cache
tokenCaches []*TokenCache
servers []*httptest.Server
oidcInstances []*TraefikOidc
httpClients []*http.Client
backgroundTasks []*BackgroundTask
goroutineWGs []*sync.WaitGroup
contexts []context.Context
cancelFuncs []context.CancelFunc
}
// NewUnifiedTestFramework creates a new test framework with automatic cleanup
func NewUnifiedTestFramework(t *testing.T) *UnifiedTestFramework {
utf := &UnifiedTestFramework{
t: t,
cleanupFuncs: make([]func(), 0),
caches: make([]*Cache, 0),
tokenCaches: make([]*TokenCache, 0),
servers: make([]*httptest.Server, 0),
oidcInstances: make([]*TraefikOidc, 0),
httpClients: make([]*http.Client, 0),
backgroundTasks: make([]*BackgroundTask, 0),
goroutineWGs: make([]*sync.WaitGroup, 0),
contexts: make([]context.Context, 0),
cancelFuncs: make([]context.CancelFunc, 0),
}
// Register cleanup to run even if test panics
t.Cleanup(func() {
utf.CleanupAll()
})
// Reset global state at start
resetGlobalState()
return utf
}
// AddCleanupFunc adds a custom cleanup function
func (utf *UnifiedTestFramework) AddCleanupFunc(fn func()) {
utf.mu.Lock()
defer utf.mu.Unlock()
utf.cleanupFuncs = append(utf.cleanupFuncs, fn)
}
// NewCache creates and tracks a cache for cleanup
func (utf *UnifiedTestFramework) NewCache() *Cache {
utf.mu.Lock()
defer utf.mu.Unlock()
// Use NewLazyCache to avoid immediate background task start
cache := NewLazyCache()
utf.caches = append(utf.caches, cache)
return cache
}
// NewCacheWithLogger creates and tracks a cache with logger
func (utf *UnifiedTestFramework) NewCacheWithLogger(logger *Logger) *Cache {
utf.mu.Lock()
defer utf.mu.Unlock()
// Use lazy cache to avoid immediate background task
cache := NewLazyCacheWithLogger(logger)
utf.caches = append(utf.caches, cache)
return cache
}
// NewTokenCache creates and tracks a token cache for cleanup
func (utf *UnifiedTestFramework) NewTokenCache() *TokenCache {
utf.mu.Lock()
defer utf.mu.Unlock()
// Create token cache with lazy initialization
tokenCache := &TokenCache{
cache: NewLazyCache(),
}
utf.tokenCaches = append(utf.tokenCaches, tokenCache)
return tokenCache
}
// NewHTTPTestServer creates and tracks an httptest server
func (utf *UnifiedTestFramework) NewHTTPTestServer(handler http.Handler) *httptest.Server {
utf.mu.Lock()
defer utf.mu.Unlock()
server := httptest.NewServer(handler)
utf.servers = append(utf.servers, server)
return server
}
// NewHTTPClient creates and tracks an HTTP client
func (utf *UnifiedTestFramework) NewHTTPClient() *http.Client {
utf.mu.Lock()
defer utf.mu.Unlock()
client := &http.Client{
Timeout: 10 * time.Second,
Transport: &http.Transport{
MaxIdleConns: 2,
IdleConnTimeout: 5 * time.Second,
},
}
utf.httpClients = append(utf.httpClients, client)
return client
}
// NewContext creates and tracks a context with cancel
func (utf *UnifiedTestFramework) NewContext() context.Context {
utf.mu.Lock()
defer utf.mu.Unlock()
ctx, cancel := context.WithCancel(context.Background())
utf.contexts = append(utf.contexts, ctx)
utf.cancelFuncs = append(utf.cancelFuncs, cancel)
return ctx
}
// NewWaitGroup creates and tracks a wait group
func (utf *UnifiedTestFramework) NewWaitGroup() *sync.WaitGroup {
utf.mu.Lock()
defer utf.mu.Unlock()
wg := &sync.WaitGroup{}
utf.goroutineWGs = append(utf.goroutineWGs, wg)
return wg
}
// NewBackgroundTask creates and tracks a background task
func (utf *UnifiedTestFramework) NewBackgroundTask(name string, interval time.Duration, taskFunc func(), logger *Logger) *BackgroundTask {
utf.mu.Lock()
defer utf.mu.Unlock()
task := NewBackgroundTask(name, interval, taskFunc, logger)
utf.backgroundTasks = append(utf.backgroundTasks, task)
return task
}
// AddOIDCInstance tracks a TraefikOidc instance for cleanup
func (utf *UnifiedTestFramework) AddOIDCInstance(oidc *TraefikOidc) *TraefikOidc {
utf.mu.Lock()
defer utf.mu.Unlock()
utf.oidcInstances = append(utf.oidcInstances, oidc)
return oidc
}
// CleanupAll performs comprehensive cleanup of all tracked resources
func (utf *UnifiedTestFramework) CleanupAll() {
utf.mu.Lock()
defer utf.mu.Unlock()
utf.t.Logf("Starting unified test framework cleanup...")
// Cancel all contexts first to signal shutdown
for _, cancel := range utf.cancelFuncs {
if cancel != nil {
cancel()
}
}
// Stop all background tasks
for _, task := range utf.backgroundTasks {
if task != nil {
task.Stop()
}
}
// Close all OIDC instances
for _, oidc := range utf.oidcInstances {
if oidc != nil {
if err := oidc.Close(); err != nil {
utf.t.Logf("Error closing OIDC instance: %v", err)
}
}
}
// Close all token caches
for _, tc := range utf.tokenCaches {
if tc != nil && tc.cache != nil {
tc.cache.Close()
}
}
// Close all caches
for _, c := range utf.caches {
if c != nil {
c.Close()
}
}
// Close all HTTP servers
for _, s := range utf.servers {
if s != nil {
s.Close()
}
}
// Close idle connections in HTTP clients
for _, client := range utf.httpClients {
if client != nil && client.Transport != nil {
if transport, ok := client.Transport.(*http.Transport); ok {
transport.CloseIdleConnections()
}
}
}
// Wait for all goroutines with timeout
for _, wg := range utf.goroutineWGs {
if wg != nil {
done := make(chan struct{})
go func(wg *sync.WaitGroup) {
wg.Wait()
close(done)
}(wg)
select {
case <-done:
// Goroutines finished
case <-time.After(5 * time.Second):
utf.t.Logf("Timeout waiting for goroutines to finish")
}
}
}
// Run custom cleanup functions
for _, fn := range utf.cleanupFuncs {
if fn != nil {
fn()
}
}
// Reset global state
resetGlobalState()
utf.t.Logf("Unified test framework cleanup completed")
}
// EnhancedTestSuite extends TestSuite with UnifiedTestFramework
type EnhancedTestSuite struct {
*TestSuite
*UnifiedTestFramework
}
// NewEnhancedTestSuite creates a test suite with automatic resource management
func NewEnhancedTestSuite(t *testing.T) *EnhancedTestSuite {
utf := NewUnifiedTestFramework(t)
ts := &TestSuite{
t: t,
utf: utf,
}
ets := &EnhancedTestSuite{
TestSuite: ts,
UnifiedTestFramework: utf,
}
return ets
}
// Setup initializes the enhanced test suite with automatic cleanup
func (ets *EnhancedTestSuite) Setup() {
var err error
// Generate RSA key
ets.rsaPrivateKey, err = rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
ets.TestSuite.t.Fatalf("Failed to generate RSA key: %v", err)
}
ets.rsaPublicKey = &ets.rsaPrivateKey.PublicKey
// Generate EC key
ets.ecPrivateKey, err = ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
if err != nil {
ets.TestSuite.t.Fatalf("Failed to generate EC key: %v", err)
}
// Create JWK for RSA public key
jwk := JWK{
Kty: "RSA",
Kid: "test-key-id",
Alg: "RS256",
N: base64.RawURLEncoding.EncodeToString(ets.rsaPublicKey.N.Bytes()),
E: base64.RawURLEncoding.EncodeToString(bigIntToBytes(big.NewInt(int64(ets.rsaPublicKey.E)))),
}
jwks := &JWKSet{
Keys: []JWK{jwk},
}
// Create mock JWK cache
ets.mockJWKCache = &MockJWKCache{
JWKS: jwks,
Err: nil,
}
// Create test JWT token
now := time.Now()
ets.token, err = createTestJWT(ets.rsaPrivateKey, "RS256", "test-key-id", map[string]interface{}{
"iss": "https://test-issuer.com",
"aud": "test-client-id",
"exp": now.Add(1 * time.Hour).Unix(),
"iat": now.Add(-2 * time.Minute).Unix(),
"nbf": now.Add(-2 * time.Minute).Unix(),
"sub": "test-subject",
"email": "[email protected]",
"nonce": "test-nonce",
"jti": generateRandomString(16),
})
if err != nil {
ets.TestSuite.t.Fatalf("Failed to create test JWT: %v", err)
}
// Create session manager
logger := NewLogger("info")
ets.sessionManager, _ = NewSessionManager("test-secret-key-that-is-at-least-32-bytes", false, "", logger)
// Create WaitGroup for the OIDC instance
goroutineWG := ets.NewWaitGroup()
// Create TraefikOidc instance with tracked resources
ets.tOidc = &TraefikOidc{
issuerURL: "https://test-issuer.com",
clientID: "test-client-id",
clientSecret: "test-client-secret",
jwkCache: ets.mockJWKCache,
jwksURL: "https://test-jwks-url.com",
revocationURL: "https://revocation-endpoint.com",
limiter: rate.NewLimiter(rate.Every(time.Second), 10),
tokenBlacklist: ets.NewCache(), // Use framework's tracked cache
tokenCache: ets.NewTokenCache(), // Use framework's tracked token cache
logger: logger,
allowedUserDomains: map[string]struct{}{"example.com": {}},
excludedURLs: map[string]struct{}{"/favicon": {}},
httpClient: ets.NewHTTPClient(), // Use framework's tracked HTTP client
redirURLPath: "/callback",
logoutURLPath: "/callback/logout",
tokenURL: "https://test-issuer.com/token",
extractClaimsFunc: extractClaims,
initComplete: make(chan struct{}),
sessionManager: ets.sessionManager,
goroutineWG: goroutineWG,
ctx: ets.NewContext(), // Use framework's tracked context
tokenCleanupStopChan: make(chan struct{}),
metadataRefreshStopChan: make(chan struct{}),
}
close(ets.tOidc.initComplete)
ets.tOidc.tokenVerifier = ets.tOidc
ets.tOidc.jwtVerifier = ets.tOidc
// Set default mock exchanger
ets.tOidc.tokenExchanger = &MockTokenExchanger{
ExchangeCodeFunc: func(ctx context.Context, grantType, codeOrToken, redirectURL, codeVerifier string) (*TokenResponse, error) {
return &TokenResponse{
IDToken: ets.token,
AccessToken: ets.token,
RefreshToken: "default-refresh-token",
ExpiresIn: 3600,
}, nil
},
RefreshTokenFunc: func(refreshToken string) (*TokenResponse, error) {
return nil, fmt.Errorf("default mock: refresh not expected")
},
RevokeTokenFunc: func(token, tokenType string) error {
return nil
},
}
// Track the OIDC instance for cleanup
ets.AddOIDCInstance(ets.tOidc)
}
// CreateMockOIDCServer creates a mock OIDC server with automatic cleanup
func (utf *UnifiedTestFramework) CreateMockOIDCServer(issuerURL string) *httptest.Server {
server := utf.NewHTTPTestServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/.well-known/openid-configuration":
w.Header().Set("Content-Type", "application/json")
fmt.Fprintf(w, `{
"issuer": "%s",
"authorization_endpoint": "%s/auth",
"token_endpoint": "%s/token",
"userinfo_endpoint": "%s/userinfo",
"jwks_uri": "%s/keys",
"revocation_endpoint": "%s/revoke"
}`, issuerURL, issuerURL, issuerURL, issuerURL, issuerURL, issuerURL)
case "/keys":
w.Header().Set("Content-Type", "application/json")
w.Write([]byte(`{
"keys": [{
"kty": "RSA",
"kid": "test-key-id",
"use": "sig",
"n": "test-n-value",
"e": "AQAB"
}]
}`))
case "/token":
w.Header().Set("Content-Type", "application/json")
w.Write([]byte(`{
"access_token": "test-access-token",
"id_token": "` + ValidIDToken + `",
"refresh_token": "test-refresh-token",
"token_type": "bearer",
"expires_in": 3600
}`))
default:
w.WriteHeader(http.StatusNotFound)
}
}))
return server
}