diff --git a/helpers.go b/helpers.go index 312289d..0dc936e 100644 --- a/helpers.go +++ b/helpers.go @@ -292,12 +292,16 @@ type TokenBlacklist struct { // mutex protects concurrent access to the blacklist mutex sync.RWMutex + + // maxSize is the maximum number of tokens in the blacklist + maxSize int } // NewTokenBlacklist creates a new TokenBlacklist instance. func NewTokenBlacklist() *TokenBlacklist { return &TokenBlacklist{ blacklist: make(map[string]time.Time), + maxSize: 1000, // Limit the size to prevent unbounded growth } } @@ -305,6 +309,30 @@ func NewTokenBlacklist() *TokenBlacklist { func (tb *TokenBlacklist) Add(tokenID string, expiration time.Time) { tb.mutex.Lock() defer tb.mutex.Unlock() + + // Clean up expired tokens if we're at capacity + if len(tb.blacklist) >= tb.maxSize { + now := time.Now() + for token, exp := range tb.blacklist { + if now.After(exp) { + delete(tb.blacklist, token) + } + } + // If still at capacity after cleanup, remove oldest token + if len(tb.blacklist) >= tb.maxSize { + var oldestToken string + var oldestTime time.Time + first := true + for token, exp := range tb.blacklist { + if first || exp.Before(oldestTime) { + oldestToken = token + oldestTime = exp + first = false + } + } + delete(tb.blacklist, oldestToken) + } + } tb.blacklist[tokenID] = expiration } diff --git a/main.go b/main.go index 8dbd3ac..33e9b7c 100644 --- a/main.go +++ b/main.go @@ -664,12 +664,42 @@ func (t *TraefikOidc) buildAuthURL(redirectURL, state, nonce string) string { // startTokenCleanup starts the token cleanup goroutine func (t *TraefikOidc) startTokenCleanup() { - ticker := time.NewTicker(1 * time.Minute) + ctx, cancel := context.WithCancel(context.Background()) + ticker := time.NewTicker(30 * time.Second) // Increased frequency to prevent memory buildup + go func() { - for range ticker.C { - t.logger.Debug("Cleaning up token cache") - t.tokenCache.Cleanup() - t.tokenBlacklist.Cleanup() + defer ticker.Stop() + defer cancel() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + t.logger.Debug("Starting token cleanup cycle") + + // Run cleanup in a separate goroutine with timeout + cleanupCtx, cleanupCancel := context.WithTimeout(ctx, 10*time.Second) + done := make(chan struct{}) + + go func() { + defer close(done) + t.tokenCache.Cleanup() + t.tokenBlacklist.Cleanup() + }() + + // Wait for cleanup to complete or timeout + select { + case <-cleanupCtx.Done(): + if cleanupCtx.Err() == context.DeadlineExceeded { + t.logger.Error("Token cleanup cycle timed out") + } + case <-done: + t.logger.Debug("Token cleanup cycle completed successfully") + } + + cleanupCancel() + } } }() }