at-container-registry/pkg/hold/pds/auth.go

package pds

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

	"atcr.io/pkg/atproto"
	"github.com/bluesky-social/indigo/atproto/atcrypto"
	"github.com/bluesky-social/indigo/atproto/syntax"
	"github.com/golang-jwt/jwt/v5"
)

// Authentication errors
var (
	ErrMissingAuthHeader = errors.New("missing Authorization header")
	ErrInvalidAuthFormat = errors.New("invalid Authorization header format")
	ErrInvalidAuthScheme = errors.New("invalid authorization scheme: expected 'Bearer' or 'DPoP'")
	ErrMissingToken      = errors.New("missing token")
	ErrMissingDPoPHeader = errors.New("missing DPoP header")
)

// JWT validation errors
var (
	ErrInvalidJWTFormat = errors.New("invalid JWT format: expected header.payload.signature")
	ErrMissingISSClaim  = errors.New("missing 'iss' claim in token")
	ErrMissingSubClaim  = errors.New("missing 'sub' claim in token")
	ErrTokenExpired     = errors.New("token has expired")
)

// AuthError provides structured authorization error information
type AuthError struct {
	Action   string   // The action being attempted: "blob:read", "blob:write", "crew:admin"
	Reason   string   // Why access was denied
	Required []string // What permission(s) would grant access
}

func (e *AuthError) Error() string {
	return fmt.Sprintf("access denied for %s: %s (required: %s)",
		e.Action, e.Reason, strings.Join(e.Required, " or "))
}

// NewAuthError creates a new AuthError
func NewAuthError(action, reason string, required ...string) *AuthError {
	return &AuthError{
		Action:   action,
		Reason:   reason,
		Required: required,
	}
}

// HTTPClient interface allows injecting a custom HTTP client for testing
type HTTPClient interface {
	Do(*http.Request) (*http.Response, error)
}

// ValidatedUser represents a successfully validated user from DPoP + OAuth
type ValidatedUser struct {
	DID        string
	Handle     string
	PDS        string
	Authorized bool
}

// ValidateDPoPRequest validates a request with DPoP + OAuth tokens
// This implements the standard ATProto token validation flow:
// 1. Extract Authorization header (DPoP <token>)
// 2. Extract DPoP header (proof JWT)
// 3. Call user's PDS to validate token via com.atproto.server.getSession
// 4. Return validated user DID
//
// The httpClient parameter is optional and defaults to http.DefaultClient if nil.
// This allows tests to inject a mock HTTP client.
func ValidateDPoPRequest(r *http.Request, httpClient HTTPClient) (*ValidatedUser, error) {
	// Extract Authorization header
	authHeader := r.Header.Get("Authorization")
	if authHeader == "" {
		return nil, ErrMissingAuthHeader
	}

	// Check for DPoP authorization scheme
	parts := strings.SplitN(authHeader, " ", 2)
	if len(parts) != 2 {
		return nil, ErrInvalidAuthFormat
	}

	if parts[0] != "DPoP" {
		return nil, fmt.Errorf("expected DPoP authorization scheme, got: %s", parts[0])
	}

	accessToken := parts[1]
	if accessToken == "" {
		return nil, ErrMissingToken
	}

	// Extract DPoP header
	dpopProof := r.Header.Get("DPoP")
	if dpopProof == "" {
		return nil, ErrMissingDPoPHeader
	}

	// TODO: We could verify the DPoP proof locally (signature, HTM, HTU, etc.)
	// For now, we'll rely on the PDS to validate everything

	// The token contains the user's DID in its claims, but we can't trust it without validation
	// We need to call the user's PDS to validate the token
	// Problem: We don't know which PDS to call yet!

	// For now, we'll parse the JWT to extract the DID/PDS hint (unverified)
	// Then validate against that PDS
	// This is safe because the PDS will verify the token is valid for that DID

	did, pds, err := extractDIDFromToken(accessToken)
	if err != nil {
		return nil, fmt.Errorf("failed to extract DID from token: %w", err)
	}

	// Validate token with the user's PDS
	session, err := validateTokenWithPDS(r.Context(), pds, accessToken, dpopProof, httpClient)
	if err != nil {
		return nil, fmt.Errorf("token validation failed: %w", err)
	}

	// Verify the DID matches
	if session.DID != did {
		return nil, fmt.Errorf("token DID mismatch: expected %s, got %s", did, session.DID)
	}

	return &ValidatedUser{
		DID:        session.DID,
		Handle:     session.Handle,
		PDS:        pds,
		Authorized: true,
	}, nil
}

// extractDIDFromToken extracts the DID and PDS from an unverified JWT token
// This is just for routing purposes - the token will be validated by the PDS
func extractDIDFromToken(token string) (string, string, error) {
	// JWT format: header.payload.signature
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return "", "", ErrInvalidJWTFormat
	}

	// Decode payload (base64url)
	payload, err := decodeBase64URL(parts[1])
	if err != nil {
		return "", "", fmt.Errorf("failed to decode payload: %w", err)
	}

	// Parse JSON
	var claims struct {
		Sub string `json:"sub"` // DID
		Iss string `json:"iss"` // PDS URL (issuer)
	}

	if err := json.Unmarshal(payload, &claims); err != nil {
		return "", "", fmt.Errorf("failed to parse claims: %w", err)
	}

	if claims.Sub == "" {
		return "", "", ErrMissingSubClaim
	}

	if claims.Iss == "" {
		return "", "", ErrMissingISSClaim
	}

	return claims.Sub, claims.Iss, nil
}

// decodeBase64URL decodes base64url (RFC 4648)
func decodeBase64URL(s string) ([]byte, error) {
	// Use Go's RawURLEncoding (base64url without padding)
	return base64.RawURLEncoding.DecodeString(s)
}

// SessionResponse represents the response from com.atproto.server.getSession
type SessionResponse struct {
	DID    string `json:"did"`
	Handle string `json:"handle"`
}

// validateTokenWithPDS calls the user's PDS to validate the token
// The httpClient parameter is optional and defaults to http.DefaultClient if nil.
func validateTokenWithPDS(ctx context.Context, pdsURL, accessToken, dpopProof string, httpClient HTTPClient) (*SessionResponse, error) {
	// Call com.atproto.server.getSession with DPoP headers
	url := fmt.Sprintf("%s%s", strings.TrimSuffix(pdsURL, "/"), atproto.ServerGetSession)

	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	// Add DPoP authorization headers
	req.Header.Set("Authorization", "DPoP "+accessToken)
	req.Header.Set("DPoP", dpopProof)

	// Use provided client or default to http.DefaultClient
	client := httpClient
	if client == nil {
		client = http.DefaultClient
	}

	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed to call PDS: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("PDS returned status %d: %s", resp.StatusCode, string(body))
	}

	var session SessionResponse
	if err := json.NewDecoder(resp.Body).Decode(&session); err != nil {
		return nil, fmt.Errorf("failed to decode session: %w", err)
	}

	return &session, nil
}

// ValidateOwnerOrCrewAdmin validates that the request has valid authentication
// and that the authenticated user is either the hold owner or a crew member with crew:admin permission.
// Supports two authentication methods:
// 1. Service tokens (Bearer tokens from com.atproto.server.getServiceAuth) - for AppView access
// 2. DPoP + OAuth tokens - for direct user access
// The httpClient parameter is optional and defaults to http.DefaultClient if nil.
func ValidateOwnerOrCrewAdmin(r *http.Request, pds *HoldPDS, httpClient HTTPClient) (*ValidatedUser, error) {
	// Try service token validation first (for AppView access)
	authHeader := r.Header.Get("Authorization")
	var user *ValidatedUser
	var err error

	if strings.HasPrefix(authHeader, "Bearer ") {
		// Service token authentication
		user, err = ValidateServiceToken(r, pds.did, httpClient)
		if err != nil {
			return nil, fmt.Errorf("service token authentication failed: %w", err)
		}
	} else if strings.HasPrefix(authHeader, "DPoP ") {
		// DPoP + OAuth authentication (direct user access)
		user, err = ValidateDPoPRequest(r, httpClient)
		if err != nil {
			return nil, fmt.Errorf("DPoP authentication failed: %w", err)
		}
	} else {
		return nil, ErrInvalidAuthScheme
	}

	// Get captain record to check owner
	_, captain, err := pds.GetCaptainRecord(r.Context())
	if err != nil {
		return nil, fmt.Errorf("failed to get captain record: %w", err)
	}

	// Check if user is the owner
	if user.DID == captain.Owner {
		return user, nil
	}

	// Check if user is crew with admin permission
	crew, err := pds.ListCrewMembers(r.Context())
	if err != nil {
		return nil, fmt.Errorf("failed to check crew membership: %w", err)
	}

	for _, member := range crew {
		if member.Record.Member == user.DID {
			// Check if this crew member has crew:admin permission
			if slices.Contains(member.Record.Permissions, "crew:admin") {
				return user, nil
			}
			// User is crew but doesn't have admin permission
			return nil, NewAuthError("crew:admin", "crew member lacks permission", "crew:admin")
		}
	}

	// User is neither owner nor authorized crew
	return nil, NewAuthError("crew:admin", "user is not a crew member", "crew:admin")
}

// ValidateBlobWriteAccess validates that the request has valid authentication
// and that the authenticated user is either the hold owner or a crew member with blob:write permission.
// Supports two authentication methods:
// 1. Service tokens (Bearer tokens from com.atproto.server.getServiceAuth) - for AppView access
// 2. DPoP + OAuth tokens - for direct user access
// The httpClient parameter is optional and defaults to http.DefaultClient if nil.
func ValidateBlobWriteAccess(r *http.Request, pds *HoldPDS, httpClient HTTPClient) (*ValidatedUser, error) {
	// Try service token validation first (for AppView access)
	authHeader := r.Header.Get("Authorization")
	var user *ValidatedUser
	var err error

	if strings.HasPrefix(authHeader, "Bearer ") {
		// Service token authentication
		user, err = ValidateServiceToken(r, pds.did, httpClient)
		if err != nil {
			return nil, fmt.Errorf("service token authentication failed: %w", err)
		}
	} else if strings.HasPrefix(authHeader, "DPoP ") {
		// DPoP + OAuth authentication (direct user access)
		user, err = ValidateDPoPRequest(r, httpClient)
		if err != nil {
			return nil, fmt.Errorf("DPoP authentication failed: %w", err)
		}
	} else {
		return nil, ErrInvalidAuthScheme
	}

	// Get captain record to check owner and public settings
	_, captain, err := pds.GetCaptainRecord(r.Context())
	if err != nil {
		return nil, fmt.Errorf("failed to get captain record: %w", err)
	}

	// Check if user is the owner (always has write access)
	if user.DID == captain.Owner {
		return user, nil
	}

	// Check if user is crew with blob:write permission
	crew, err := pds.ListCrewMembers(r.Context())
	if err != nil {
		return nil, fmt.Errorf("failed to check crew membership: %w", err)
	}

	for _, member := range crew {
		if member.Record.Member == user.DID {
			// Check if this crew member has blob:write permission
			if slices.Contains(member.Record.Permissions, "blob:write") {
				return user, nil
			}
			// User is crew but doesn't have write permission
			return nil, NewAuthError("blob:write", "crew member lacks permission", "blob:write")
		}
	}

	// User is neither owner nor authorized crew
	return nil, NewAuthError("blob:write", "user is not a crew member", "blob:write")
}

// ValidateBlobReadAccess validates that the request has read access to blobs
// If captain.public = true: No auth required (returns nil user to indicate public access)
// If captain.public = false: Requires valid DPoP + OAuth and (captain OR crew with blob:read or blob:write permission).
// Note: blob:write implicitly grants blob:read access.
// The httpClient parameter is optional and defaults to http.DefaultClient if nil.
// If scannerSecret is non-empty, a Bearer token matching it grants full read access (for scanner blob fetches).
func ValidateBlobReadAccess(r *http.Request, pds *HoldPDS, httpClient HTTPClient, scannerSecret string) (*ValidatedUser, error) {
	// Get captain record to check public setting
	_, captain, err := pds.GetCaptainRecord(r.Context())
	if err != nil {
		return nil, fmt.Errorf("failed to get captain record: %w", err)
	}

	// If hold is public, allow access without authentication
	if captain.Public {
		return nil, nil // nil user indicates public access
	}

	// Private hold - require authentication (accept both service tokens and DPoP)
	authHeader := r.Header.Get("Authorization")
	var user *ValidatedUser

	if strings.HasPrefix(authHeader, "Bearer ") {
		// Check if this is a scanner using the shared secret
		if scannerSecret != "" && strings.TrimPrefix(authHeader, "Bearer ") == scannerSecret {
			return &ValidatedUser{DID: "scanner"}, nil
		}
		// Service token authentication (from AppView via getServiceAuth)
		user, err = ValidateServiceToken(r, pds.did, httpClient)
		if err != nil {
			return nil, fmt.Errorf("service token authentication failed: %w", err)
		}
	} else if strings.HasPrefix(authHeader, "DPoP ") {
		// DPoP + OAuth authentication (direct user access)
		user, err = ValidateDPoPRequest(r, httpClient)
		if err != nil {
			return nil, fmt.Errorf("DPoP authentication failed: %w", err)
		}
	} else {
		return nil, ErrInvalidAuthScheme
	}

	// Check if user is the owner (always has read access)
	if user.DID == captain.Owner {
		return user, nil
	}

	// Check if user is crew with blob:read or blob:write permission
	// Note: blob:write implicitly grants blob:read access
	crew, err := pds.ListCrewMembers(r.Context())
	if err != nil {
		return nil, fmt.Errorf("failed to check crew membership: %w", err)
	}

	for _, member := range crew {
		if member.Record.Member == user.DID {
			// Check if this crew member has blob:read or blob:write permission
			// blob:write implicitly grants read access (can't push without pulling)
			if slices.Contains(member.Record.Permissions, "blob:read") ||
				slices.Contains(member.Record.Permissions, "blob:write") {
				return user, nil
			}
			// User is crew but doesn't have read or write permission
			return nil, NewAuthError("blob:read", "crew member lacks permission", "blob:read", "blob:write")
		}
	}

	// User is neither owner nor authorized crew
	return nil, NewAuthError("blob:read", "user is not a crew member", "blob:read", "blob:write")
}

// ServiceTokenClaims represents the claims in a service token JWT
type ServiceTokenClaims struct {
	jwt.RegisteredClaims
}

// ValidateServiceToken validates a service token JWT from com.atproto.server.getServiceAuth
// This validates the JWT signature using the issuer's (PDS) public key from their DID document
// Returns the user DID from the iss claim if validation succeeds
func ValidateServiceToken(r *http.Request, holdDID string, httpClient HTTPClient) (*ValidatedUser, error) {
	// Extract Authorization header
	authHeader := r.Header.Get("Authorization")
	if authHeader == "" {
		return nil, ErrMissingAuthHeader
	}

	// Check for Bearer authorization scheme
	parts := strings.SplitN(authHeader, " ", 2)
	if len(parts) != 2 {
		return nil, ErrInvalidAuthFormat
	}

	if parts[0] != "Bearer" {
		return nil, fmt.Errorf("expected Bearer authorization scheme, got: %s", parts[0])
	}

	tokenString := parts[1]
	if tokenString == "" {
		return nil, ErrMissingToken
	}

	slog.Debug("Validating service token", "holdDID", holdDID)

	// Manually parse JWT (bypass golang-jwt since it doesn't support ES256K algorithm used by ATProto)
	// Split token: header.payload.signature
	tokenParts := strings.Split(tokenString, ".")
	if len(tokenParts) != 3 {
		return nil, ErrInvalidJWTFormat
	}

	// Decode payload (second part) to extract claims
	payloadBytes, err := base64.RawURLEncoding.DecodeString(tokenParts[1])
	if err != nil {
		return nil, fmt.Errorf("failed to decode JWT payload: %w", err)
	}

	// Parse claims from JSON
	var claims ServiceTokenClaims
	if err := json.Unmarshal(payloadBytes, &claims); err != nil {
		return nil, fmt.Errorf("failed to unmarshal claims: %w", err)
	}

	// Get issuer (user DID)
	issuerDID := claims.Issuer
	if issuerDID == "" {
		return nil, ErrMissingISSClaim
	}

	// Verify audience matches this hold service
	audiences, err := claims.GetAudience()
	if err != nil {
		return nil, fmt.Errorf("failed to get audience: %w", err)
	}
	if len(audiences) == 0 || audiences[0] != holdDID {
		return nil, fmt.Errorf("token audience mismatch: expected %s, got %v", holdDID, audiences)
	}

	// Verify expiration
	exp, err := claims.GetExpirationTime()
	if err != nil {
		return nil, fmt.Errorf("failed to get expiration: %w", err)
	}
	if exp != nil && time.Now().After(exp.Time) {
		return nil, ErrTokenExpired
	}

	// Verify JWT signature using ATProto's secp256k1 crypto

	// Signature is over "header.payload"
	signedData := []byte(tokenParts[0] + "." + tokenParts[1])

	// Decode signature (base64url)
	signature, err := base64.RawURLEncoding.DecodeString(tokenParts[2])
	if err != nil {
		return nil, fmt.Errorf("failed to decode signature: %w", err)
	}

	// Fetch public key from issuer's DID document
	publicKey, err := fetchPublicKeyFromDID(r.Context(), issuerDID)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch public key for issuer %s: %w", issuerDID, err)
	}

	// Verify signature using indigo's crypto (handles secp256k1)
	if err := publicKey.HashAndVerify(signedData, signature); err != nil {
		return nil, fmt.Errorf("signature verification failed: %w", err)
	}

	slog.Debug("Successfully validated service token", "userDID", issuerDID)

	// Return validated user
	return &ValidatedUser{
		DID:        issuerDID,
		Handle:     "", // Not available in service token
		PDS:        "", // Not needed for authorization
		Authorized: true,
	}, nil
}

// ValidateAppviewToken validates a JWT signed by the trusted appview using ES256 (P-256).
// It resolves the appview's DID document to extract the P-256 public key, then verifies
// the JWT signature, issuer (iss), and audience (aud).
//
// Returns the subject (sub) claim which is the user DID being acted upon.
func ValidateAppviewToken(r *http.Request, appviewDID, holdDID string) (string, error) {
	// Extract Authorization header
	authHeader := r.Header.Get("Authorization")
	if authHeader == "" {
		return "", ErrMissingAuthHeader
	}

	parts := strings.SplitN(authHeader, " ", 2)
	if len(parts) != 2 || parts[0] != "Bearer" {
		return "", fmt.Errorf("expected Bearer authorization scheme")
	}

	tokenString := parts[1]
	if tokenString == "" {
		return "", ErrMissingToken
	}

	// Manually parse JWT
	tokenParts := strings.Split(tokenString, ".")
	if len(tokenParts) != 3 {
		return "", ErrInvalidJWTFormat
	}

	// Decode and parse claims
	payloadBytes, err := base64.RawURLEncoding.DecodeString(tokenParts[1])
	if err != nil {
		return "", fmt.Errorf("failed to decode JWT payload: %w", err)
	}

	var claims ServiceTokenClaims
	if err := json.Unmarshal(payloadBytes, &claims); err != nil {
		return "", fmt.Errorf("failed to unmarshal claims: %w", err)
	}

	// Verify issuer matches configured appview DID
	if claims.Issuer != appviewDID {
		return "", fmt.Errorf("token issuer mismatch: expected %s, got %s", appviewDID, claims.Issuer)
	}

	// Verify audience matches this hold's DID
	audiences, err := claims.GetAudience()
	if err != nil {
		return "", fmt.Errorf("failed to get audience: %w", err)
	}
	if len(audiences) == 0 || audiences[0] != holdDID {
		return "", fmt.Errorf("token audience mismatch: expected %s, got %v", holdDID, audiences)
	}

	// Verify expiration
	exp, err := claims.GetExpirationTime()
	if err != nil {
		return "", fmt.Errorf("failed to get expiration: %w", err)
	}
	if exp != nil && time.Now().After(exp.Time) {
		return "", ErrTokenExpired
	}

	// Get subject (user DID)
	subject, err := claims.GetSubject()
	if err != nil || subject == "" {
		return "", ErrMissingSubClaim
	}

	// Fetch P-256 public key from appview DID document
	pubKey, err := fetchP256PublicKeyFromDID(r.Context(), appviewDID)
	if err != nil {
		return "", fmt.Errorf("failed to fetch appview public key: %w", err)
	}

	// Verify JWT signature with P-256 key
	signedData := []byte(tokenParts[0] + "." + tokenParts[1])
	signature, err := base64.RawURLEncoding.DecodeString(tokenParts[2])
	if err != nil {
		return "", fmt.Errorf("failed to decode signature: %w", err)
	}

	if err := pubKey.HashAndVerifyLenient(signedData, signature); err != nil {
		return "", fmt.Errorf("signature verification failed: %w", err)
	}

	slog.Debug("Validated appview service token", "appviewDID", appviewDID, "userDID", subject)
	return subject, nil
}

// fetchP256PublicKeyFromDID fetches a P-256 public key from a did:web DID document.
// It resolves the DID document and looks for a Multikey verification method with P-256 prefix.
func fetchP256PublicKeyFromDID(ctx context.Context, did string) (*atcrypto.PublicKeyP256, error) {
	if !strings.HasPrefix(did, "did:web:") {
		return nil, fmt.Errorf("only did:web is supported for appview DID, got %s", did)
	}

	// Resolve did:web to URL
	host := strings.TrimPrefix(did, "did:web:")
	host = strings.ReplaceAll(host, "%3A", ":")
	scheme := "https"
	if atproto.IsTestMode() {
		scheme = "http"
	}
	didDocURL := fmt.Sprintf("%s://%s/.well-known/did.json", scheme, host)

	req, err := http.NewRequestWithContext(ctx, "GET", didDocURL, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch DID document: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("DID document fetch returned status %d", resp.StatusCode)
	}

	var doc struct {
		VerificationMethod []struct {
			ID                 string `json:"id"`
			Type               string `json:"type"`
			PublicKeyMultibase string `json:"publicKeyMultibase"`
		} `json:"verificationMethod"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&doc); err != nil {
		return nil, fmt.Errorf("failed to decode DID document: %w", err)
	}

	// Find a Multikey verification method with P-256 public key
	for _, vm := range doc.VerificationMethod {
		if vm.Type != "Multikey" || vm.PublicKeyMultibase == "" {
			continue
		}

		// Try parsing as P-256 key via atcrypto's multibase parser
		pubKey, err := atcrypto.ParsePublicMultibase(vm.PublicKeyMultibase)
		if err != nil {
			continue
		}

		p256Key, ok := pubKey.(*atcrypto.PublicKeyP256)
		if ok {
			return p256Key, nil
		}
	}

	return nil, fmt.Errorf("no P-256 public key found in DID document for %s", did)
}

// fetchPublicKeyFromDID fetches the public key from a DID document
// Supports did:plc and did:web
// Returns the atcrypto.PublicKey for signature verification
func fetchPublicKeyFromDID(ctx context.Context, did string) (atcrypto.PublicKey, error) {
	// Use indigo's identity resolution
	directory := atproto.GetDirectory()
	atID, err := syntax.ParseAtIdentifier(did)
	if err != nil {
		return nil, fmt.Errorf("invalid DID format: %w", err)
	}

	ident, err := directory.Lookup(ctx, atID)
	if err != nil {
		return nil, fmt.Errorf("failed to resolve DID: %w", err)
	}

	// Get the public key using indigo's built-in method
	// This returns an atcrypto.PublicKey (secp256k1)
	publicKey, err := ident.PublicKey()
	if err != nil {
		return nil, fmt.Errorf("failed to get public key from DID: %w", err)
	}

	return publicKey, nil
}