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

package oauth

import (
	"context"
	"database/sql"
	"fmt"
	"html/template"
	"net/http"
	"strings"
	"time"

	"atcr.io/pkg/appview/db"
	"atcr.io/pkg/atproto"
	"github.com/bluesky-social/indigo/atproto/syntax"
)

// UISessionStore is the interface for UI session management
type UISessionStore interface {
	Create(did, handle, pdsEndpoint string, duration time.Duration) (string, error)
}

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

// Server handles OAuth authorization for the AppView
type Server struct {
	app                 *App
	refresher           *Refresher
	uiSessionStore      UISessionStore
	db                  *sql.DB
	defaultHoldEndpoint string
}

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

// SetDefaultHoldEndpoint sets the default hold endpoint for profile creation
func (s *Server) SetDefaultHoldEndpoint(endpoint string) {
	s.defaultHoldEndpoint = endpoint
}

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

// SetDatabase sets the database for user management
func (s *Server) SetDatabase(db *sql.DB) {
	s.db = db
}

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

	fmt.Printf("DEBUG [oauth/server]: Starting OAuth flow for handle=%s\n", handle)

	// Start auth flow via indigo
	authURL, err := s.app.StartAuthFlow(r.Context(), handle)
	if err != nil {
		fmt.Printf("ERROR [oauth/server]: Failed to start auth flow: %v\n", err)

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

	fmt.Printf("DEBUG [oauth/server]: Generated authURL=%s\n", 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

	fmt.Printf("DEBUG [oauth/server]: OAuth callback successful for DID=%s, sessionID=%s\n", did, sessionID)

	// Invalidate cached session (if any) since we have a new session with new tokens
	if s.refresher != nil {
		s.refresher.InvalidateSession(did)
		fmt.Printf("DEBUG [oauth/server]: Invalidated cached session for DID=%s after creating new session\n", did)
	}

	// Look up identity
	ident, err := s.app.directory.LookupDID(r.Context(), sessionData.AccountDID)
	handle := ident.Handle.String()
	if err != nil {
		fmt.Printf("WARNING [oauth/server]: Failed to resolve DID to handle: %v, using DID as handle\n", err)
		handle = did // Fallback to DID if resolution fails
	}

	// Fetch user's Bluesky profile (including avatar) and store in database
	if s.db != nil {
		s.fetchAndStoreAvatar(r.Context(), did, sessionID, handle, sessionData.HostURL)
	}

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

// fetchAndStoreAvatar fetches the user's Bluesky profile and stores avatar in database
func (s *Server) fetchAndStoreAvatar(ctx context.Context, did, sessionID, handle, pdsEndpoint string) {
	fmt.Printf("DEBUG [oauth/server]: Fetching avatar for DID=%s from PDS=%s\n", did, pdsEndpoint)

	// Parse DID for session resume
	didParsed, err := syntax.ParseDID(did)
	if err != nil {
		fmt.Printf("WARNING [oauth/server]: Failed to parse DID %s: %v\n", did, err)
		return
	}

	// Resume OAuth session to get authenticated client
	session, err := s.app.ResumeSession(ctx, didParsed, sessionID)
	if err != nil {
		fmt.Printf("WARNING [oauth/server]: Failed to resume session for DID=%s: %v\n", did, err)
		// Fallback: update user without avatar
		_ = db.UpsertUser(s.db, &db.User{
			DID:         did,
			Handle:      handle,
			PDSEndpoint: pdsEndpoint,
			Avatar:      "",
			LastSeen:    time.Now(),
		})
		return
	}

	// Create authenticated atproto client using the indigo session's API client
	client := atproto.NewClientWithIndigoClient(pdsEndpoint, did, session.APIClient())

	// Ensure sailor profile exists (creates with default hold if configured, or empty profile if not)
	fmt.Printf("DEBUG [oauth/server]: Ensuring profile exists for %s (defaultHold=%s)\n", did, s.defaultHoldEndpoint)
	if err := atproto.EnsureProfile(ctx, client, s.defaultHoldEndpoint); err != nil {
		fmt.Printf("WARNING [oauth/server]: Failed to ensure profile for %s: %v\n", did, err)
		// Continue anyway - profile creation is not critical for avatar fetch
	} else {
		fmt.Printf("DEBUG [oauth/server]: Profile ensured for %s\n", did)
	}

	// Fetch user's profile record from PDS (contains blob references)
	profileRecord, err := client.GetProfileRecord(ctx, did)
	if err != nil {
		fmt.Printf("WARNING [oauth/server]: Failed to fetch profile record for DID=%s: %v\n", did, err)
		// Still update user without avatar
		_ = db.UpsertUser(s.db, &db.User{
			DID:         did,
			Handle:      handle,
			PDSEndpoint: pdsEndpoint,
			Avatar:      "",
			LastSeen:    time.Now(),
		})
		return
	}

	// Construct avatar URL from blob CID using imgs.blue CDN
	var avatarURL string
	if profileRecord.Avatar != nil && profileRecord.Avatar.Ref.Link != "" {
		avatarURL = atproto.BlobCDNURL(did, profileRecord.Avatar.Ref.Link)
		fmt.Printf("DEBUG [oauth/server]: Constructed avatar URL: %s\n", avatarURL)
	}

	// Store user with avatar in database
	err = db.UpsertUser(s.db, &db.User{
		DID:         did,
		Handle:      handle,
		PDSEndpoint: pdsEndpoint,
		Avatar:      avatarURL,
		LastSeen:    time.Now(),
	})
	if err != nil {
		fmt.Printf("WARNING [oauth/server]: Failed to store user in database: %v\n", err)
		return
	}

	fmt.Printf("DEBUG [oauth/server]: Stored user with avatar for DID=%s\n", did)
}

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