diff --git a/autocleanup.go b/autocleanup.go index 7926309..e4980da 100644 --- a/autocleanup.go +++ b/autocleanup.go @@ -9,12 +9,16 @@ import ( // 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 - wg *sync.WaitGroup + 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. @@ -29,37 +33,54 @@ type BackgroundTask struct { // Returns: // - A configured BackgroundTask ready to be started. func NewBackgroundTask(name string, interval time.Duration, taskFunc func(), logger *Logger, wg ...*sync.WaitGroup) *BackgroundTask { - var waitGroup *sync.WaitGroup + var externalWG *sync.WaitGroup if len(wg) > 0 { - waitGroup = wg[0] + externalWG = wg[0] } return &BackgroundTask{ - name: name, - interval: interval, - stopChan: make(chan struct{}), - taskFunc: taskFunc, - logger: logger, - wg: waitGroup, + 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() { - if bt.wg != nil { - bt.wg.Add(1) + 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. -// If a WaitGroup was provided, it waits for the goroutine to complete. +// 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) - if bt.wg != nil { - bt.wg.Wait() - } + 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. @@ -67,8 +88,11 @@ func (bt *BackgroundTask) Stop() { // until the stop signal is received. func (bt *BackgroundTask) run() { defer func() { - if bt.wg != nil { - bt.wg.Done() + // Always decrement internal WaitGroup + bt.internalWG.Done() + // Decrement external WaitGroup if provided + if bt.externalWG != nil { + bt.externalWG.Done() } }() ticker := time.NewTicker(bt.interval) diff --git a/azure_oidc_test.go b/azure_oidc_test.go index a295da7..925c70e 100644 --- a/azure_oidc_test.go +++ b/azure_oidc_test.go @@ -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{}), diff --git a/cache_test.go b/cache_test.go index 0ad4150..dffa221 100644 --- a/cache_test.go +++ b/cache_test.go @@ -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") } } diff --git a/enhanced_memory_leak_test.go b/enhanced_memory_leak_test.go index 6bc36e2..5bc221c 100644 --- a/enhanced_memory_leak_test.go +++ b/enhanced_memory_leak_test.go @@ -1,7 +1,6 @@ package traefikoidc import ( - "context" "encoding/json" "fmt" "net/http" @@ -32,39 +31,24 @@ func TestBackgroundGoroutineLeaks(t *testing.T) { t.Logf("Baseline: Memory=%d KB, Goroutines=%d", baselineAlloc/1024, baselineGoroutines) - // Create middleware instance - config := CreateConfig() - config.ProviderURL = "https://example.com" - config.ClientID = "test-client" - config.ClientSecret = "test-secret" - config.SessionEncryptionKey = "test-encryption-key-that-is-long-enough-32bytes" + // Create test cleanup + tc := newTestCleanup(t) + + // Create middleware instance with mock setup + config := createTestConfig() config.LogLevel = "error" // Reduce noise - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusOK) - }) - - handler, err := New(ctx, next, config, "test-middleware") - if err != nil { - t.Fatal(err) - } - - // Cast to TraefikOidc to access internals - middleware, ok := handler.(*TraefikOidc) - if !ok { - t.Fatal("Failed to cast to TraefikOidc") - } + 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 30 seconds...") + t.Log("Letting middleware sit idle for 3 seconds...") - // Take measurements every 5 seconds - for i := 0; i < 6; i++ { - time.Sleep(5 * time.Second) + // Take measurements every 1 second + for i := 0; i < 3; i++ { + time.Sleep(1 * time.Second) runtime.GC() runtime.ReadMemStats(&m) @@ -75,12 +59,12 @@ func TestBackgroundGoroutineLeaks(t *testing.T) { goroutineIncrease := currentGoroutines - baselineGoroutines t.Logf("After %d seconds: Memory increase=%.2f MB, Goroutine increase=%d", - (i+1)*5, allocIncrease, goroutineIncrease) + (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)*5) + allocIncrease, (i + 1)) } // Check for goroutine leaks (more than 10 extra goroutines) @@ -91,8 +75,7 @@ func TestBackgroundGoroutineLeaks(t *testing.T) { } // Clean up - err = middleware.Close() - if err != nil { + if err := middleware.Close(); err != nil { t.Errorf("Failed to close middleware: %v", err) } @@ -171,8 +154,8 @@ func TestHTTPClientConnectionLeaks(t *testing.T) { } } - // Let connections sit idle (now only 5s with our optimizations) - time.Sleep(6 * time.Second) // Slightly longer than new IdleConnTimeout (5s) + // Let connections sit idle briefly to test cleanup + time.Sleep(1 * time.Second) // Reduced for faster tests // Force cleanup for _, client := range clients { @@ -188,10 +171,25 @@ func TestHTTPClientConnectionLeaks(t *testing.T) { // 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 @@ -220,13 +218,25 @@ func TestCacheBackgroundTaskLeaks(t *testing.T) { cache.Close() } - // Wait for goroutines to stop - time.Sleep(500 * time.Millisecond) + // 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 remainingGoroutines > 5 { // Allow small tolerance + if !success && remainingGoroutines > 5 { // Allow small tolerance t.Errorf("Cache cleanup goroutines not stopped properly: %d extra goroutines remain", remainingGoroutines) } @@ -295,7 +305,15 @@ func TestGlobalSingletonMemoryGrowth(t *testing.T) { runtime.GC() runtime.ReadMemStats(&m) finalAlloc := m.Alloc - finalAllocIncrease := float64(finalAlloc-baselineAlloc) / 1024 / 1024 + + // 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) @@ -592,8 +610,8 @@ func TestConcurrentMemoryLeaks(t *testing.T) { } // Let it run and measure periodically - for i := 0; i < 5; i++ { - time.Sleep(5 * time.Second) + for i := 0; i < 3; i++ { + time.Sleep(1 * time.Second) runtime.GC() runtime.ReadMemStats(&m) @@ -604,12 +622,12 @@ func TestConcurrentMemoryLeaks(t *testing.T) { goroutineIncrease := currentGoroutines - baselineGoroutines t.Logf("After %d seconds of concurrent load: Memory increase=%.2f MB, Goroutine increase=%d", - (i+1)*5, allocIncrease, goroutineIncrease) + (i + 1), allocIncrease, goroutineIncrease) // Under sustained concurrent load, memory should stabilize - if i > 2 && allocIncrease > 20.0 { + if i > 1 && allocIncrease > 20.0 { t.Errorf("Memory leak under concurrent load: %.2f MB after %d seconds", - allocIncrease, (i+1)*5) + allocIncrease, (i + 1)) } } diff --git a/google_session_test.go b/google_session_test.go index 3a64f4d..8e27029 100644 --- a/google_session_test.go +++ b/google_session_test.go @@ -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, diff --git a/jwt.go b/jwt.go index 58face5..166dc15 100644 --- a/jwt.go +++ b/jwt.go @@ -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() } diff --git a/main.go b/main.go index b1dd81b..4501323 100644 --- a/main.go +++ b/main.go @@ -1074,7 +1074,7 @@ func fetchMetadata(wellKnownURL string, httpClient *http.Client) (*ProviderMetad // - 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 !t.firstRequestReceived && !strings.HasPrefix(req.URL.Path, "/health") { + if !strings.HasPrefix(req.URL.Path, "/health") { t.firstRequestMutex.Lock() if !t.firstRequestReceived { t.firstRequestReceived = true diff --git a/main_bench_test.go b/main_bench_test.go index fa7c40a..8bbb3a8 100644 --- a/main_bench_test.go +++ b/main_bench_test.go @@ -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" diff --git a/main_test.go b/main_test.go index 99478c2..9a067a0 100644 --- a/main_test.go +++ b/main_test.go @@ -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 { @@ -92,6 +107,9 @@ func (ts *TestSuite) Setup() { logger := NewLogger("info") 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{ issuerURL: "https://test-issuer.com", @@ -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. @@ -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 @@ -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 { @@ -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 diff --git a/robustness_test.go b/robustness_test.go index d12b34f..6186ddd 100644 --- a/robustness_test.go +++ b/robustness_test.go @@ -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++ { @@ -41,12 +44,14 @@ 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": {}}, @@ -56,13 +61,7 @@ func TestConcurrentTokenVerification(t *testing.T) { } 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 @@ -132,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 @@ -374,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 @@ -443,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 { @@ -498,14 +497,19 @@ 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": {}}, @@ -515,13 +519,7 @@ func TestMaliciousInputValidation(t *testing.T) { } 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) @@ -650,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 @@ -711,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) @@ -736,12 +737,14 @@ 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": {}}, @@ -751,13 +754,7 @@ func TestPerformanceUnderLoad(t *testing.T) { } 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 diff --git a/scopes_test.go b/scopes_test.go index fd0d37c..0b0ffc9 100644 --- a/scopes_test.go +++ b/scopes_test.go @@ -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 diff --git a/security_edge_cases_test.go b/security_edge_cases_test.go index 5ee5f1b..311e475 100644 --- a/security_edge_cases_test.go +++ b/security_edge_cases_test.go @@ -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,6 +463,9 @@ 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) if err != nil { @@ -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": {}}, @@ -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": {}}, diff --git a/security_monitoring.go b/security_monitoring.go index ddfddf2..6dcb7dd 100644 --- a/security_monitoring.go +++ b/security_monitoring.go @@ -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 } } diff --git a/session.go b/session.go index 19660fe..6aaf027 100644 --- a/session.go +++ b/session.go @@ -659,9 +659,11 @@ func (sm *SessionManager) CleanupOldCookies(w http.ResponseWriter, r *http.Reque } // Mark cleanup as done for this session manager instance - if !sm.cleanupDone && len(processedCookies) > 0 { + if len(processedCookies) > 0 { sm.cleanupMutex.Lock() - sm.cleanupDone = true + if !sm.cleanupDone { + sm.cleanupDone = true + } sm.cleanupMutex.Unlock() } } diff --git a/template_processing_test.go b/template_processing_test.go index 24e83a4..be74dce 100644 --- a/template_processing_test.go +++ b/template_processing_test.go @@ -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) diff --git a/templated_header_integration_test.go b/templated_header_integration_test.go index af1754d..600c3d6 100644 --- a/templated_header_integration_test.go +++ b/templated_header_integration_test.go @@ -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 { diff --git a/templated_header_traefik_integration_test.go b/templated_header_traefik_integration_test.go index e5fe63d..b55f469 100644 --- a/templated_header_traefik_integration_test.go +++ b/templated_header_traefik_integration_test.go @@ -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)) diff --git a/test_helpers_adapter.go b/test_helpers_adapter_test.go similarity index 78% rename from test_helpers_adapter.go rename to test_helpers_adapter_test.go index 1258a93..a6f97aa 100644 --- a/test_helpers_adapter.go +++ b/test_helpers_adapter_test.go @@ -37,14 +37,8 @@ func resetGlobalState() { cacheManagerOnce = sync.Once{} cacheManagerMutex.Unlock() - // Reset replay cache - replayCacheMu.Lock() - if replayCache != nil { - replayCache.Close() - replayCache = nil - } - replayCacheOnce = sync.Once{} - replayCacheMu.Unlock() + // Reset and cleanup replay cache + cleanupReplayCache() // Reset memory pools memoryPoolMutex.Lock() @@ -53,6 +47,104 @@ func resetGlobalState() { 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() diff --git a/test_tokens.go b/test_tokens_test.go similarity index 100% rename from test_tokens.go rename to test_tokens_test.go diff --git a/token_handling_test.go b/token_handling_test.go index 4be9008..9c1d47b 100644 --- a/token_handling_test.go +++ b/token_handling_test.go @@ -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": {}}, diff --git a/unified_test_framework_test.go b/unified_test_framework_test.go new file mode 100644 index 0000000..fc25fa8 --- /dev/null +++ b/unified_test_framework_test.go @@ -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": "user@example.com", + "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 +}