package tokenmanager import ( "context" "sync" "sync/atomic" "time" "go.uber.org/zap" "codeberg.org/mapleopentech/monorepo/cloud/maplefile-backend/pkg/maplefile/client" "codeberg.org/mapleopentech/monorepo/native/desktop/maplefile/internal/service/auth" "codeberg.org/mapleopentech/monorepo/native/desktop/maplefile/internal/usecase/session" ) // Manager handles automatic token refresh with graceful shutdown type Manager struct { config Config client *client.Client authService *auth.Service getSession *session.GetByIdUseCase logger *zap.Logger // Lifecycle management ctx context.Context cancel context.CancelFunc stopCh chan struct{} // Signal to stop stoppedCh chan struct{} // Confirmation of stopped running atomic.Bool // Thread-safe running flag // Refresh state management mu sync.Mutex consecutiveFailures int } // New creates a new token manager func New( config Config, client *client.Client, authService *auth.Service, getSession *session.GetByIdUseCase, logger *zap.Logger, ) *Manager { return &Manager{ config: config, client: client, authService: authService, getSession: getSession, logger: logger.Named("token-manager"), } } // Start begins the token refresh background process // Safe to call multiple times - will only start once func (m *Manager) Start() { // Only start if not already running if !m.running.CompareAndSwap(false, true) { m.logger.Debug("Token manager already running, skipping start") return } m.ctx, m.cancel = context.WithCancel(context.Background()) m.stopCh = make(chan struct{}) m.stoppedCh = make(chan struct{}) m.consecutiveFailures = 0 m.logger.Info("Token manager starting") go m.refreshLoop() } // Stop gracefully stops the token refresh background process // Blocks until stopped or context deadline exceeded func (m *Manager) Stop(ctx context.Context) error { if !m.running.Load() { m.logger.Debug("Token manager not running, nothing to stop") return nil } m.logger.Info("Token manager stopping...") // Signal stop close(m.stopCh) // Wait for goroutine to finish or timeout select { case <-m.stoppedCh: m.logger.Info("Token manager stopped gracefully") return nil case <-ctx.Done(): m.logger.Warn("Token manager stop timeout, forcing cancellation") m.cancel() // Wait a bit more for cancellation to take effect select { case <-m.stoppedCh: m.logger.Info("Token manager stopped after forced cancellation") return nil case <-time.After(100 * time.Millisecond): m.logger.Error("Token manager failed to stop cleanly") return ctx.Err() } } } // IsRunning returns true if the token manager is currently running func (m *Manager) IsRunning() bool { return m.running.Load() } // refreshLoop is the background goroutine that checks and refreshes tokens func (m *Manager) refreshLoop() { defer close(m.stoppedCh) defer m.running.Store(false) defer m.logger.Info("Token refresh loop exited") ticker := time.NewTicker(m.config.CheckInterval) defer ticker.Stop() m.logger.Info("Token refresh loop started", zap.Duration("check_interval", m.config.CheckInterval), zap.Duration("refresh_before_expiry", m.config.RefreshBeforeExpiry)) // Do initial check immediately if err := m.checkAndRefresh(); err != nil { m.logger.Error("Initial token refresh check failed", zap.Error(err)) } for { select { case <-m.stopCh: m.logger.Info("Token refresh loop received stop signal") return case <-m.ctx.Done(): m.logger.Info("Token refresh loop context cancelled") return case <-ticker.C: if err := m.checkAndRefresh(); err != nil { m.logger.Error("Token refresh check failed", zap.Error(err)) } } } } // checkAndRefresh checks if token refresh is needed and performs it func (m *Manager) checkAndRefresh() error { m.mu.Lock() defer m.mu.Unlock() // Get current session sess, err := m.getSession.Execute() if err != nil { m.logger.Debug("No session found, skipping refresh check", zap.Error(err)) return nil // Not an error - user might not be logged in } if sess == nil { m.logger.Debug("Session is nil, skipping refresh check") return nil } // Check if session is still valid if sess.IsExpired() { m.logger.Warn("Session has expired, forcing logout") return m.forceLogout() } // Check if refresh is needed timeUntilExpiry := time.Until(sess.ExpiresAt) if timeUntilExpiry > m.config.RefreshBeforeExpiry { // No refresh needed yet if m.consecutiveFailures > 0 { // Reset failure counter on successful check m.logger.Info("Session valid, resetting failure counter") m.consecutiveFailures = 0 } m.logger.Debug("Token refresh not needed yet", zap.Duration("time_until_expiry", timeUntilExpiry)) return nil } // Refresh needed m.logger.Info("Token refresh needed", zap.Duration("time_until_expiry", timeUntilExpiry)) // Attempt refresh (with background context, not the manager's context) refreshCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() if err := m.client.RefreshToken(refreshCtx); err != nil { m.consecutiveFailures++ m.logger.Error("Token refresh failed", zap.Error(err), zap.Int("consecutive_failures", m.consecutiveFailures), zap.Int("max_failures", m.config.MaxConsecutiveFailures)) if m.consecutiveFailures >= m.config.MaxConsecutiveFailures { m.logger.Error("Max consecutive refresh failures reached, forcing logout") return m.forceLogout() } return err } // Success - reset failure counter m.consecutiveFailures = 0 m.logger.Info("Token refreshed successfully", zap.Duration("time_until_old_expiry", timeUntilExpiry)) return nil } // forceLogout forces a logout due to refresh failures func (m *Manager) forceLogout() error { m.logger.Warn("Forcing logout due to token refresh issues") // Use background context since manager might be shutting down ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() if err := m.authService.Logout(ctx); err != nil { m.logger.Error("Failed to force logout", zap.Error(err)) return err } m.logger.Info("Force logout completed successfully") return nil }