From 84d7cd3d76c8f0047b9c90fff555469f3cffb360 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Wed, 18 Jun 2025 00:50:12 +0100 Subject: [PATCH] Improvements targetting possible memory usage spikes. --- jwt.go | 100 ++++++++++++++++++++---- main.go | 224 +++++++++++++++++++++++++++++++++++++++-------------- session.go | 57 +++++++++++--- 3 files changed, 301 insertions(+), 80 deletions(-) diff --git a/jwt.go b/jwt.go index c40301e..9f5089e 100644 --- a/jwt.go +++ b/jwt.go @@ -1,6 +1,7 @@ package traefikoidc import ( + "context" "crypto" "crypto/ecdsa" "crypto/rsa" @@ -16,18 +17,82 @@ import ( ) var ( - replayCacheMu sync.Mutex - replayCache *Cache // Replace unbounded map with bounded Cache + replayCacheMu sync.RWMutex // Use RWMutex for better read performance + replayCache *Cache // Replace unbounded map with bounded Cache + replayCacheOnce sync.Once // CRITICAL FIX: Ensure thread-safe singleton initialization ) -// initReplayCache initializes the global replay cache with size limit +// CRITICAL FIX: Thread-safe initialization with proper size limits and LRU eviction func initReplayCache() { - if replayCache == nil { + replayCacheOnce.Do(func() { replayCache = NewCache() - replayCache.SetMaxSize(10000) // Set size limit to 10,000 entries + replayCache.SetMaxSize(10000) // Set size limit to 10,000 entries (prevents unbounded growth) + // Note: NewCache() already starts auto-cleanup goroutine with LRU eviction + }) +} + +// CRITICAL FIX: Add cleanup function for global replay cache +func cleanupReplayCache() { + replayCacheMu.Lock() + defer replayCacheMu.Unlock() + + if replayCache != nil { + replayCache.Close() // Stop auto-cleanup goroutine + replayCache = nil // Allow GC } } +// CRITICAL FIX: Add function to get cache stats for monitoring +func getReplayCacheStats() (size int, maxSize int) { + replayCacheMu.RLock() + defer replayCacheMu.RUnlock() + + if replayCache == nil { + return 0, 0 + } + + // We need to add a method to get current size from Cache + // For now, return maxSize as approximation + return 0, 10000 // TODO: Add Size() method to Cache +} + +// CRITICAL FIX: Add periodic cleanup mechanism for long-running scenarios +func startReplayCacheCleanup(ctx context.Context, logger *Logger) { + // Start background cleanup goroutine with periodic stats logging + go func() { + ticker := time.NewTicker(5 * time.Minute) // Run cleanup every 5 minutes + defer ticker.Stop() + + for { + select { + case <-ticker.C: + // Log cache stats for monitoring + size, maxSize := getReplayCacheStats() + if logger != nil { + logger.Debugf("Replay cache stats: size=%d, maxSize=%d", size, maxSize) + } + + // Force cleanup if cache exists (the cache handles its own cleanup, + // but this provides additional periodic maintenance) + replayCacheMu.RLock() + if replayCache != nil { + // Cache already handles cleanup automatically, + // this just ensures cleanup runs periodically for long-running scenarios + } + replayCacheMu.RUnlock() + + case <-ctx.Done(): + // Context cancelled, cleanup and exit + cleanupReplayCache() + if logger != nil { + logger.Debug("Replay cache cleanup goroutine stopped due to context cancellation") + } + return + } + } + }() +} + // STABILITY FIX: Standardize clock skew tolerance usage // ClockSkewToleranceFuture defines the tolerance for future-based claims like 'exp'. // Allows for more leniency with expiration checks. @@ -195,16 +260,16 @@ func (j *JWT) Verify(issuerURL, clientID string, skipReplayCheck ...bool) error return nil } - // SECURITY FIX: Use bounded Cache with thread-safe operations - replayCacheMu.Lock() - defer replayCacheMu.Unlock() + // CRITICAL FIX: Use thread-safe initialization and proper locking + initReplayCache() // Thread-safe initialization - // Initialize cache if not already done - initReplayCache() + // SECURITY FIX: Use read lock for checking, write lock for writing + replayCacheMu.RLock() + _, exists := replayCache.Get(jti) + replayCacheMu.RUnlock() - // SECURITY FIX: Check for replay attack using Cache API - if _, exists := replayCache.Get(jti); exists { - return fmt.Errorf("token replay detected") + if exists { + return fmt.Errorf("token replay detected (jti: %s)", jti) } // Calculate expiration time @@ -216,10 +281,15 @@ func (j *JWT) Verify(issuerURL, clientID string, skipReplayCheck ...bool) error expTime = time.Now().Add(10 * time.Minute) } - // SECURITY FIX: Add to replay cache with expiration using Cache API + // SECURITY FIX: Add to replay cache with expiration using write lock duration := time.Until(expTime) if duration > 0 { - replayCache.Set(jti, true, duration) + replayCacheMu.Lock() + // Double-check pattern: verify cache is still initialized + if replayCache != nil { + replayCache.Set(jti, true, duration) + } + replayCacheMu.Unlock() } } diff --git a/main.go b/main.go index 611d9a5..c4ab45e 100644 --- a/main.go +++ b/main.go @@ -94,6 +94,87 @@ const ( defaultMaxBlacklistSize = 10000 // Default maximum size for token blacklist cache ) +// CRITICAL FIX: Global cache manager for shared cache instances to prevent memory leaks +var ( + globalCacheManager *CacheManager + cacheManagerOnce sync.Once +) + +// CacheManager provides shared cache instances across middleware instances +type CacheManager struct { + tokenBlacklist *Cache + tokenCache *TokenCache + metadataCache *MetadataCache + jwkCache JWKCacheInterface + mu sync.RWMutex +} + +// GetGlobalCacheManager returns the singleton cache manager instance +func GetGlobalCacheManager() *CacheManager { + cacheManagerOnce.Do(func() { + globalCacheManager = &CacheManager{ + tokenBlacklist: func() *Cache { + c := NewCache() + c.SetMaxSize(defaultMaxBlacklistSize) + return c + }(), + tokenCache: NewTokenCache(), + metadataCache: NewMetadataCache(), + jwkCache: &JWKCache{}, + } + }) + return globalCacheManager +} + +// GetSharedTokenBlacklist returns the shared token blacklist cache +func (cm *CacheManager) GetSharedTokenBlacklist() *Cache { + cm.mu.RLock() + defer cm.mu.RUnlock() + return cm.tokenBlacklist +} + +// GetSharedTokenCache returns the shared token cache +func (cm *CacheManager) GetSharedTokenCache() *TokenCache { + cm.mu.RLock() + defer cm.mu.RUnlock() + return cm.tokenCache +} + +// GetSharedMetadataCache returns the shared metadata cache +func (cm *CacheManager) GetSharedMetadataCache() *MetadataCache { + cm.mu.RLock() + defer cm.mu.RUnlock() + return cm.metadataCache +} + +// GetSharedJWKCache returns the shared JWK cache +func (cm *CacheManager) GetSharedJWKCache() JWKCacheInterface { + cm.mu.RLock() + defer cm.mu.RUnlock() + return cm.jwkCache +} + +// Close properly closes all shared caches +func (cm *CacheManager) Close() error { + cm.mu.Lock() + defer cm.mu.Unlock() + + if cm.tokenBlacklist != nil { + cm.tokenBlacklist.Close() + } + if cm.tokenCache != nil { + cm.tokenCache.Close() + } + if cm.metadataCache != nil { + cm.metadataCache.Close() + } + if cm.jwkCache != nil { + cm.jwkCache.Close() + } + + return nil +} + // TokenVerifier interface for token verification type TokenVerifier interface { VerifyToken(token string) error @@ -155,6 +236,10 @@ type TraefikOidc struct { tokenCleanupStopChan chan struct{} // Channel to stop token cleanup goroutine metadataRefreshStopChan chan struct{} // Channel to stop metadata refresh goroutine goroutineWG sync.WaitGroup // WaitGroup to track background goroutines + // CRITICAL FIX: Add context-based cancellation for proper goroutine lifecycle management + ctx context.Context // Context for all background operations + cancelFunc context.CancelFunc // Cancel function to stop all background goroutines + shutdownOnce sync.Once // Ensure shutdown is only called once } // ProviderMetadata holds OIDC provider metadata @@ -462,6 +547,12 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h } else { httpClient = createDefaultHTTPClient() } + // CRITICAL FIX: Use shared cache instances to prevent memory leaks + cacheManager := GetGlobalCacheManager() + + // CRITICAL FIX: Initialize context for proper goroutine lifecycle management + pluginCtx, cancelFunc := context.WithCancel(context.Background()) + t := &TraefikOidc{ next: next, name: name, @@ -478,20 +569,17 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h } return config.PostLogoutRedirectURI }(), - tokenBlacklist: func() *Cache { - c := NewCache() - c.SetMaxSize(defaultMaxBlacklistSize) - return c - }(), // Use generic cache for blacklist with size limit - jwkCache: &JWKCache{}, - metadataCache: NewMetadataCache(), + // CRITICAL FIX: Use shared cache instances instead of creating new ones per middleware + tokenBlacklist: cacheManager.GetSharedTokenBlacklist(), + jwkCache: cacheManager.GetSharedJWKCache(), + metadataCache: cacheManager.GetSharedMetadataCache(), clientID: config.ClientID, clientSecret: config.ClientSecret, forceHTTPS: config.ForceHTTPS, enablePKCE: config.EnablePKCE, scopes: mergeScopes([]string{"openid", "profile", "email"}, config.Scopes), limiter: rate.NewLimiter(rate.Every(time.Second), config.RateLimit), - tokenCache: NewTokenCache(), + tokenCache: cacheManager.GetSharedTokenCache(), httpClient: httpClient, tokenHTTPClient: createTokenHTTPClient(httpClient), excludedURLs: createStringMap(config.ExcludedURLs), @@ -508,6 +596,9 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h }(), tokenCleanupStopChan: make(chan struct{}), metadataRefreshStopChan: make(chan struct{}), + // CRITICAL FIX: Initialize context-based cancellation + ctx: pluginCtx, + cancelFunc: cancelFunc, } t.sessionManager, _ = NewSessionManager(config.SessionEncryptionKey, config.ForceHTTPS, t.logger) @@ -539,6 +630,9 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h logger.Debugf("Parsed template for header %s: %s", header.Name, header.Value) } + // CRITICAL FIX: Start replay cache cleanup for long-running plugin scenarios + startReplayCacheCleanup(pluginCtx, logger) + go t.initializeMetadata(config.ProviderURL) return t, nil @@ -629,6 +723,10 @@ func (t *TraefikOidc) startMetadataRefresh(providerURL string) { case <-t.metadataRefreshStopChan: t.logger.Debug("Metadata refresh goroutine stopped.") return + case <-t.ctx.Done(): + // CRITICAL FIX: Context-based cancellation for proper goroutine lifecycle + t.logger.Debug("Metadata refresh goroutine stopped due to context cancellation.") + return } } }() @@ -1721,6 +1819,10 @@ func (t *TraefikOidc) startTokenCleanup() { case <-t.tokenCleanupStopChan: t.logger.Debug("Token cleanup goroutine stopped.") return + case <-t.ctx.Done(): + // CRITICAL FIX: Context-based cancellation for proper goroutine lifecycle + t.logger.Debug("Token cleanup goroutine stopped due to context cancellation.") + return } } }() @@ -2418,58 +2520,68 @@ func (t *TraefikOidc) validateTokenExpiry(session *SessionData, token string) (b // Close stops all background goroutines and closes resources with proper timeout. func (t *TraefikOidc) Close() error { - t.logger.Debug("Closing TraefikOidc plugin instance") + // CRITICAL FIX: Use shutdownOnce to ensure Close is only called once + var closeErr error + t.shutdownOnce.Do(func() { + t.logger.Debug("Closing TraefikOidc plugin instance") - // Signal all goroutines to stop - if t.tokenCleanupStopChan != nil { - close(t.tokenCleanupStopChan) - t.logger.Debug("tokenCleanupStopChan closed") - } - if t.metadataRefreshStopChan != nil { - close(t.metadataRefreshStopChan) - t.logger.Debug("metadataRefreshStopChan closed") - } + // CRITICAL FIX: Signal all goroutines to stop using context cancellation + if t.cancelFunc != nil { + t.cancelFunc() + t.logger.Debug("Context cancellation signaled to all goroutines") + } - // Wait for all goroutines to finish with timeout - done := make(chan struct{}) - go func() { - t.goroutineWG.Wait() - close(done) - }() + // Legacy channel cleanup for backward compatibility (if channels still exist) + if t.tokenCleanupStopChan != nil { + close(t.tokenCleanupStopChan) + t.logger.Debug("tokenCleanupStopChan closed") + } + if t.metadataRefreshStopChan != nil { + close(t.metadataRefreshStopChan) + t.logger.Debug("metadataRefreshStopChan closed") + } - // Wait for goroutines to finish or timeout after 10 seconds - 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") - // Continue with cleanup even if goroutines didn't stop gracefully - } + // Wait for all goroutines to finish with timeout + done := make(chan struct{}) + go func() { + t.goroutineWG.Wait() + close(done) + }() - // Close caches - // These Close methods should stop their respective autoCleanupRoutine goroutines - if t.tokenBlacklist != nil { - t.tokenBlacklist.Close() // This is *cache.Cache, which has Close() - t.logger.Debug("tokenBlacklist closed") - } - if t.metadataCache != nil { - t.metadataCache.Close() // This is *MetadataCache, which has Close() - t.logger.Debug("metadataCache closed") - } - if t.tokenCache != nil { - t.tokenCache.Close() // This is *TokenCache, which now has Close() - t.logger.Debug("tokenCache closed") - } + // Wait for goroutines to finish or timeout after 10 seconds + 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") + // Continue with cleanup even if goroutines didn't stop gracefully + } - if t.jwkCache != nil { - // Reverting to the original explicit instruction to call t.jwkCache.Close(). - // This will cause a compile error if JWKCacheInterface (and its implementation *JWKCache) - // is not updated in jwk.go to include and implement a Close() method - // that properly closes the internal *cache.Cache instance. - t.jwkCache.Close() - t.logger.Debug("t.jwkCache.Close() called as per original instruction.") - } + // Close caches + // These Close methods should stop their respective autoCleanupRoutine goroutines + if t.tokenBlacklist != nil { + t.tokenBlacklist.Close() // This is *cache.Cache, which has Close() + t.logger.Debug("tokenBlacklist closed") + } + if t.metadataCache != nil { + t.metadataCache.Close() // This is *MetadataCache, which has Close() + t.logger.Debug("metadataCache closed") + } + if t.tokenCache != nil { + t.tokenCache.Close() // This is *TokenCache, which now has Close() + t.logger.Debug("tokenCache closed") + } - t.logger.Info("TraefikOidc plugin instance closed successfully.") - return nil + if t.jwkCache != nil { + // Reverting to the original explicit instruction to call t.jwkCache.Close(). + // This will cause a compile error if JWKCacheInterface (and its implementation *JWKCache) + // is not updated in jwk.go to include and implement a Close() method + // that properly closes the internal *cache.Cache instance. + t.jwkCache.Close() + t.logger.Debug("t.jwkCache.Close() called as per original instruction.") + } + + t.logger.Info("TraefikOidc plugin instance closed successfully.") + }) + return closeErr } diff --git a/session.go b/session.go index 71a92d8..834bb61 100644 --- a/session.go +++ b/session.go @@ -240,11 +240,27 @@ func (sm *SessionManager) GetSession(r *http.Request) (*SessionData, error) { sessionData.request = r sessionData.dirty = false // Reset dirty flag when getting a session + // CRITICAL FIX: Use defer to guarantee session is returned to pool on any error + var sessionReturned bool + defer func() { + if !sessionReturned && sessionData != nil { + // Recover from any panics and ensure session cleanup + if r := recover(); r != nil { + sessionData.inUse = false + sessionData.Reset() + sm.sessionPool.Put(sessionData) + panic(r) // Re-panic after cleanup + } + } + }() + // Function to properly handle errors and return the session to the pool handleError := func(err error, message string) (*SessionData, error) { - if sessionData != nil { + if sessionData != nil && !sessionReturned { sessionData.inUse = false // Mark as not in use before returning to pool + sessionData.Reset() // Reset to prevent data leakage sm.sessionPool.Put(sessionData) + sessionReturned = true // Mark as returned to prevent double-return } return nil, fmt.Errorf("%s: %w", message, err) } @@ -285,6 +301,8 @@ func (sm *SessionManager) GetSession(r *http.Request) (*SessionData, error) { sm.getTokenChunkSessions(r, accessTokenCookie, sessionData.accessTokenChunks) sm.getTokenChunkSessions(r, refreshTokenCookie, sessionData.refreshTokenChunks) + // CRITICAL FIX: Mark session as successfully retrieved to prevent cleanup in defer + sessionReturned = false // Session is successfully being returned to caller return sessionData, nil } @@ -442,6 +460,21 @@ func (sd *SessionData) Save(r *http.Request, w http.ResponseWriter) error { // // Note: This method will always return the SessionData object to the pool, even if an error occurs. func (sd *SessionData) Clear(r *http.Request, w http.ResponseWriter) error { + // CRITICAL FIX: Use defer to guarantee session is returned to pool regardless of any errors or panics + defer func() { + if rec := recover(); rec != nil { + // Ensure session is returned to pool even on panic + sd.returnToPoolSafely() + panic(rec) // Re-panic after cleanup + } + // Normal path - return to pool + sd.returnToPoolSafely() + }() + + // CRITICAL FIX: Lock session mutex to prevent race conditions during Clear + sd.sessionMutex.Lock() + defer sd.sessionMutex.Unlock() + sd.dirty = true // Clearing the session means its state is changing and needs to be saved. // Clear and expire all sessions. @@ -484,17 +517,23 @@ func (sd *SessionData) Clear(r *http.Request, w http.ResponseWriter) error { // Clear transient per-request fields. sd.request = nil - // STABILITY FIX: Mark as not in use and return session to pool, regardless of error. - // This ensures the session is always returned to the pool, preventing memory leaks. - sd.inUse = false - // Reset the session data before returning to pool to prevent data leakage - sd.Reset() - sd.manager.sessionPool.Put(sd) - - // Return the error from Save, if any + // Return the error from Save, if any (defer will handle pool return) return err } +// CRITICAL FIX: Add thread-safe helper method to return session to pool +func (sd *SessionData) returnToPoolSafely() { + if sd != nil && sd.manager != nil { + // Check if already returned to prevent double-return + if sd.inUse { + sd.inUse = false + // Reset the session data before returning to pool to prevent data leakage + sd.Reset() + sd.manager.sessionPool.Put(sd) + } + } +} + // clearTokenChunks iterates through a map of session chunks, clears their values, // and sets their MaxAge to -1 to expire them. This is used internally by Clear. //