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

// Package oauth provides OAuth client configuration and helper functions for ATCR.
// It provides helpers for setting up indigo's OAuth library with ATCR-specific
// configuration, including default scopes, confidential client setup, and
// interactive browser-based authentication flows.
package oauth

import (
	"context"
	"fmt"
	"log/slog"
	"strings"
	"sync"
	"time"

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

// NewClientApp creates an indigo OAuth ClientApp with ATCR-specific configuration
// Automatically configures confidential client for production deployments
// keyPath specifies where to store/load the OAuth client P-256 key (ignored for localhost)
// clientName is added to OAuth client metadata (currently unused, reserved for future)
func NewClientApp(baseURL string, store oauth.ClientAuthStore, scopes []string, keyPath string, clientName string) (*oauth.ClientApp, error) {
	var config oauth.ClientConfig
	redirectURI := RedirectURI(baseURL)

	// If production (not localhost), automatically set up confidential client
	if !isLocalhost(baseURL) {
		clientID := baseURL + "/oauth-client-metadata.json"
		config = oauth.NewPublicConfig(clientID, redirectURI, scopes)

		// Generate or load P-256 key
		privateKey, err := GenerateOrLoadClientKey(keyPath)
		if err != nil {
			return nil, fmt.Errorf("failed to load OAuth client key: %w", err)
		}

		// Generate key ID from public key
		keyID, err := GenerateKeyID(privateKey)
		if err != nil {
			return nil, fmt.Errorf("failed to generate key ID: %w", err)
		}

		// Upgrade to confidential client
		if err := config.SetClientSecret(privateKey, keyID); err != nil {
			return nil, fmt.Errorf("failed to configure confidential client: %w", err)
		}

		// Log clock information for debugging timestamp issues
		now := time.Now()
		slog.Info("Configured confidential OAuth client",
			"key_id", keyID,
			"key_path", keyPath,
			"system_time_unix", now.Unix(),
			"system_time_rfc3339", now.Format(time.RFC3339),
			"timezone", now.Location().String())
	} else {
		config = oauth.NewLocalhostConfig(redirectURI, scopes)

		slog.Info("Using public OAuth client (localhost development)")
	}

	clientApp := oauth.NewClientApp(&config, store)
	clientApp.Dir = atproto.GetDirectory()

	return clientApp, nil
}

// RedirectURI returns the OAuth redirect URI for ATCR
func RedirectURI(baseURL string) string {
	return baseURL + "/auth/oauth/callback"
}

// GetDefaultScopes returns the default OAuth scopes for ATCR registry operations.
// Includes io.atcr.authFullApp permission-set plus individual scopes for PDS compatibility.
// Blob scopes are listed explicitly (not supported in Lexicon permission-sets).
func GetDefaultScopes(did string) []string {
	return []string{
		"atproto",
		// Permission-set (for future PDS support)
		// See lexicons/io/atcr/authFullApp.json for definition
		// Uses "include:" prefix per ATProto permission spec
		"include:io.atcr.authFullApp",
		// com.atproto scopes must be separate (permission-sets are namespace-limited)
		"rpc:com.atproto.repo.getRecord?aud=*",
		// Blob scopes (not supported in Lexicon permission-sets)
		// Image manifest types (single-arch)
		"blob:application/vnd.oci.image.manifest.v1+json",
		"blob:application/vnd.docker.distribution.manifest.v2+json",
		// Manifest list/index types (multi-arch)
		"blob:application/vnd.oci.image.index.v1+json",
		"blob:application/vnd.docker.distribution.manifest.list.v2+json",
		// OCI artifact manifests (for cosign signatures, SBOMs, attestations)
		"blob:application/vnd.cncf.oras.artifact.manifest.v1+json",
		// Image avatars
		"blob:image/*",
	}
}

// ScopesMatch checks if two scope lists are equivalent (order-independent)
// Returns true if both lists contain the same scopes, regardless of order
func ScopesMatch(stored, desired []string) bool {
	// Handle nil/empty cases
	if len(stored) == 0 && len(desired) == 0 {
		return true
	}
	if len(stored) != len(desired) {
		return false
	}

	// Build map of desired scopes for O(1) lookup
	desiredMap := make(map[string]bool, len(desired))
	for _, scope := range desired {
		desiredMap[scope] = true
	}

	// Check if all stored scopes exist in desired
	for _, scope := range stored {
		if !desiredMap[scope] {
			return false
		}
	}

	return true
}

// isLocalhost checks if a base URL is a localhost address
func isLocalhost(baseURL string) bool {
	return strings.Contains(baseURL, "127.0.0.1") || strings.Contains(baseURL, "localhost")
}

// ----------------------------------------------------------------------------
// Session Management
// ----------------------------------------------------------------------------

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

// UISessionStore interface for managing UI sessions
// Shared between refresher and server
type UISessionStore interface {
	Create(did, handle, pdsEndpoint string, duration time.Duration) (string, error)
	DeleteByDID(did string)
}

// Refresher manages OAuth sessions and token refresh for AppView
// Sessions are loaded fresh from database on every request (database is source of truth)
type Refresher struct {
	clientApp      *oauth.ClientApp
	uiSessionStore UISessionStore // For invalidating UI sessions on OAuth failures
	didLocks       sync.Map       // Per-DID mutexes to prevent concurrent DPoP nonce races
}

// NewRefresher creates a new session refresher
func NewRefresher(clientApp *oauth.ClientApp) *Refresher {
	return &Refresher{
		clientApp: clientApp,
	}
}

// SetUISessionStore sets the UI session store for invalidating sessions on OAuth failures
func (r *Refresher) SetUISessionStore(store UISessionStore) {
	r.uiSessionStore = store
}

// DoWithSession executes a function with a locked OAuth session.
// The lock is held for the entire duration of the function, preventing DPoP nonce races.
//
// This is the preferred way to make PDS requests that require OAuth/DPoP authentication.
// The lock is held through the entire PDS interaction, ensuring that:
//  1. Only one goroutine at a time can negotiate DPoP nonces with the PDS for a given DID
//  2. The session's PersistSessionCallback saves the updated nonce before other goroutines load
//  3. Concurrent layer uploads don't race on stale nonces
//
// Why locking is critical:
// During docker push, multiple layers upload concurrently. Each layer creates a new
// ClientSession by loading from database. Without locking, this race condition occurs:
//  1. Layer A loads session with stale DPoP nonce from DB
//  2. Layer B loads session with same stale nonce (A hasn't updated DB yet)
//  3. Layer A makes request → 401 "use_dpop_nonce" → gets fresh nonce → saves to DB
//  4. Layer B makes request → 401 "use_dpop_nonce" (using stale nonce from step 2)
//  5. DPoP nonce thrashing continues, eventually causing 500 errors
//
// With per-DID locking:
//  1. Layer A acquires lock, loads session, handles nonce negotiation, saves, releases lock
//  2. Layer B acquires lock AFTER A releases, loads fresh nonce from DB, succeeds
//
// Example usage:
//
//	var result MyResult
//	err := refresher.DoWithSession(ctx, did, func(session *oauth.ClientSession) error {
//	    resp, err := session.DoWithAuth(session.Client, req, "com.atproto.server.getServiceAuth")
//	    if err != nil {
//	        return err
//	    }
//	    // Parse response into result...
//	    return nil
//	})
func (r *Refresher) DoWithSession(ctx context.Context, did string, fn func(session *oauth.ClientSession) error) error {
	// Get or create a mutex for this DID
	mutexInterface, _ := r.didLocks.LoadOrStore(did, &sync.Mutex{})
	mutex := mutexInterface.(*sync.Mutex)

	// Hold the lock for the ENTIRE operation (load + PDS request + nonce save)
	mutex.Lock()
	defer mutex.Unlock()

	slog.Debug("Acquired session lock for DoWithSession",
		"component", "oauth/refresher",
		"did", did)

	// Load session while holding lock
	session, err := r.resumeSession(ctx, did)
	if err != nil {
		return err
	}

	// Execute the function (PDS request) while still holding lock
	// The session's PersistSessionCallback will save nonce updates to DB
	err = fn(session)

	// If request failed with auth error, delete session to force re-auth
	if err != nil && isAuthError(err) {
		slog.Warn("Auth error detected, deleting session to force re-auth",
			"component", "oauth/refresher",
			"did", did,
			"error", err)
		// Don't hold the lock while deleting - release first
		mutex.Unlock()
		_ = r.DeleteSession(ctx, did)
		mutex.Lock() // Re-acquire for the deferred unlock
	}

	slog.Debug("Released session lock for DoWithSession",
		"component", "oauth/refresher",
		"did", did,
		"success", err == nil)

	return err
}

// isAuthError checks if an error looks like an OAuth/auth failure
func isAuthError(err error) bool {
	if err == nil {
		return false
	}
	errStr := strings.ToLower(err.Error())
	return strings.Contains(errStr, "unauthorized") ||
		strings.Contains(errStr, "invalid_token") ||
		strings.Contains(errStr, "insufficient_scope") ||
		strings.Contains(errStr, "token expired") ||
		strings.Contains(errStr, "401")
}

// resumeSession loads a session from storage
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.clientApp.Store.(sessionGetter)
	if !ok {
		return nil, fmt.Errorf("store must implement GetLatestSessionForDID (SQLite store required)")
	}

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

	// Log scope differences for debugging, but don't delete session
	// The PDS will reject requests if scopes are insufficient
	// (Permission-sets get expanded by PDS, so exact matching doesn't work)
	desiredScopes := r.clientApp.Config.Scopes
	if !ScopesMatch(sessionData.Scopes, desiredScopes) {
		slog.Debug("Session scopes differ from desired (may be permission-set expansion)",
			"did", did,
			"storedScopes", sessionData.Scopes,
			"desiredScopes", desiredScopes)
	}

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

	// Set up callback to persist token updates to SQLite
	// This ensures that when indigo automatically refreshes tokens or updates DPoP nonces,
	// the new state is saved to the database immediately
	session.PersistSessionCallback = func(callbackCtx context.Context, updatedData *oauth.ClientSessionData) {
		if err := r.clientApp.Store.SaveSession(callbackCtx, *updatedData); err != nil {
			slog.Error("Failed to persist OAuth session update",
				"component", "oauth/refresher",
				"did", did,
				"sessionID", sessionID,
				"error", err)
		} else {
			// Log session updates (token refresh, DPoP nonce updates, etc.)
			// Note: updatedData contains the full session state including DPoP nonce,
			// but we don't log sensitive data like tokens or nonces themselves
			slog.Debug("Persisted OAuth session update to database",
				"component", "oauth/refresher",
				"did", did,
				"sessionID", sessionID,
				"hint", "This includes token refresh and DPoP nonce updates")
		}
	}
	return session, nil
}

// DeleteSession removes an OAuth session from storage and optionally invalidates the UI session
// This is called when OAuth authentication fails to force re-authentication
func (r *Refresher) DeleteSession(ctx context.Context, did string) error {
	// Parse DID
	accountDID, err := syntax.ParseDID(did)
	if err != nil {
		return fmt.Errorf("failed to parse DID: %w", err)
	}

	// Get the session ID before deleting (for logging)
	type sessionGetter interface {
		GetLatestSessionForDID(ctx context.Context, did string) (*oauth.ClientSessionData, string, error)
	}

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

	_, sessionID, err := getter.GetLatestSessionForDID(ctx, did)
	if err != nil {
		// No session to delete - this is fine
		slog.Debug("No OAuth session to delete", "did", did)
		return nil
	}

	// Delete OAuth session from database
	if err := r.clientApp.Store.DeleteSession(ctx, accountDID, sessionID); err != nil {
		slog.Warn("Failed to delete OAuth session", "did", did, "sessionID", sessionID, "error", err)
		return fmt.Errorf("failed to delete OAuth session: %w", err)
	}

	slog.Info("Deleted stale OAuth session",
		"component", "oauth/refresher",
		"did", did,
		"sessionID", sessionID,
		"reason", "OAuth authentication failed")

	// Also invalidate the UI session if store is configured
	if r.uiSessionStore != nil {
		r.uiSessionStore.DeleteByDID(did)
		slog.Info("Invalidated UI session for DID",
			"component", "oauth/refresher",
			"did", did,
			"reason", "OAuth session deleted")
	}

	return nil
}

// ValidateSession checks if an OAuth session is usable by attempting to load it.
// This triggers token refresh if needed (via indigo's auto-refresh in DoWithSession).
// Returns nil if session is valid, error if session is invalid/expired/needs re-auth.
//
// This is used by the token handler to validate OAuth sessions before issuing JWTs,
// preventing the flood of errors that occurs when a stale session is discovered
// during parallel layer uploads.
func (r *Refresher) ValidateSession(ctx context.Context, did string) error {
	return r.DoWithSession(ctx, did, func(session *oauth.ClientSession) error {
		// Session loaded and refreshed successfully
		// DoWithSession already handles token refresh if needed
		slog.Debug("OAuth session validated successfully",
			"component", "oauth/refresher",
			"did", did)
		return nil
	})
}