at-container-registry/pkg/auth/oauth/refresher.go

package oauth

import (
	"context"
	"fmt"
	"sync"

	"github.com/bluesky-social/indigo/atproto/auth/oauth"
	"github.com/bluesky-social/indigo/atproto/syntax"
)

// SessionCache represents a cached OAuth session
type SessionCache struct {
	Session   *oauth.ClientSession
	SessionID string
}

// Refresher manages OAuth sessions and token refresh for AppView
type Refresher struct {
	app           *App
	sessions      map[string]*SessionCache // Key: DID string
	mu            sync.RWMutex
	refreshLocks  map[string]*sync.Mutex // Per-DID locks for refresh operations
	refreshLockMu sync.Mutex             // Protects refreshLocks map
}

// NewRefresher creates a new session refresher
func NewRefresher(app *App) *Refresher {
	return &Refresher{
		app:          app,
		sessions:     make(map[string]*SessionCache),
		refreshLocks: make(map[string]*sync.Mutex),
	}
}

// GetSession gets a fresh OAuth session for a DID
// Returns cached session if still valid, otherwise resumes from store
func (r *Refresher) GetSession(ctx context.Context, did string) (*oauth.ClientSession, error) {
	// Check cache first (fast path)
	r.mu.RLock()
	cached, ok := r.sessions[did]
	r.mu.RUnlock()

	if ok && cached.Session != nil {
		// Session cached, tokens will auto-refresh if needed
		return cached.Session, nil
	}

	// Session not cached, need to resume from store
	// Get or create per-DID lock to prevent concurrent resume operations
	r.refreshLockMu.Lock()
	didLock, ok := r.refreshLocks[did]
	if !ok {
		didLock = &sync.Mutex{}
		r.refreshLocks[did] = didLock
	}
	r.refreshLockMu.Unlock()

	// Acquire DID-specific lock
	didLock.Lock()
	defer didLock.Unlock()

	// Double-check cache after acquiring lock (another goroutine might have loaded it)
	r.mu.RLock()
	cached, ok = r.sessions[did]
	r.mu.RUnlock()

	if ok && cached.Session != nil {
		return cached.Session, nil
	}

	// Actually resume the session
	return r.resumeSession(ctx, did)
}

// resumeSession loads a session from storage and caches it
func (r *Refresher) resumeSession(ctx context.Context, did string) (*oauth.ClientSession, error) {
	// Parse DID
	accountDID, err := syntax.ParseDID(did)
	if err != nil {
		return nil, fmt.Errorf("failed to parse DID: %w", err)
	}

	// Get the latest session for this DID from SQLite store
	// The store must implement GetLatestSessionForDID (returns newest by updated_at)
	type sessionGetter interface {
		GetLatestSessionForDID(ctx context.Context, did string) (*oauth.ClientSessionData, string, error)
	}

	getter, ok := r.app.clientApp.Store.(sessionGetter)
	if !ok {
		return nil, fmt.Errorf("store must implement GetLatestSessionForDID (SQLite store required)")
	}

	_, sessionID, err := getter.GetLatestSessionForDID(ctx, did)
	if err != nil {
		return nil, fmt.Errorf("no session found for DID: %s", did)
	}

	// Resume session
	session, err := r.app.ResumeSession(ctx, accountDID, sessionID)
	if err != nil {
		return nil, fmt.Errorf("failed to resume session: %w", err)
	}

	// Cache the session
	r.mu.Lock()
	r.sessions[did] = &SessionCache{
		Session:   session,
		SessionID: sessionID,
	}
	r.mu.Unlock()

	return session, nil
}

// InvalidateSession removes a cached session for a DID
// This is useful when a new OAuth flow creates a fresh session
func (r *Refresher) InvalidateSession(did string) {
	r.mu.Lock()
	delete(r.sessions, did)
	r.mu.Unlock()
}

// GetSessionID returns the sessionID for a cached session
// Returns empty string if session not cached
func (r *Refresher) GetSessionID(did string) string {
	r.mu.RLock()
	defer r.mu.RUnlock()

	cached, ok := r.sessions[did]
	if !ok || cached == nil {
		return ""
	}

	return cached.SessionID
}