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

package oauth

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

	"atcr.io/pkg/atproto"
)

// UISessionStore is the interface for UI session management
// UISessionStore is defined in refresher.go to avoid duplication

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

// Server handles OAuth authorization for the AppView
type Server struct {
	app              *App
	refresher        *Refresher
	uiSessionStore   UISessionStore
	postAuthCallback PostAuthCallback
}

// NewServer creates a new OAuth server
func NewServer(app *App) *Server {
	return &Server{
		app: app,
	}
}

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

// 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.app.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, "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, fmt.Sprintf("OAuth error: %s - %s", errorParam, errorDesc))
		return
	}

	// Process OAuth callback via indigo (handles state validation internally)
	sessionData, err := s.app.ProcessCallback(r.Context(), r.URL.Query())
	if err != nil {
		s.renderError(w, 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)

	// Invalidate cached session (if any) since we have a new session with new tokens
	if s.refresher != nil {
		s.refresher.InvalidateSession(did)
		slog.Debug("Invalidated cached session after creating new session", "did", did)
	}

	// 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 {
			uiSessionID, err := store.CreateWithOAuth(did, handle, sessionData.HostURL, sessionID, 30*24*time.Hour)
			if err != nil {
				s.renderError(w, fmt.Sprintf("Failed to create UI session: %v", err))
				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
			uiSessionID, err := s.uiSessionStore.Create(did, handle, sessionData.HostURL, 30*24*time.Hour)
			if err != nil {
				s.renderError(w, fmt.Sprintf("Failed to create UI session: %v", err))
				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,
		})

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

// renderRedirectToSettings redirects to the settings page to generate an API key
func (s *Server) renderRedirectToSettings(w http.ResponseWriter, handle string) {
	tmpl := template.Must(template.New("redirect").Parse(redirectToSettingsTemplate))
	data := struct {
		Handle string
	}{
		Handle: handle,
	}

	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	if err := tmpl.Execute(w, data); err != nil {
		http.Error(w, "failed to render template", http.StatusInternalServerError)
	}
}

// renderError renders an error page
func (s *Server) renderError(w http.ResponseWriter, message string) {
	tmpl := template.Must(template.New("error").Parse(errorTemplate))
	data := struct {
		Message string
	}{
		Message: message,
	}

	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.WriteHeader(http.StatusBadRequest)
	if err := tmpl.Execute(w, data); err != nil {
		http.Error(w, "failed to render template", http.StatusInternalServerError)
	}
}

// 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>
`