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

package token

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"strings"
	"time"

	"atcr.io/pkg/appview/db"
	"atcr.io/pkg/atproto"
	"atcr.io/pkg/auth"
	"github.com/distribution/distribution/v3/registry/api/errcode"
	"github.com/go-chi/render"
)

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

// OAuthSessionValidator validates OAuth sessions before issuing tokens
// This interface allows the token handler to verify OAuth sessions are usable
// (not just that they exist) without depending directly on the OAuth implementation.
type OAuthSessionValidator interface {
	// ValidateSession checks if OAuth session is usable by attempting to load/refresh it
	// Returns nil if session is valid, error if session is invalid/expired/needs re-auth
	ValidateSession(ctx context.Context, did string) error
}

// Authorizer gates issuance of registry JWTs at the auth phase.
// Implementations decide which sub-checks apply based on `access`:
// crew reconciliation runs for any token request (so first-time CLI users
// can pull from a private hold), while membership and quota enforcement
// only apply to non-wildcard push scopes. The Docker spec model is to
// embed authorization in the JWT and trust it for its short lifetime; this
// is the single point where those gates run.
type Authorizer interface {
	// Authorize is called once per token issuance, after the requester's
	// ATProto credentials have been validated. Returns nil if `access` is
	// allowed, or an error describing the denial reason. The error message
	// surfaces to the OCI client inside the distribution error JSON body.
	//
	// authMethod is one of AuthMethodOAuth or AuthMethodAppPassword and lets
	// the implementation pick the right service-token fetcher when it needs
	// to talk to the hold (e.g. for crew reconciliation).
	Authorize(ctx context.Context, did, authMethod string, access []auth.AccessEntry) error
}

// ServiceAuthFetcher pre-mints the AppView↔hold service-auth at /auth/token
// time so the registry JWT can be bound to its lifetime. JWT and service-auth
// then expire concurrently; when Docker hits 401, the next /auth/token call
// mints both fresh in lockstep.
type ServiceAuthFetcher interface {
	// Fetch ensures a service-auth exists for (did, hold derived from did)
	// and returns its expiry. Returns (zero time, nil) when the user has no
	// hold configured — caller falls back to the issuer's default exp.
	Fetch(ctx context.Context, did, authMethod string) (expiresAt time.Time, err error)
}

// Handler handles /auth/token requests
type Handler struct {
	issuer                *Issuer
	validator             *auth.SessionValidator
	deviceStore           *db.DeviceStore // For validating device secrets
	postAuthCallback      PostAuthCallback
	oauthSessionValidator OAuthSessionValidator
	authorizer            Authorizer
	serviceAuthFetcher    ServiceAuthFetcher
	// services is the set of registry domains this AppView fronts, keyed by
	// normalized hostname. Nil means single-domain: the lookups in
	// resolveService miss and every token gets the issuer's own service.
	services map[string]bool
}

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

// SetPostAuthCallback sets the callback to be invoked after successful Basic Auth authentication
// This allows AppView to inject business logic without coupling the token package
func (h *Handler) SetPostAuthCallback(callback PostAuthCallback) {
	h.postAuthCallback = callback
}

// SetOAuthSessionValidator sets the OAuth session validator for validating device auth
// When set, the handler will validate OAuth sessions are usable before issuing tokens for device auth
// This prevents the flood of errors that occurs when a stale session is discovered during push
func (h *Handler) SetOAuthSessionValidator(validator OAuthSessionValidator) {
	h.oauthSessionValidator = validator
}

// SetAuthorizer wires the auth-phase gate. When set, every token request
// runs Authorize after credentials validate; a non-nil error is returned to
// the client as a 403 (distribution error JSON) and the JWT is not issued.
func (h *Handler) SetAuthorizer(authorizer Authorizer) {
	h.authorizer = authorizer
}

// SetServiceAuthFetcher binds JWT issuance to the AppView↔hold service-auth.
// When set, the handler pre-mints the service-auth and stamps the JWT's exp
// from the cached expiry, so both tokens expire concurrently.
func (h *Handler) SetServiceAuthFetcher(fetcher ServiceAuthFetcher) {
	h.serviceAuthFetcher = fetcher
}

// SetServices declares the registry domains this AppView fronts, e.g.
// ["buoy.cr", "seamark.cr", "atcr.io"]. Each issued JWT is stamped with
// whichever of these the client is authenticating against, so the audience
// names the front door actually used (see pkg/appview/registryauth). Unset
// leaves every token on the issuer's configured service, which is correct for
// a single-domain deployment.
func (h *Handler) SetServices(services []string) {
	if len(services) == 0 {
		h.services = nil
		return
	}
	set := make(map[string]bool, len(services))
	for _, s := range services {
		if n := NormalizeService(s); n != "" {
			set[n] = true
		}
	}
	h.services = set
}

// resolveService picks the registry domain to stamp as the JWT's audience.
//
// The requested service is the primary signal: Docker echoes back whatever the
// WWW-Authenticate challenge advertised, and that challenge is built per front
// door. It arrives in the ?service= query parameter on the GET form and in the
// service form field on the POST form, so callers extract it per method and
// pass it in. The request's own host is the fallback, which covers clients that
// reach /auth/token directly on a registry domain rather than via the realm.
// Both are client-influenced, so both are only honoured when they name a
// configured registry domain; anything else falls back to the issuer's service.
// That makes the worst case a token scoped to the primary domain, not a
// caller-chosen audience.
func (h *Handler) resolveService(r *http.Request, requested string) string {
	if s := NormalizeService(requested); h.services[s] {
		return s
	}
	if s := NormalizeService(r.Host); h.services[s] {
		return s
	}
	return h.issuer.service
}

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

// AuthErrorResponse is returned when authentication fails in a way the credential helper can handle
type AuthErrorResponse struct {
	Error    string `json:"error"`
	Message  string `json:"message"`
	LoginURL string `json:"login_url,omitempty"`
}

// sendOAuthSessionExpiredError sends a JSON error response when OAuth session is missing
// This allows the credential helper to detect this specific error and open the browser
func sendOAuthSessionExpiredError(w http.ResponseWriter, r *http.Request) {
	baseURL := getBaseURL(r)
	loginURL := baseURL + "/auth/oauth/login"

	w.Header().Set("WWW-Authenticate", `Basic realm="ATCR Registry"`)
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusUnauthorized)

	resp := AuthErrorResponse{
		Error:    "oauth_session_expired",
		Message:  "OAuth session expired or invalidated. Please re-authenticate in your browser.",
		LoginURL: loginURL,
	}
	render.JSON(w, r, resp)
}

// slowPhaseThreshold flags /auth/token phases that take long enough to risk
// the Docker client's ~15s deadline; production incidents showed ~14s stalls
// before the first PDS call, and these phase timings exist to attribute them.
const slowPhaseThreshold = 5 * time.Second

// oauthError is the RFC 6749 section 5.2 error body used by the POST form.
type oauthError struct {
	Error            string `json:"error"`
	ErrorDescription string `json:"error_description,omitempty"`
}

// writeOAuthError responds in the shape the OAuth2 token spec defines. Only the
// POST form uses it; the GET form keeps its plain-text guidance, which is what
// users actually see when they run docker login by hand.
func writeOAuthError(w http.ResponseWriter, status int, code, description string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	if err := json.NewEncoder(w).Encode(oauthError{Error: code, ErrorDescription: description}); err != nil {
		slog.Warn("failed to write OAuth error response", "error", err)
	}
}

// ServeHTTP handles the token request.
//
// Both Docker token specs land here. The original spec is GET with HTTP Basic
// and the scope in the query string. The OAuth2 spec is POST with a form-encoded
// body; containerd and Docker try it first whenever they hold a secret and only
// fall back to the GET form on 404/401/405, so answering it saves every such
// client a wasted round trip. Once credentials and scope are extracted the two
// paths are identical.
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	phaseStart := time.Now()
	slog.Debug("Received token request", "method", r.Method, "path", r.URL.Path)

	var username, password, scopeParam, requestedService string

	switch r.Method {
	case http.MethodGet:
		var ok bool
		username, password, ok = r.BasicAuth()
		if !ok {
			slog.Debug("No Basic auth credentials provided")
			sendAuthError(w, r, "authentication required")
			return
		}
		scopeParam = r.URL.Query().Get("scope")
		requestedService = r.URL.Query().Get("service")

	case http.MethodPost:
		if err := r.ParseForm(); err != nil {
			writeOAuthError(w, http.StatusBadRequest, "invalid_request", "malformed request body")
			return
		}
		// Only the password grant is supported, and no refresh token is issued.
		// The registry JWT's lifetime is deliberately pinned to the AppView<->hold
		// service-auth (see SetServiceAuthFetcher), so a refresh token would be a
		// fourth long-lived credential with its own storage and revocation.
		// Clients handle its absence by continuing to use the credential they hold.
		//
		// The status here is deliberately 401 rather than the 400 that RFC 6749
		// section 5.2 prescribes. containerd only sends grant_type=refresh_token
		// when it has no username, which is the same condition that disables its
		// 405 fallback, so a 400 would hard-fail those clients. 401 is on its
		// retry list, sending them to the GET form, where a device secret
		// authenticates off the password alone and succeeds.
		if grant := r.PostFormValue("grant_type"); grant != "" && grant != "password" {
			writeOAuthError(w, http.StatusUnauthorized, "unsupported_grant_type",
				fmt.Sprintf("grant_type %q is not supported, use password", grant))
			return
		}
		username = r.PostFormValue("username")
		password = r.PostFormValue("password")
		if username == "" || password == "" {
			// 401 is one of the statuses clients retry on the GET form, where the
			// plain-text guidance in sendAuthError is waiting for them.
			writeOAuthError(w, http.StatusUnauthorized, "invalid_client",
				"username and password are required")
			return
		}
		scopeParam = r.PostFormValue("scope")
		requestedService = r.PostFormValue("service")

	default:
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// Reconstruct DID usernames that were mangled by BasicAuth's colon split.
	// Form bodies carry the username as a discrete field so only the
	// hyphen-encoding case can fire there, but running both keeps the two paths
	// accepting exactly the same set of usernames.
	username, password = parseBasicAuthDID(username, password)

	slog.Debug("Got credentials", "username", username, "passwordLength", len(password))

	// The service names the front door the client is authenticating against and
	// becomes the JWT's audience; resolveService validates it against the
	// configured registry domains.
	service := h.resolveService(r, requestedService)

	// 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
	var authMethod 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 {
			slog.Debug("Device secret validation failed", "error", err)
			sendAuthError(w, r, "authentication failed")
			return
		}

		// Validate OAuth session is usable (not just exists)
		// Device secrets are permanent, but they require a working OAuth session to push
		// By validating here, we prevent the flood of errors that occurs when a stale
		// session is discovered during parallel layer uploads
		if h.oauthSessionValidator != nil {
			if err := h.oauthSessionValidator.ValidateSession(r.Context(), device.DID); err != nil {
				slog.Debug("OAuth session validation failed", "did", device.DID, "error", err)
				sendOAuthSessionExpiredError(w, r)
				return
			}
		}

		did = device.DID
		handle = device.Handle
		authMethod = AuthMethodOAuth
		// 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)
		slog.Debug("Trying app password authentication", "username", username)
		did, handle, accessToken, err = h.validator.CreateSessionAndGetToken(r.Context(), username, password)
		if err != nil {
			// Log at WARN level with specific error type
			if errors.Is(err, auth.ErrIdentityResolution) {
				slog.Warn("Identity resolution failed", "error", err, "username", username)
				sendAuthError(w, r, "authentication failed: could not resolve handle")
			} else if errors.Is(err, auth.ErrInvalidCredentials) {
				slog.Warn("Invalid credentials", "username", username)
				sendAuthError(w, r, "authentication failed: invalid credentials")
			} else if errors.Is(err, auth.ErrPDSUnavailable) {
				slog.Warn("PDS unavailable", "error", err, "username", username)
				sendAuthError(w, r, "authentication failed: PDS unavailable")
			} else {
				slog.Warn("Authentication failed", "error", err, "username", username)
				sendAuthError(w, r, "authentication failed")
			}
			return
		}

		authMethod = AuthMethodAppPassword

		slog.Debug("App password validated successfully",
			"did", did,
			"handle", handle,
			"accessTokenLength", 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)
		slog.Debug("Cached access token", "did", did)

		// Call post-auth callback for AppView business logic (profile management, etc.)
		if h.postAuthCallback != nil {
			// Resolve PDS endpoint for callback
			_, _, pdsEndpoint, err := atproto.ResolveIdentity(r.Context(), username)
			if err != nil {
				// Log error but don't fail auth - profile management is not critical
				slog.Warn("Failed to resolve PDS for callback", "error", err, "username", username)
			} else {
				if err := h.postAuthCallback(r.Context(), did, handle, pdsEndpoint, accessToken); err != nil {
					// Log error but don't fail auth - business logic is non-critical
					slog.Warn("Post-auth callback failed", "error", err, "did", did)
				}
			}
		}
	}

	// Credential phase covers device-secret/app-password validation including
	// any OAuth session validation (and its lazy token refresh).
	credDur := time.Since(phaseStart)
	if credDur > slowPhaseThreshold {
		slog.Warn("slow /auth/token phase",
			"phase", "credentials",
			"did", did,
			"authMethod", authMethod,
			"duration", credDur.Round(time.Millisecond))
	}

	// 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 {
		slog.Debug("Access validation failed", "error", err, "did", did)
		http.Error(w, fmt.Sprintf("access denied: %v", err), http.StatusForbidden)
		return
	}

	// Auth-phase gate (crew reconciliation, plus membership/quota for
	// non-wildcard push) and service-auth pre-mint are independent network
	// paths — one talks to the hold, the other to the PDS. Run them in
	// parallel so a cold /auth/token pays max(both) instead of sum(both).
	// Nil-authorizer / nil-fetcher configurations short-circuit each branch
	// independently. The gate runs for pull too: first-time CLI users on a
	// private hold need crew reconciliation before the hold-side read check
	// will accept the resulting JWT.
	runGate := h.authorizer != nil
	runFetch := h.serviceAuthFetcher != nil

	type gateRes struct{ err error }
	type fetchRes struct {
		expiresAt time.Time
		err       error
	}
	var (
		gateCh  chan gateRes
		fetchCh chan fetchRes
	)
	if runGate {
		gateCh = make(chan gateRes, 1)
		go func() {
			gateCh <- gateRes{err: h.authorizer.Authorize(r.Context(), did, authMethod, access)}
		}()
	}
	if runFetch {
		fetchCh = make(chan fetchRes, 1)
		go func() {
			exp, err := h.serviceAuthFetcher.Fetch(r.Context(), did, authMethod)
			fetchCh <- fetchRes{expiresAt: exp, err: err}
		}()
	}

	// Drain the gate first so a denial wins over a transient fetch error.
	drainStart := time.Now()
	var gateDur time.Duration
	if runGate {
		res := <-gateCh
		gateDur = time.Since(drainStart)
		if gateDur > slowPhaseThreshold {
			slog.Warn("slow /auth/token phase",
				"phase", "gate",
				"did", did,
				"duration", gateDur.Round(time.Millisecond))
		}
		if res.err != nil {
			slog.Info("Authorization denied", "did", did, "error", res.err)
			_ = errcode.ServeJSON(w, errcode.ErrorCodeDenied.WithMessage(res.err.Error()))
			return
		}
	}

	// Bind JWT lifetime to the AppView↔hold service-auth: pre-mint it now and
	// stamp the JWT's exp from the cached expiry. They expire concurrently;
	// when Docker hits 401, the next /auth/token call mints both fresh.
	issueExp := h.issuer.expiration
	var fetchDur time.Duration
	if runFetch {
		res := <-fetchCh
		// Both goroutines started together, so measure from the drain start,
		// not after the gate drain, to reflect the fetch's real wall time.
		fetchDur = time.Since(drainStart)
		if fetchDur > slowPhaseThreshold {
			slog.Warn("slow /auth/token phase",
				"phase", "service-auth-fetch",
				"did", did,
				"duration", fetchDur.Round(time.Millisecond))
		}
		if res.err != nil {
			slog.Warn("service-auth pre-mint failed", "did", did, "error", res.err)
			_ = errcode.ServeJSON(w, errcode.ErrorCodeUnavailable.WithMessage(fmt.Sprintf("service-auth fetch failed: %v", res.err)))
			return
		}
		if !res.expiresAt.IsZero() {
			// Cap JWT lifetime at the service-auth's expiry. The cache's
			// expiresAt already includes a 10s safety margin
			// (pkg/auth/cache.go:71), so this guarantees the service-auth
			// is still cache-valid for any /v2/* request the JWT can
			// authorize. We never extend beyond the configured default.
			until := time.Until(res.expiresAt)
			if until < issueExp {
				issueExp = until
			}
		}
	}

	// Issue JWT token
	tokenString, err := h.issuer.IssueWithExpiration(did, access, authMethod, issueExp, service)
	if err != nil {
		slog.Error("Failed to issue token", "error", err, "did", did)
		http.Error(w, fmt.Sprintf("failed to issue token: %v", err), http.StatusInternalServerError)
		return
	}

	slog.Debug("Issued JWT token",
		"tokenLength", len(tokenString),
		"did", did,
		"authMethod", authMethod,
		"credentialsDur", credDur.Round(time.Millisecond),
		"gateDur", gateDur.Round(time.Millisecond),
		"fetchDur", fetchDur.Round(time.Millisecond),
		"totalDur", time.Since(phaseStart).Round(time.Millisecond))

	// Return token response
	now := time.Now()
	expiresIn := int(issueExp.Seconds())

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

	render.JSON(w, r, resp)
}

// parseBasicAuthDID fixes DID usernames that are mangled by HTTP Basic Auth.
//
// This handles two cases:
//  1. Hyphen-encoded DIDs (did-plc-abc123) — converted to did:plc:abc123.
//     This is the recommended format for tools like helm that reject colons in usernames.
//  2. Raw DIDs split by BasicAuth — "did:plc:abc123" gets split on the first colon
//     into username="did", password="plc:abc123:<real-password>". Reconstructed here.
func parseBasicAuthDID(username, password string) (string, string) {
	// Case 1: Hyphen-encoded DID (e.g., did-plc-abc123 or did-web-example.com)
	if did, ok := auth.DecodeDIDFromHyphens(username); ok {
		return did, password
	}

	// Case 2: Raw DID was split by BasicAuth on the first colon
	// username="did", password="plc:<id>:<real-password>" or "web:<host>:<real-password>"
	if username != "did" {
		return username, password
	}

	if after, ok := strings.CutPrefix(password, "plc:"); ok {
		rest := after
		if idx := strings.Index(rest, ":"); idx > 0 {
			return "did:plc:" + rest[:idx], rest[idx+1:]
		}
	} else if after, ok := strings.CutPrefix(password, "web:"); ok {
		rest := after
		if idx := strings.Index(rest, ":"); idx > 0 {
			return "did:web:" + rest[:idx], rest[idx+1:]
		}
	}

	return username, password
}