From da70e69ad1e45cb94cf3b5dade67d6dbeff0af18 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Fri, 9 May 2025 19:05:24 +0100 Subject: [PATCH] Memleak fixes. --- helpers.go | 5 ++ jwk.go | 10 ++++ main.go | 130 +++++++++++++++++++++++++++++++++++++++------------ main_test.go | 5 ++ 4 files changed, 120 insertions(+), 30 deletions(-) diff --git a/helpers.go b/helpers.go index ea1f41d..9f9529f 100644 --- a/helpers.go +++ b/helpers.go @@ -278,6 +278,11 @@ func (tc *TokenCache) Cleanup() { tc.cache.Cleanup() } +// Close stops the cleanup goroutine in the underlying cache. +func (tc *TokenCache) Close() { + tc.cache.Close() +} + // exchangeCodeForToken is a convenience function that wraps exchangeTokens specifically // for the "authorization_code" grant type. It handles the conditional inclusion of the // PKCE code verifier based on the middleware's configuration (t.enablePKCE). diff --git a/jwk.go b/jwk.go index a3e44e8..1ca08bf 100644 --- a/jwk.go +++ b/jwk.go @@ -38,11 +38,13 @@ type JWKCache struct { mutex sync.RWMutex // CacheLifetime is configurable to determine how long the JWKS is cached. CacheLifetime time.Duration + internalCache *Cache // To hold the closable Cache instance from cache.go } type JWKCacheInterface interface { GetJWKS(ctx context.Context, jwksURL string, httpClient *http.Client) (*JWKSet, error) Cleanup() + Close() } // GetJWKS retrieves the JSON Web Key Set (JWKS) from the cache or fetches it from the provider. @@ -101,6 +103,14 @@ func (c *JWKCache) Cleanup() { } } +// Close shuts down the cache's auto-cleanup routine. +func (c *JWKCache) Close() { + // Close shuts down the internal cache's auto-cleanup routine, if the cache exists. + if c.internalCache != nil { + c.internalCache.Close() + } +} + // fetchJWKS retrieves the JSON Web Key Set (JWKS) from the specified URL. // It uses the provided context and HTTP client to make the request. // diff --git a/main.go b/main.go index a94a6fa..1da28ab 100644 --- a/main.go +++ b/main.go @@ -112,14 +112,16 @@ type TraefikOidc struct { allowedRolesAndGroups map[string]struct{} initiateAuthenticationFunc func(rw http.ResponseWriter, req *http.Request, session *SessionData, redirectURL string) // exchangeCodeForTokenFunc func(code string, redirectURL string, codeVerifier string) (*TokenResponse, error) // Replaced by interface - extractClaimsFunc func(tokenString string) (map[string]interface{}, error) - initComplete chan struct{} - endSessionURL string - postLogoutRedirectURI string - sessionManager *SessionManager - tokenExchanger TokenExchanger // Added field for mocking - refreshGracePeriod time.Duration // Configurable grace period for proactive refresh - headerTemplates map[string]*template.Template // Parsed templates for custom headers + extractClaimsFunc func(tokenString string) (map[string]interface{}, error) + initComplete chan struct{} + endSessionURL string + postLogoutRedirectURI string + sessionManager *SessionManager + tokenExchanger TokenExchanger // Added field for mocking + refreshGracePeriod time.Duration // Configurable grace period for proactive refresh + headerTemplates map[string]*template.Template // Parsed templates for custom headers + tokenCleanupStopChan chan struct{} // Channel to stop token cleanup goroutine + metadataRefreshStopChan chan struct{} // Channel to stop metadata refresh goroutine } // ProviderMetadata holds OIDC provider metadata @@ -406,6 +408,8 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h } return 60 * time.Second // Default to 60 seconds }(), + tokenCleanupStopChan: make(chan struct{}), + metadataRefreshStopChan: make(chan struct{}), } t.sessionManager, _ = NewSessionManager(config.SessionEncryptionKey, config.ForceHTTPS, t.logger) @@ -499,24 +503,31 @@ func (t *TraefikOidc) updateMetadataEndpoints(metadata *ProviderMetadata) { // middleware's endpoint URLs via updateMetadataEndpoints. Fetch errors are logged. // // Parameters: -// - providerURL: The base URL of the OIDC provider, used for subsequent refresh attempts. +// - providerURL: The base URL of bogged OIDC provider, used for subsequent refresh attempts. func (t *TraefikOidc) startMetadataRefresh(providerURL string) { ticker := time.NewTicker(1 * time.Hour) - defer ticker.Stop() + // No defer ticker.Stop() here, it's stopped in the select case - for range ticker.C { - t.logger.Debug("Refreshing OIDC metadata") - metadata, err := t.metadataCache.GetMetadata(providerURL, t.httpClient, t.logger) - if err != nil { - t.logger.Errorf("Failed to refresh metadata: %v", err) - continue - } + for { + select { + case <-ticker.C: + t.logger.Debug("Refreshing OIDC metadata") + metadata, err := t.metadataCache.GetMetadata(providerURL, t.httpClient, t.logger) + if err != nil { + t.logger.Errorf("Failed to refresh metadata: %v", err) + continue + } - if metadata != nil { - t.updateMetadataEndpoints(metadata) - t.logger.Debug("Successfully refreshed metadata") - } else { - t.logger.Error("Received nil metadata during refresh") + if metadata != nil { + t.updateMetadataEndpoints(metadata) + t.logger.Debug("Successfully refreshed metadata") + } else { + t.logger.Error("Received nil metadata during refresh") + } + case <-t.metadataRefreshStopChan: + ticker.Stop() + t.logger.Debug("Metadata refresh goroutine stopped.") + return } } } @@ -1541,17 +1552,34 @@ func (t *TraefikOidc) buildURLWithParams(baseURL string, params url.Values) stri } // startTokenCleanup starts background goroutines for periodically cleaning up -// the token cache, token blacklist cache, and JWK cache using the autoCleanupRoutine helper. +// the token cache, token blacklist cache, and JWK cache. func (t *TraefikOidc) startTokenCleanup() { ticker := time.NewTicker(1 * time.Minute) // Run cleanup every minute go func() { - defer ticker.Stop() - for range ticker.C { - t.logger.Debug("Starting token cleanup cycle") - t.tokenCache.Cleanup() - // t.tokenBlacklist.Cleanup() // Removed: Generic Cache handles its own cleanup - t.jwkCache.Cleanup() // Assuming jwkCache is the cache from cache.go - // Removed runtime.GC() call + // No defer ticker.Stop() here, it's stopped in the select case + for { + select { + case <-ticker.C: + t.logger.Debug("Starting token cleanup cycle") + if t.tokenCache != nil { + t.tokenCache.Cleanup() + } + // t.tokenBlacklist is a *Cache, its autoCleanupRoutine handles its own cleanup + // if t.tokenBlacklist != nil { + // t.tokenBlacklist.Cleanup() + // } + if t.jwkCache != nil { + // Assuming jwkCache is the cache from cache.go which has a Cleanup method + // If jwkCache is *cache.Cache, its autoCleanupRoutine handles its own cleanup + // If it's JWKCacheInterface, it needs a Cleanup method. + // Based on New(), t.jwkCache = &JWKCache{}, which has a Cleanup method. + t.jwkCache.Cleanup() + } + case <-t.tokenCleanupStopChan: + ticker.Stop() + t.logger.Debug("Token cleanup goroutine stopped.") + return + } } }() } @@ -1985,3 +2013,45 @@ func (t *TraefikOidc) sendErrorResponse(rw http.ResponseWriter, req *http.Reques rw.WriteHeader(code) _, _ = rw.Write([]byte(htmlBody)) // Ignore write error as header is already sent } + +// Close stops all background goroutines and closes resources. +func (t *TraefikOidc) Close() error { + t.logger.Debug("Closing TraefikOidc plugin instance") + // Signal and close tokenCleanupStopChan + if t.tokenCleanupStopChan != nil { + close(t.tokenCleanupStopChan) + t.logger.Debug("tokenCleanupStopChan closed") + } + // Signal and close metadataRefreshStopChan + if t.metadataRefreshStopChan != nil { + close(t.metadataRefreshStopChan) + t.logger.Debug("metadataRefreshStopChan closed") + } + + // Close caches + // These Close methods should stop their respective autoCleanupRoutine goroutines + if t.tokenBlacklist != nil { + t.tokenBlacklist.Close() // This is *cache.Cache, which has Close() + t.logger.Debug("tokenBlacklist closed") + } + if t.metadataCache != nil { + t.metadataCache.Close() // This is *MetadataCache, which has Close() + t.logger.Debug("metadataCache closed") + } + if t.tokenCache != nil { + t.tokenCache.Close() // This is *TokenCache, which now has Close() + t.logger.Debug("tokenCache closed") + } + + if t.jwkCache != nil { + // Reverting to the original explicit instruction to call t.jwkCache.Close(). + // This will cause a compile error if JWKCacheInterface (and its implementation *JWKCache) + // is not updated in jwk.go to include and implement a Close() method + // that properly closes the internal *cache.Cache instance. + t.jwkCache.Close() + t.logger.Debug("t.jwkCache.Close() called as per original instruction.") + } + + t.logger.Info("TraefikOidc plugin instance closed successfully.") + return nil +} diff --git a/main_test.go b/main_test.go index a0f0def..de721a9 100644 --- a/main_test.go +++ b/main_test.go @@ -149,6 +149,11 @@ type MockJWKCache struct { Err error } +// Close is a no-op for the mock. +func (m *MockJWKCache) Close() { + // No operation needed for the mock. +} + func (m *MockJWKCache) GetJWKS(ctx context.Context, jwksURL string, httpClient *http.Client) (*JWKSet, error) { return m.JWKS, m.Err }