diff --git a/internal/api/oidc_service.go b/internal/api/oidc_service.go index c33b950..ccb787c 100644 --- a/internal/api/oidc_service.go +++ b/internal/api/oidc_service.go @@ -173,8 +173,9 @@ func (s *OIDCService) exchangeCode(ctx context.Context, code string, entry *oidc // oidcStateStore keeps short-lived authorization state tokens. type oidcStateStore struct { - mu sync.RWMutex - entries map[string]*oidcStateEntry + mu sync.RWMutex + entries map[string]*oidcStateEntry + stopCleanup chan struct{} } type oidcStateEntry struct { @@ -186,7 +187,45 @@ type oidcStateEntry struct { } func newOIDCStateStore() *oidcStateStore { - return &oidcStateStore{entries: make(map[string]*oidcStateEntry)} + s := &oidcStateStore{ + entries: make(map[string]*oidcStateEntry), + stopCleanup: make(chan struct{}), + } + + // Start cleanup routine to prevent memory leak from abandoned OIDC flows + go func() { + ticker := time.NewTicker(5 * time.Minute) + defer ticker.Stop() + + for { + select { + case <-ticker.C: + s.cleanup() + case <-s.stopCleanup: + return + } + } + }() + + return s +} + +// cleanup removes expired state entries +func (s *oidcStateStore) cleanup() { + s.mu.Lock() + defer s.mu.Unlock() + + now := time.Now() + for state, entry := range s.entries { + if now.After(entry.ExpiresAt) { + delete(s.entries, state) + } + } +} + +// Stop stops the cleanup routine +func (s *oidcStateStore) Stop() { + close(s.stopCleanup) } func (s *oidcStateStore) Put(state string, entry *oidcStateEntry) { diff --git a/internal/api/ratelimit.go b/internal/api/ratelimit.go index d2efc10..8791897 100644 --- a/internal/api/ratelimit.go +++ b/internal/api/ratelimit.go @@ -7,17 +7,19 @@ import ( ) type RateLimiter struct { - attempts map[string][]time.Time - mu sync.RWMutex - limit int - window time.Duration + attempts map[string][]time.Time + mu sync.RWMutex + limit int + window time.Duration + stopCleanup chan struct{} } func NewRateLimiter(limit int, window time.Duration) *RateLimiter { rl := &RateLimiter{ - attempts: make(map[string][]time.Time), - limit: limit, - window: window, + attempts: make(map[string][]time.Time), + limit: limit, + window: window, + stopCleanup: make(chan struct{}), } // Clean up old entries periodically @@ -25,14 +27,24 @@ func NewRateLimiter(limit int, window time.Duration) *RateLimiter { ticker := time.NewTicker(1 * time.Minute) defer ticker.Stop() - for range ticker.C { - rl.cleanup() + for { + select { + case <-ticker.C: + rl.cleanup() + case <-rl.stopCleanup: + return + } } }() return rl } +// Stop stops the cleanup routine +func (rl *RateLimiter) Stop() { + close(rl.stopCleanup) +} + func (rl *RateLimiter) Allow(ip string) bool { rl.mu.Lock() defer rl.mu.Unlock() diff --git a/internal/api/recovery_tokens.go b/internal/api/recovery_tokens.go index 04a2069..fa0a2dc 100644 --- a/internal/api/recovery_tokens.go +++ b/internal/api/recovery_tokens.go @@ -25,9 +25,10 @@ type RecoveryToken struct { // RecoveryTokenStore manages recovery tokens type RecoveryTokenStore struct { - tokens map[string]*RecoveryToken - mu sync.RWMutex - dataPath string + tokens map[string]*RecoveryToken + mu sync.RWMutex + dataPath string + stopCleanup chan struct{} } var ( @@ -39,8 +40,9 @@ var ( func InitRecoveryTokenStore(dataPath string) { recoveryStoreOnce.Do(func() { recoveryStore = &RecoveryTokenStore{ - tokens: make(map[string]*RecoveryToken), - dataPath: dataPath, + tokens: make(map[string]*RecoveryToken), + dataPath: dataPath, + stopCleanup: make(chan struct{}), } recoveryStore.load() @@ -182,11 +184,22 @@ func (r *RecoveryTokenStore) cleanupRoutine() { ticker := time.NewTicker(1 * time.Hour) defer ticker.Stop() - for range ticker.C { - r.cleanup() + for { + select { + case <-ticker.C: + r.cleanup() + case <-r.stopCleanup: + log.Debug().Msg("Recovery token cleanup routine stopped") + return + } } } +// Stop stops the cleanup routine +func (r *RecoveryTokenStore) Stop() { + close(r.stopCleanup) +} + // cleanup removes expired and used tokens func (r *RecoveryTokenStore) cleanup() { r.mu.Lock()