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

package oauth

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"html/template"
	"log/slog"
	"net/http"
	"net/url"
	"strings"
	"time"

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

// retryOnBusy retries a function up to maxAttempts times if the error
// contains "database is locked" or "database table is locked" (transient
// SQLite contention). Returns the last error if all attempts fail.
func retryOnBusy(maxAttempts int, fn func() error) error {
	var err error
	for i := range maxAttempts {
		err = fn()
		if err == nil {
			return nil
		}
		msg := err.Error()
		if !strings.Contains(msg, "database is locked") && !strings.Contains(msg, "database table is locked") {
			return err
		}
		if i < maxAttempts-1 {
			time.Sleep(time.Duration(50*(i+1)) * time.Millisecond)
		}
	}
	return err
}

// UISessionStore is the interface for UI session management
// UISessionStore is defined in client.go (session management section)

// getOAuthErrorHint provides troubleshooting hints for OAuth errors during token exchange
func getOAuthErrorHint(apiErr *atclient.APIError) string {
	switch apiErr.Name {
	case "invalid_client":
		if strings.Contains(apiErr.Message, "iat") && strings.Contains(apiErr.Message, "timestamp") {
			return "JWT timestamp validation failed - AppView system clock may be ahead of PDS clock. Check NTP sync: timedatectl status. Typical tolerance is ±30 seconds."
		}
		return "OAuth client authentication failed during token exchange - check client key and PDS OAuth configuration"
	case "invalid_grant":
		return "Authorization code is invalid, expired, or already used - user should retry OAuth flow from beginning"
	case "use_dpop_nonce":
		return "DPoP nonce challenge during token exchange - indigo should retry automatically, persistent failures indicate PDS issue"
	case "invalid_dpop_proof":
		return "DPoP proof validation failed - check system clock sync between AppView and PDS"
	case "unauthorized_client":
		return "PDS rejected the client - check client metadata URL is accessible and scopes are supported"
	case "invalid_request":
		return "Malformed token request - check OAuth flow parameters (code, redirect_uri, state)"
	case "server_error":
		return "PDS internal error during token exchange - check PDS logs for root cause"
	default:
		if apiErr.StatusCode == 400 {
			return "Bad request during OAuth token exchange - check error details and PDS logs"
		}
		return "OAuth token exchange failed - see errorName and errorMessage for PDS response"
	}
}

// UserStore is the interface for user management
type UserStore interface {
	UpsertUser(did, handle, pdsEndpoint, avatar string) error
}

// PostAuthCallback is called after successful OAuth authentication.
// Parameters: ctx, did, handle, pdsEndpoint, sessionID
// This allows AppView to perform business logic (profile creation, avatar fetch, etc.)
// without coupling the OAuth package to AppView-specific dependencies.
type PostAuthCallback func(ctx context.Context, did, handle, pdsEndpoint, sessionID string) error

// PageRenderer renders the HTML pages the OAuth endpoints show to a person in
// a browser, so an embedding application (the appview) can draw them in its
// own site layout. Each method returns the complete page body; Server owns the
// status code and headers. When a method returns an error, or no renderer is
// set, Server falls back to the plain inline templates in this file.
type PageRenderer interface {
	// RenderOAuthError renders a failed authorization. message is plain text.
	RenderOAuthError(r *http.Request, message string) ([]byte, error)
	// RenderMissingScopes renders the refusal shown when the PDS granted
	// fewer scopes than requested. retryURL restarts the flow.
	RenderMissingScopes(r *http.Request, missing []string, retryURL string) ([]byte, error)
	// RenderAuthorized renders the success page for the non-UI flow, which
	// forwards to /settings.
	RenderAuthorized(r *http.Request, handle string) ([]byte, error)
}

// Server handles OAuth authorization for the AppView
type Server struct {
	clientApp        *oauth.ClientApp
	refresher        *Refresher
	uiSessionStore   UISessionStore
	postAuthCallback PostAuthCallback
	pages            PageRenderer
}

// NewServer creates a new OAuth server
func NewServer(clientApp *oauth.ClientApp) *Server {
	return &Server{
		clientApp: clientApp,
	}
}

// SetRefresher sets the refresher for invalidating session cache
func (s *Server) SetRefresher(refresher *Refresher) {
	s.refresher = refresher
}

// SetUISessionStore sets the UI session store for web login
func (s *Server) SetUISessionStore(store UISessionStore) {
	s.uiSessionStore = store
}

// SetPostAuthCallback sets the callback to be invoked after successful OAuth authentication
// This allows AppView to inject business logic without coupling the OAuth package
func (s *Server) SetPostAuthCallback(callback PostAuthCallback) {
	s.postAuthCallback = callback
}

// SetPageRenderer sets the renderer for the browser-facing OAuth pages. Pass
// nil to use the built-in inline templates.
func (s *Server) SetPageRenderer(pages PageRenderer) {
	s.pages = pages
}

// ServeAuthorize handles GET /auth/oauth/authorize
func (s *Server) ServeAuthorize(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// Get handle from query parameter
	handle := r.URL.Query().Get("handle")
	if handle == "" {
		http.Error(w, "handle parameter required", http.StatusBadRequest)
		return
	}

	slog.Debug("Starting OAuth flow", "handle", handle)

	// Start auth flow via indigo
	authURL, err := s.clientApp.StartAuthFlow(r.Context(), handle)
	if err != nil {
		slog.Error("Failed to start auth flow", "error", err, "handle", handle)

		// Check if error is about invalid_client_metadata (usually means PDS doesn't support required scopes)
		errMsg := err.Error()
		if strings.Contains(errMsg, "invalid_client_metadata") {
			s.renderError(w, r, "OAuth authorization failed: Your PDS does not support one or more required OAuth scopes (likely the 'blob:' scope). Please update your PDS to the latest version and try again.")
			return
		}

		http.Error(w, fmt.Sprintf("failed to start auth flow: %v", err), http.StatusInternalServerError)
		return
	}

	slog.Debug("Generated OAuth authorization URL", "authURL", authURL)

	// Redirect to PDS authorization page
	// Note: indigo handles state internally via the auth store
	http.Redirect(w, r, authURL, http.StatusFound)
}

// ServeCallback handles GET /auth/oauth/callback
func (s *Server) ServeCallback(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// Check for OAuth error
	if errorParam := r.URL.Query().Get("error"); errorParam != "" {
		errorDesc := r.URL.Query().Get("error_description")
		s.renderError(w, r, fmt.Sprintf("OAuth error: %s - %s", errorParam, errorDesc))
		return
	}

	// Process OAuth callback via indigo (handles state validation internally)
	// This performs token exchange with the PDS using authorization code
	sessionData, err := s.clientApp.ProcessCallback(r.Context(), r.URL.Query())
	if err != nil {
		// Detailed error logging for token exchange failures
		var apiErr *atclient.APIError
		if errors.As(err, &apiErr) {
			slog.Error("OAuth callback failed - token exchange error",
				"component", "oauth/server",
				"error", err,
				"httpStatus", apiErr.StatusCode,
				"errorName", apiErr.Name,
				"errorMessage", apiErr.Message,
				"hint", getOAuthErrorHint(apiErr),
				"queryParams", r.URL.Query().Encode())
		} else {
			slog.Error("OAuth callback failed - unknown error",
				"component", "oauth/server",
				"error", err,
				"errorType", fmt.Sprintf("%T", err),
				"queryParams", r.URL.Query().Encode())
		}

		s.renderError(w, r, fmt.Sprintf("Failed to process OAuth callback: %v", err))
		return
	}

	did := sessionData.AccountDID.String()
	sessionID := sessionData.SessionID

	slog.Debug("OAuth callback successful", "did", did, "sessionID", sessionID)

	// Refuse a partial grant. Some PDSes let the user untick permissions, or
	// leave one off the consent screen, and a session missing one works until
	// the first push that needs it. Signing in again only helps once the PDS
	// offers the full set, so say what's missing now. This runs before the
	// old-session cleanup below so a refused login leaves any working session
	// the user already had untouched.
	if missing := MissingScopes(sessionData.Scopes, s.clientApp.Config.Scopes); len(missing) > 0 {
		slog.Warn("Refusing OAuth login: PDS did not grant all requested scopes",
			"component", "oauth/server",
			"did", did,
			"host", sessionData.HostURL,
			"missing", missing,
			"granted", sessionData.Scopes)
		if err := s.clientApp.Store.DeleteSession(r.Context(), sessionData.AccountDID, sessionID); err != nil {
			slog.Warn("Failed to delete refused OAuth session", "did", did, "error", err)
		}
		s.renderMissingScopes(w, r, missing)
		return
	}

	// Clean up old OAuth sessions for this DID BEFORE invalidating cache
	// This prevents accumulation of stale sessions with expired refresh tokens
	// Order matters: delete from DB first, then invalidate cache, so when cache reloads
	// it will only find the new session
	type sessionCleaner interface {
		DeleteOldSessionsForDID(ctx context.Context, did string, keepSessionID string) error
	}
	if cleaner, ok := s.clientApp.Store.(sessionCleaner); ok {
		if err := cleaner.DeleteOldSessionsForDID(r.Context(), did, sessionID); err != nil {
			slog.Warn("Failed to clean up old OAuth sessions", "did", did, "error", err)
			// Non-fatal - log and continue
		} else {
			slog.Debug("Cleaned up old OAuth sessions", "did", did, "kept", sessionID)
		}
	}

	// Look up identity (resolve DID to handle)
	_, handle, _, err := atproto.ResolveIdentity(r.Context(), did)
	if err != nil {
		slog.Warn("Failed to resolve DID to handle, using DID as fallback", "error", err, "did", did)
		handle = did // Fallback to DID if resolution fails
	}

	// Call post-auth callback for AppView business logic (profile, avatar, etc.)
	if s.postAuthCallback != nil {
		if err := s.postAuthCallback(r.Context(), did, handle, sessionData.HostURL, sessionID); err != nil {
			// Log error but don't fail OAuth flow - business logic is non-critical
			slog.Warn("Post-auth callback failed", "error", err, "did", did)
		}
	}

	// Check if this is a UI login (has oauth_return_to cookie)
	if cookie, err := r.Cookie("oauth_return_to"); err == nil && s.uiSessionStore != nil {
		// Create UI session (30 days to match OAuth refresh token lifetime)
		// Store OAuth sessionID so we can resume it on next login
		if store, ok := s.uiSessionStore.(interface {
			CreateWithOAuth(did, handle, pdsEndpoint, oauthSessionID string, duration time.Duration) (string, error)
		}); ok {
			var uiSessionID string
			err := retryOnBusy(3, func() error {
				var createErr error
				uiSessionID, createErr = store.CreateWithOAuth(did, handle, sessionData.HostURL, sessionID, 30*24*time.Hour)
				return createErr
			})
			if err != nil {
				slog.Error("Failed to create UI session", "error", err, "did", did)
				s.renderError(w, r, "Something went wrong while logging you in. Please try again.")
				return
			}
			// Set UI session cookie and redirect (code below)
			// Note: Secure flag depends on the request scheme (HTTP vs HTTPS)
			http.SetCookie(w, &http.Cookie{
				Name:     "atcr_session",
				Value:    uiSessionID,
				Path:     "/",
				MaxAge:   30 * 86400, // 30 days
				HttpOnly: true,
				Secure:   r.URL.Scheme == "https" || r.Header.Get("X-Forwarded-Proto") == "https",
				SameSite: http.SameSiteLaxMode,
			})
		} else {
			// Fallback for stores that don't support OAuth sessionID
			var uiSessionID string
			err := retryOnBusy(3, func() error {
				var createErr error
				uiSessionID, createErr = s.uiSessionStore.Create(did, handle, sessionData.HostURL, 30*24*time.Hour)
				return createErr
			})
			if err != nil {
				slog.Error("Failed to create UI session", "error", err, "did", did)
				s.renderError(w, r, "Something went wrong while logging you in. Please try again.")
				return
			}
			// Set UI session cookie
			// Note: Secure flag depends on the request scheme (HTTP vs HTTPS)
			http.SetCookie(w, &http.Cookie{
				Name:     "atcr_session",
				Value:    uiSessionID,
				Path:     "/",
				MaxAge:   30 * 86400, // 30 days
				HttpOnly: true,
				Secure:   r.URL.Scheme == "https" || r.Header.Get("X-Forwarded-Proto") == "https",
				SameSite: http.SameSiteLaxMode,
			})
		}

		// Clear the return_to cookie
		http.SetCookie(w, &http.Cookie{
			Name:     "oauth_return_to",
			Value:    "",
			Path:     "/",
			MaxAge:   -1,
			HttpOnly: true,
		})

		// Set a JS-readable cookie with the handle for "recent accounts" feature
		// Frontend will read this, save to localStorage, and delete the cookie
		http.SetCookie(w, &http.Cookie{
			Name:     "atcr_login_handle",
			Value:    handle,
			Path:     "/",
			MaxAge:   60, // Short-lived, just for the redirect
			HttpOnly: false,
			Secure:   r.URL.Scheme == "https" || r.Header.Get("X-Forwarded-Proto") == "https",
			SameSite: http.SameSiteLaxMode,
		})

		// Redirect to return URL
		returnTo := cookie.Value
		if returnTo == "" {
			returnTo = "/"
		}
		http.Redirect(w, r, returnTo, http.StatusFound)
		return
	}

	// Non-UI flow: redirect to settings to get API key
	s.renderRedirectToSettings(w, r, handle)
}

// writePage writes an HTML page with the given status. It prefers the
// injected renderer and falls back to the inline template when there is none
// or it fails.
func (s *Server) writePage(w http.ResponseWriter, status int, render func(PageRenderer) ([]byte, error), fallbackName, fallback string, data any) {
	var body []byte
	if s.pages != nil {
		b, err := render(s.pages)
		if err == nil {
			body = b
		} else {
			slog.Warn("OAuth page renderer failed, using fallback template",
				"component", "oauth/server", "page", fallbackName, "error", err)
		}
	}
	if body == nil {
		var buf bytes.Buffer
		tmpl := template.Must(template.New(fallbackName).Parse(fallback))
		if err := tmpl.Execute(&buf, data); err != nil {
			http.Error(w, "failed to render template", http.StatusInternalServerError)
			return
		}
		body = buf.Bytes()
	}

	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.WriteHeader(status)
	_, _ = w.Write(body)
}

// renderRedirectToSettings renders the success page for the non-UI flow,
// which forwards to the settings page.
func (s *Server) renderRedirectToSettings(w http.ResponseWriter, r *http.Request, handle string) {
	data := struct {
		Handle string
	}{
		Handle: handle,
	}
	s.writePage(w, http.StatusOK, func(p PageRenderer) ([]byte, error) {
		return p.RenderAuthorized(r, handle)
	}, "redirect", redirectToSettingsTemplate, data)
}

// renderError renders an error page
func (s *Server) renderError(w http.ResponseWriter, r *http.Request, message string) {
	data := struct {
		Message string
	}{
		Message: message,
	}
	s.writePage(w, http.StatusBadRequest, func(p PageRenderer) ([]byte, error) {
		return p.RenderOAuthError(r, message)
	}, "error", errorTemplate, data)
}

// renderMissingScopes renders the page shown when the PDS granted fewer
// permissions than ATCR asked for.
//
// "Try again" goes back through the login page, not straight to authorize: the
// login form is what sets the oauth_return_to cookie that gets a UI session
// created, and it carries any return_to in flight (the device page, during a
// credential-helper login) through to the retry.
func (s *Server) renderMissingScopes(w http.ResponseWriter, r *http.Request, missing []string) {
	retryURL := "/auth/oauth/login"
	if c, err := r.Cookie("oauth_return_to"); err == nil && c.Value != "" {
		retryURL += "?return_to=" + url.QueryEscape(c.Value)
	}
	data := struct {
		Missing  []string
		RetryURL string
	}{
		Missing:  missing,
		RetryURL: retryURL,
	}
	s.writePage(w, http.StatusForbidden, func(p PageRenderer) ([]byte, error) {
		return p.RenderMissingScopes(r, data.Missing, data.RetryURL)
	}, "missing-scopes", missingScopesTemplate, data)
}

// HTML templates

const redirectToSettingsTemplate = `
<!DOCTYPE html>
<html>
<head>
    <title>Authorization Successful - ATCR</title>
    <meta http-equiv="refresh" content="3;url=/settings">
    <style>
        body { font-family: sans-serif; max-width: 600px; margin: 50px auto; padding: 20px; }
        .success { background: #d4edda; border: 1px solid #c3e6cb; padding: 20px; border-radius: 5px; }
        .info { background: #d1ecf1; border: 1px solid #bee5eb; padding: 15px; border-radius: 5px; margin-top: 15px; }
        a { color: #007bff; text-decoration: none; }
        a:hover { text-decoration: underline; }
    </style>
</head>
<body>
    <div class="success">
        <h1>✓ Authorization Successful!</h1>
        <p>You have successfully authorized ATCR to access your ATProto account: <strong>{{.Handle}}</strong></p>
        <p>Redirecting to settings page to generate your API key...</p>
        <p>If not redirected, <a href="/settings">click here</a>.</p>
    </div>
    <div class="info">
        <h3>Next Steps:</h3>
        <ol>
            <li>Generate an API key on the settings page</li>
            <li>Copy the API key (shown once!)</li>
            <li>Use it with: <code>docker login atcr.io -u {{.Handle}} -p [your-api-key]</code></li>
        </ol>
    </div>
</body>
</html>
`

const errorTemplate = `
<!DOCTYPE html>
<html>
<head>
    <title>Authorization Failed - ATCR</title>
    <style>
        body { font-family: sans-serif; max-width: 600px; margin: 50px auto; padding: 20px; }
        .error { background: #f8d7da; border: 1px solid #f5c6cb; padding: 20px; border-radius: 5px; }
    </style>
</head>
<body>
    <div class="error">
        <h1>✗ Authorization Failed</h1>
        <p>{{.Message}}</p>
        <p><a href="/">Return to home</a></p>
    </div>
</body>
</html>
`

const missingScopesTemplate = `
<!DOCTYPE html>
<html>
<head>
    <title>Missing Permissions - ATCR</title>
    <style>
        body { font-family: sans-serif; max-width: 600px; margin: 50px auto; padding: 20px; }
        .error { background: #f8d7da; border: 1px solid #f5c6cb; padding: 20px; border-radius: 5px; }
        code { font-size: 0.9em; word-break: break-all; }
    </style>
</head>
<body>
    <div class="error">
        <h1>Some permissions weren't granted</h1>
        <p>Your PDS signed you in, but didn't grant everything ATCR needs to push and manage images. You haven't been logged in.</p>
        <p>Missing:</p>
        <ul>
            {{range .Missing}}<li><code>{{.}}</code></li>{{end}}
        </ul>
        <p>If your PDS let you untick permissions, try again and approve all of them. If a permission wasn't shown to you at all, your PDS may not support it yet. Updating the PDS usually fixes this.</p>
        <p><a href="{{.RetryURL}}">Try again</a> · <a href="/">Return to home</a></p>
    </div>
</body>
</html>
`