at-container-registry/pkg/auth/token/handler.go

package token

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/bluesky-social/indigo/atproto/identity"
	"github.com/bluesky-social/indigo/atproto/syntax"

	"atcr.io/pkg/appview/db"
	mainAtproto "atcr.io/pkg/atproto"
	"atcr.io/pkg/auth"
	"atcr.io/pkg/auth/atproto"
)

// Handler handles /auth/token requests
type Handler struct {
	issuer              *Issuer
	validator           *atproto.SessionValidator
	deviceStore         *db.DeviceStore // For validating device secrets
	defaultHoldEndpoint string
}

// NewHandler creates a new token handler
func NewHandler(issuer *Issuer, deviceStore *db.DeviceStore, defaultHoldEndpoint string) *Handler {
	return &Handler{
		issuer:              issuer,
		validator:           atproto.NewSessionValidator(),
		deviceStore:         deviceStore,
		defaultHoldEndpoint: defaultHoldEndpoint,
	}
}

// TokenResponse represents the response from /auth/token
type TokenResponse struct {
	Token       string `json:"token,omitempty"`        // Legacy field
	AccessToken string `json:"access_token,omitempty"` // Standard field
	ExpiresIn   int    `json:"expires_in,omitempty"`
	IssuedAt    string `json:"issued_at,omitempty"`
}

// getBaseURL extracts the base URL from the request, handling proxies
func getBaseURL(r *http.Request) string {
	baseURL := r.Header.Get("X-Forwarded-Host")
	if baseURL == "" {
		baseURL = r.Host
	}
	if !strings.HasPrefix(baseURL, "http") {
		// Add scheme
		if r.TLS != nil || r.Header.Get("X-Forwarded-Proto") == "https" {
			baseURL = "https://" + baseURL
		} else {
			baseURL = "http://" + baseURL
		}
	}
	return baseURL
}

// sendAuthError sends a formatted authentication error response
func sendAuthError(w http.ResponseWriter, r *http.Request, message string) {
	baseURL := getBaseURL(r)
	w.Header().Set("WWW-Authenticate", `Basic realm="ATCR Registry"`)
	http.Error(w, fmt.Sprintf(`%s

To authenticate:
  1. Install credential helper: %s/install
  2. Or run: docker login %s
     (use your ATProto handle + app-password)`, message, baseURL, r.Host), http.StatusUnauthorized)
}

// ServeHTTP handles the token request
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	fmt.Printf("DEBUG [token/handler]: Received %s request to %s\n", r.Method, r.URL.Path)

	// Only accept GET requests (per Docker spec)
	if r.Method != http.MethodGet {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// Extract Basic auth credentials
	username, password, ok := r.BasicAuth()
	if !ok {
		fmt.Printf("DEBUG [token/handler]: No Basic auth credentials provided\n")
		sendAuthError(w, r, "authentication required")
		return
	}

	fmt.Printf("DEBUG [token/handler]: Got Basic auth for username=%s, password length=%d\n", username, len(password))

	// Parse query parameters
	_ = r.URL.Query().Get("service") // service parameter - validated by issuer
	scopeParam := r.URL.Query().Get("scope")

	// Parse scopes
	var scopes []string
	if scopeParam != "" {
		scopes = strings.Split(scopeParam, " ")
	}

	access, err := auth.ParseScope(scopes)
	if err != nil {
		http.Error(w, fmt.Sprintf("invalid scope: %v", err), http.StatusBadRequest)
		return
	}

	var did string
	var handle string
	var accessToken string

	// 1. Check if it's a device secret (starts with "atcr_device_")
	if strings.HasPrefix(password, "atcr_device_") {
		device, err := h.deviceStore.ValidateDeviceSecret(password)
		if err != nil {
			fmt.Printf("DEBUG [token/handler]: Device secret validation failed: %v\n", err)
			sendAuthError(w, r, "authentication failed")
			return
		}

		did = device.DID
		handle = device.Handle
		// Device is linked to OAuth session via DID
		// OAuth refresher will provide access token when needed via middleware
	} else {
		// 2. Try app password (direct PDS authentication)
		fmt.Printf("DEBUG [token/handler]: Not a device secret, trying app password for %s\n", username)
		did, handle, accessToken, err = h.validator.CreateSessionAndGetToken(r.Context(), username, password)
		if err != nil {
			fmt.Printf("DEBUG [token/handler]: App password validation failed: %v\n", err)
			sendAuthError(w, r, "authentication failed")
			return
		}

		fmt.Printf("DEBUG [token/handler]: App password validated successfully, DID=%s, handle=%s, AccessToken length=%d\n", did, handle, len(accessToken))

		// Cache the access token for later use (e.g., when pushing manifests)
		// TTL of 2 hours (ATProto tokens typically last longer)
		auth.GetGlobalTokenCache().Set(did, accessToken, 2*time.Hour)
		fmt.Printf("DEBUG [token/handler]: Cached access token for DID=%s\n", did)

		// Ensure user profile exists (creates with default hold if needed)
		// Resolve PDS endpoint for profile management
		directory := identity.DefaultDirectory()
		atID, err := syntax.ParseAtIdentifier(username)
		if err == nil {
			ident, err := directory.Lookup(r.Context(), *atID)
			if err != nil {
				// Log error but don't fail auth - profile management is not critical
				fmt.Printf("WARNING: failed to resolve PDS for profile management: %v\n", err)
			} else {
				pdsEndpoint := ident.PDSEndpoint()
				if pdsEndpoint != "" {
					// Create ATProto client with validated token
					atprotoClient := mainAtproto.NewClient(pdsEndpoint, did, accessToken)

					// Ensure profile exists (will create with default hold if not exists and default is configured)
					if err := mainAtproto.EnsureProfile(r.Context(), atprotoClient, h.defaultHoldEndpoint); err != nil {
						// Log error but don't fail auth - profile management is not critical
						fmt.Printf("WARNING: failed to ensure profile for %s: %v\n", did, err)
					}
				}
			}
		}
	}

	// Validate that the user has permission for the requested access
	// Use the actual handle from the validated credentials, not the Basic Auth username
	if err := auth.ValidateAccess(did, handle, access); err != nil {
		fmt.Printf("DEBUG [token/handler]: Access validation failed: %v\n", err)
		http.Error(w, fmt.Sprintf("access denied: %v", err), http.StatusForbidden)
		return
	}

	// Issue JWT token
	tokenString, err := h.issuer.Issue(did, access)
	if err != nil {
		fmt.Printf("DEBUG [token/handler]: Failed to issue token: %v\n", err)
		http.Error(w, fmt.Sprintf("failed to issue token: %v", err), http.StatusInternalServerError)
		return
	}

	fmt.Printf("DEBUG [token/handler]: Issued JWT token (length=%d) for DID=%s\n", len(tokenString), did)

	// Return token response
	now := time.Now()
	expiresIn := int(h.issuer.expiration.Seconds())

	resp := TokenResponse{
		Token:       tokenString,
		AccessToken: tokenString,
		ExpiresIn:   expiresIn,
		IssuedAt:    now.Format(time.RFC3339),
	}

	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
		return
	}
}

// RegisterRoutes registers the token handler with the provided mux
func (h *Handler) RegisterRoutes(mux *http.ServeMux) {
	mux.Handle("/auth/token", h)
}