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"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"sync"
	"time"

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

// permissionSetExpansions maps lexicon IDs to their expanded scope format.
// These must match the collections defined in lexicons/io/atcr/authFullApp.json
// Collections are sorted alphabetically for consistent comparison with PDS-expanded scopes.
var permissionSetExpansions = map[string]string{
	"io.atcr.authFullApp": "repo?" +
		"collection=io.atcr.manifest&" +
		"collection=io.atcr.repo.page&" +
		"collection=io.atcr.sailor.profile&" +
		"collection=io.atcr.sailor.star&" +
		"collection=io.atcr.tag",
}

// ExpandIncludeScopes expands any "include:" prefixed scopes to their full form
// by looking up the corresponding permission-set in the embedded lexicon files.
// For example, "include:io.atcr.authFullApp" expands to "repo?collection=io.atcr.manifest&..."
func ExpandIncludeScopes(scopes []string) []string {
	var expanded []string
	for _, scope := range scopes {
		if after, ok := strings.CutPrefix(scope, "include:"); ok {
			lexiconID := after
			if exp, ok := permissionSetExpansions[lexiconID]; ok {
				expanded = append(expanded, exp)
			} else {
				expanded = append(expanded, scope) // Keep original if unknown
			}
		} else {
			expanded = append(expanded, scope)
		}
	}
	return expanded
}

// 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)
		}

		slog.Info("Configured confidential OAuth client",
			"key_id", keyID,
			"key_path", keyPath,
		)
	} 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
}

// NewClientAppWithKey creates an indigo OAuth ClientApp with a pre-loaded P-256 key.
// Used by AppView when loading keys from the database instead of disk.
// For localhost development, privateKey is ignored (public client).
func NewClientAppWithKey(baseURL string, store oauth.ClientAuthStore, scopes []string, privateKey *atcrypto.PrivateKeyP256, clientName string) (*oauth.ClientApp, error) {
	var config oauth.ClientConfig
	redirectURI := RedirectURI(baseURL)

	if !isLocalhost(baseURL) {
		clientID := baseURL + "/oauth-client-metadata.json"
		config = oauth.NewPublicConfig(clientID, redirectURI, scopes)

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

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

		slog.Info("Configured confidential OAuth client", "key_id", keyID)
	} 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
		// See lexicons/io/atcr/authFullApp.json for definition
		"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",
		// Helm chart support
		"blob:application/vnd.cncf.helm.config.v1+json",
		"blob:application/vnd.cncf.helm.chart.content.v1.tar+gzip",
		// 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.
// Expands any "include:" prefixed scopes in the desired list before comparing,
// since the PDS returns expanded scopes in the stored session.
func ScopesMatch(stored, desired []string) bool {
	// Expand any include: scopes in desired before comparing
	expandedDesired := ExpandIncludeScopes(desired)

	// Handle nil/empty cases
	if len(stored) == 0 && len(expandedDesired) == 0 {
		return true
	}
	if len(stored) != len(expandedDesired) {
		return false
	}

	// Build map of desired scopes for O(1) lookup
	desiredMap := make(map[string]bool, len(expandedDesired))
	for _, scope := range expandedDesired {
		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)
	lockStart := time.Now()
	mutex.Lock()
	defer mutex.Unlock()
	lockWait := time.Since(lockStart)

	slog.Debug("Acquired session lock for DoWithSession",
		"component", "oauth/refresher",
		"did", did,
		"lockWait", lockWait.Round(time.Millisecond))
	if lockWait > 5*time.Second {
		slog.Warn("Slow per-DID session lock acquisition",
			"component", "oauth/refresher",
			"did", did,
			"lockWait", lockWait.Round(time.Millisecond))
	}

	// Load session while holding lock
	resumeStart := time.Now()
	session, err := r.resumeSession(ctx, did)
	if err != nil {
		return err
	}
	if resumeDur := time.Since(resumeStart); resumeDur > 5*time.Second {
		slog.Warn("Slow OAuth session resume",
			"component", "oauth/refresher",
			"did", did,
			"duration", resumeDur.Round(time.Millisecond))
	}

	// 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. Detached
		// context: once we decide to delete, the cleanup must finish even if
		// the inbound request is canceled mid-way.
		mutex.Unlock()
		_ = r.DeleteSession(context.WithoutCancel(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
}

// IsSessionInvalidError reports whether err indicates the OAuth session
// itself is invalid or revoked, i.e. deleting it (and signing the user out)
// is the right response. It is deliberately false for cancellation, deadline,
// and transport errors: deleting a session over a transient failure signs the
// user out everywhere for nothing.
func IsSessionInvalidError(err error) bool {
	if err == nil {
		return false
	}

	// A canceled or timed-out request says nothing about session validity.
	// Checked first so wrapped chains never fall through to string matching.
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
		return false
	}

	var xrpcErr *xrpc.Error
	if errors.As(err, &xrpcErr) && xrpcErr.StatusCode == 401 {
		return true
	}
	var apiErr *atclient.APIError
	if errors.As(err, &apiErr) {
		if apiErr.StatusCode == 401 {
			return true
		}
		switch apiErr.Name {
		case "InvalidToken", "InvalidGrant", "InsufficientScope":
			return true
		}
	}

	// The token-refresh failure from indigo arrives as a plain wrapped error
	// ("auth server request failed (HTTP 400): invalid_grant"), not an
	// APIError, so a string fallback is required. These substrings are
	// auth-specific and won't appear in digests or URIs.
	errStr := strings.ToLower(err.Error())
	return strings.Contains(errStr, "invalid_grant") ||
		strings.Contains(errStr, "invalid_token") ||
		strings.Contains(errStr, "insufficient_scope") ||
		strings.Contains(errStr, "token expired")
}

// isAuthError checks if an error looks like an OAuth/auth failure
// Uses structured error types to avoid false positives from substring matching
// (e.g., a digest hash containing "401" in a RecordNotFound error)
func isAuthError(err error) bool {
	if err == nil {
		return false
	}

	// Never treat cancellation/deadline as an auth failure, no matter what
	// the wrapped chain's text looks like.
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
		return false
	}

	// Check structured error types first
	var xrpcErr *xrpc.Error
	if errors.As(err, &xrpcErr) && xrpcErr.StatusCode == 401 {
		return true
	}
	var apiErr *atclient.APIError
	if errors.As(err, &apiErr) {
		if apiErr.StatusCode == 401 {
			return true
		}
		if apiErr.Name == "InvalidToken" || apiErr.Name == "InsufficientScope" {
			return true
		}
	}

	// Fallback: check for known auth-specific error strings that won't
	// appear in digests or URIs
	errStr := strings.ToLower(err.Error())
	return strings.Contains(errStr, "invalid_token") ||
		strings.Contains(errStr, "insufficient_scope") ||
		strings.Contains(errStr, "token expired")
}

// 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)
	}

	// Token-refresh POSTs rotate the refresh token on the auth server; an
	// inbound request cancellation must never abort one mid-flight, or the
	// rotated token is stranded server-side and the session dies with
	// invalid_grant "Refresh token replayed" on the next refresh.
	session.Client = newRefreshDetachClient(session.Client, session.Data.AuthServerTokenEndpoint)

	// 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) {
		// Indigo invokes this with the context of whatever request triggered
		// the refresh — possibly already canceled. Once tokens have rotated,
		// this save must complete or the session is bricked.
		saveCtx, cancel := context.WithTimeout(context.WithoutCancel(callbackCtx), 10*time.Second)
		defer cancel()
		if err := r.clientApp.Store.SaveSession(saveCtx, *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
	})
}