at-container-registry/pkg/appview/middleware/registry.go

package middleware

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"strings"
	"sync"

	"github.com/distribution/distribution/v3"
	"github.com/distribution/distribution/v3/registry/api/errcode"
	registrymw "github.com/distribution/distribution/v3/registry/middleware/registry"
	"github.com/distribution/distribution/v3/registry/storage/driver"
	"github.com/distribution/reference"

	"atcr.io/pkg/appview/storage"
	"atcr.io/pkg/atproto"
	"atcr.io/pkg/auth"
	"atcr.io/pkg/auth/oauth"
	"atcr.io/pkg/auth/token"
)

// holdDIDKey is the context key for storing hold DID
const holdDIDKey contextKey = "hold.did"

// Global variables for initialization only
// These are set by main.go during startup and copied into NamespaceResolver instances.
// After initialization, request handling uses the NamespaceResolver's instance fields.
var (
	globalRefresher   *oauth.Refresher
	globalDatabase    storage.DatabaseMetrics
	globalAuthorizer  auth.HoldAuthorizer
	globalReadmeCache storage.ReadmeCache
)

// SetGlobalRefresher sets the OAuth refresher instance during initialization
// Must be called before the registry starts serving requests
func SetGlobalRefresher(refresher *oauth.Refresher) {
	globalRefresher = refresher
}

// SetGlobalDatabase sets the database instance during initialization
// Must be called before the registry starts serving requests
func SetGlobalDatabase(database storage.DatabaseMetrics) {
	globalDatabase = database
}

// SetGlobalAuthorizer sets the authorizer instance during initialization
// Must be called before the registry starts serving requests
func SetGlobalAuthorizer(authorizer auth.HoldAuthorizer) {
	globalAuthorizer = authorizer
}

// SetGlobalReadmeCache sets the readme cache instance during initialization
// Must be called before the registry starts serving requests
func SetGlobalReadmeCache(readmeCache storage.ReadmeCache) {
	globalReadmeCache = readmeCache
}

func init() {
	// Register the name resolution middleware
	registrymw.Register("atproto-resolver", initATProtoResolver)
}

// NamespaceResolver wraps a namespace and resolves names
type NamespaceResolver struct {
	distribution.Namespace
	defaultHoldDID string                  // Default hold DID (e.g., "did:web:hold01.atcr.io")
	baseURL        string                  // Base URL for error messages (e.g., "https://atcr.io")
	testMode       bool                    // If true, fallback to default hold when user's hold is unreachable
	repositories   sync.Map                // Cache of RoutingRepository instances by key (did:reponame)
	refresher      *oauth.Refresher        // OAuth session manager (copied from global on init)
	database       storage.DatabaseMetrics // Metrics database (copied from global on init)
	authorizer     auth.HoldAuthorizer     // Hold authorization (copied from global on init)
	readmeCache    storage.ReadmeCache     // README cache (copied from global on init)
}

// initATProtoResolver initializes the name resolution middleware
func initATProtoResolver(ctx context.Context, ns distribution.Namespace, _ driver.StorageDriver, options map[string]any) (distribution.Namespace, error) {
	// Get default hold DID from config (required)
	// Expected format: "did:web:hold01.atcr.io"
	defaultHoldDID := ""
	if holdDID, ok := options["default_hold_did"].(string); ok {
		defaultHoldDID = holdDID
	}

	// Get base URL from config (for error messages)
	baseURL := ""
	if url, ok := options["base_url"].(string); ok {
		baseURL = url
	}

	// Check test mode from options (passed via env var)
	testMode := false
	if tm, ok := options["test_mode"].(bool); ok {
		testMode = tm
	}

	// Copy shared services from globals into the instance
	// This avoids accessing globals during request handling
	return &NamespaceResolver{
		Namespace:      ns,
		defaultHoldDID: defaultHoldDID,
		baseURL:        baseURL,
		testMode:       testMode,
		refresher:      globalRefresher,
		database:       globalDatabase,
		authorizer:     globalAuthorizer,
		readmeCache:    globalReadmeCache,
	}, nil
}

// authErrorMessage creates a user-friendly auth error with login URL
func (nr *NamespaceResolver) authErrorMessage(message string) error {
	loginURL := fmt.Sprintf("%s/auth/oauth/login", nr.baseURL)
	fullMessage := fmt.Sprintf("%s - please re-authenticate at %s", message, loginURL)
	return errcode.ErrorCodeUnauthorized.WithMessage(fullMessage)
}

// Repository resolves the repository name and delegates to underlying namespace
// Handles names like:
// - atcr.io/alice/myimage → resolve alice to DID
// - atcr.io/did:plc:xyz123/myimage → use DID directly
func (nr *NamespaceResolver) Repository(ctx context.Context, name reference.Named) (distribution.Repository, error) {
	// Extract the first part of the name (username or DID)
	repoPath := name.Name()
	parts := strings.SplitN(repoPath, "/", 2)

	if len(parts) < 2 {
		// No user specified, use default or return error
		return nil, fmt.Errorf("repository name must include user: %s", repoPath)
	}

	identityStr := parts[0]
	imageName := parts[1]

	// Resolve identity to DID, handle, and PDS endpoint
	did, handle, pdsEndpoint, err := atproto.ResolveIdentity(ctx, identityStr)
	if err != nil {
		return nil, err
	}

	slog.Debug("Resolved identity", "component", "registry/middleware", "did", did, "pds", pdsEndpoint, "handle", handle)

	// Query for hold DID - either user's hold or default hold service
	holdDID := nr.findHoldDID(ctx, did, pdsEndpoint)
	if holdDID == "" {
		// This is a fatal configuration error - registry cannot function without a hold service
		return nil, fmt.Errorf("no hold DID configured: ensure default_hold_did is set in middleware config")
	}
	ctx = context.WithValue(ctx, holdDIDKey, holdDID)

	// Auto-reconcile crew membership on first push/pull
	// This ensures users can push immediately after docker login without web sign-in
	// EnsureCrewMembership is best-effort and logs errors without failing the request
	if holdDID != "" && nr.refresher != nil {
		slog.Debug("Auto-reconciling crew membership", "component", "registry/middleware", "did", did, "hold_did", holdDID)
		client := atproto.NewClient(pdsEndpoint, did, "")
		storage.EnsureCrewMembership(ctx, client, nr.refresher, holdDID)
	}

	// Get service token for hold authentication
	var serviceToken string
	if nr.refresher != nil {
		var err error
		serviceToken, err = token.GetOrFetchServiceToken(ctx, nr.refresher, did, holdDID, pdsEndpoint)
		if err != nil {
			slog.Error("Failed to get service token", "component", "registry/middleware", "did", did, "error", err)
			slog.Error("User needs to re-authenticate via credential helper", "component", "registry/middleware")
			return nil, nr.authErrorMessage("OAuth session expired")
		}
	}

	// Create a new reference with identity/image format
	// Use the identity (or DID) as the namespace to ensure canonical format
	// This transforms: evan.jarrett.net/debian -> evan.jarrett.net/debian (keeps full path)
	canonicalName := fmt.Sprintf("%s/%s", identityStr, imageName)
	ref, err := reference.ParseNamed(canonicalName)
	if err != nil {
		return nil, fmt.Errorf("invalid image name %s: %w", imageName, err)
	}

	// Delegate to underlying namespace with modified name
	repo, err := nr.Namespace.Repository(ctx, ref)
	if err != nil {
		return nil, err
	}

	// Get access token for PDS operations
	// Try OAuth refresher first (for users who authorized via AppView OAuth)
	// Fall back to Basic Auth token cache (for users who used app passwords)
	var atprotoClient *atproto.Client

	if nr.refresher != nil {
		// Try OAuth flow first
		session, err := nr.refresher.GetSession(ctx, did)
		if err == nil {
			// OAuth session available - use indigo's API client (handles DPoP automatically)
			apiClient := session.APIClient()
			atprotoClient = atproto.NewClientWithIndigoClient(pdsEndpoint, did, apiClient)
		} else {
			slog.Debug("OAuth refresh failed, falling back to Basic Auth", "component", "registry/middleware", "did", did, "error", err)
		}
	}

	// Fall back to Basic Auth token cache if OAuth not available
	if atprotoClient == nil {
		accessToken, ok := auth.GetGlobalTokenCache().Get(did)
		if !ok {
			slog.Debug("No cached access token found (neither OAuth nor Basic Auth)", "component", "registry/middleware", "did", did)
			accessToken = "" // Will fail on manifest push, but let it try
		} else {
			slog.Debug("Using Basic Auth access token", "component", "registry/middleware", "did", did, "token_length", len(accessToken))
		}
		atprotoClient = atproto.NewClient(pdsEndpoint, did, accessToken)
	}

	// IMPORTANT: Use only the image name (not identity/image) for ATProto storage
	// ATProto records are scoped to the user's DID, so we don't need the identity prefix
	// Example: "evan.jarrett.net/debian" -> store as "debian"
	repositoryName := imageName

	// Cache key is DID + repository name
	cacheKey := did + ":" + repositoryName

	// Check cache first and update service token
	if cached, ok := nr.repositories.Load(cacheKey); ok {
		cachedRepo := cached.(*storage.RoutingRepository)
		// Always update the service token even for cached repos (token may have been renewed)
		cachedRepo.Ctx.ServiceToken = serviceToken
		return cachedRepo, nil
	}

	// Create routing repository - routes manifests to ATProto, blobs to hold service
	// The registry is stateless - no local storage is used
	// Bundle all context into a single RegistryContext struct
	registryCtx := &storage.RegistryContext{
		DID:           did,
		Handle:        handle,
		HoldDID:       holdDID,
		PDSEndpoint:   pdsEndpoint,
		Repository:    repositoryName,
		ServiceToken:  serviceToken, // Cached service token from middleware validation
		ATProtoClient: atprotoClient,
		Database:      nr.database,
		Authorizer:    nr.authorizer,
		Refresher:     nr.refresher,
		ReadmeCache:   nr.readmeCache,
	}
	routingRepo := storage.NewRoutingRepository(repo, registryCtx)

	// Cache the repository
	nr.repositories.Store(cacheKey, routingRepo)

	return routingRepo, nil
}

// Repositories delegates to underlying namespace
func (nr *NamespaceResolver) Repositories(ctx context.Context, repos []string, last string) (int, error) {
	return nr.Namespace.Repositories(ctx, repos, last)
}

// Blobs delegates to underlying namespace
func (nr *NamespaceResolver) Blobs() distribution.BlobEnumerator {
	return nr.Namespace.Blobs()
}

// BlobStatter delegates to underlying namespace
func (nr *NamespaceResolver) BlobStatter() distribution.BlobStatter {
	return nr.Namespace.BlobStatter()
}

// findHoldDID determines which hold DID to use for blob storage
// Priority order:
// 1. User's sailor profile defaultHold (if set)
// 2. User's own hold record (io.atcr.hold)
// 3. AppView's default hold DID
// Returns a hold DID (e.g., "did:web:hold01.atcr.io"), or empty string if none configured
func (nr *NamespaceResolver) findHoldDID(ctx context.Context, did, pdsEndpoint string) string {
	// Create ATProto client (without auth - reading public records)
	client := atproto.NewClient(pdsEndpoint, did, "")

	// Check for sailor profile
	profile, err := storage.GetProfile(ctx, client)
	if err != nil {
		// Error reading profile (not a 404) - log and continue
		slog.Warn("Failed to read profile", "did", did, "error", err)
	}

	if profile != nil && profile.DefaultHold != "" {
		// Profile exists with defaultHold set
		// In test mode, verify it's reachable before using it
		if nr.testMode {
			if nr.isHoldReachable(ctx, profile.DefaultHold) {
				return profile.DefaultHold
			}
			slog.Debug("User's defaultHold unreachable, falling back to default", "component", "registry/middleware/testmode", "default_hold", profile.DefaultHold)
			return nr.defaultHoldDID
		}
		return profile.DefaultHold
	}

	// Profile doesn't exist or defaultHold is null/empty
	// Check for user's own hold records
	records, err := client.ListRecords(ctx, atproto.HoldCollection, 10)
	if err != nil {
		// Failed to query holds, use default
		return nr.defaultHoldDID
	}

	// Find the first hold record
	for _, record := range records {
		var holdRecord atproto.HoldRecord
		if err := json.Unmarshal(record.Value, &holdRecord); err != nil {
			continue
		}

		// Return the endpoint from the first hold (normalize to DID if URL)
		if holdRecord.Endpoint != "" {
			return atproto.ResolveHoldDIDFromURL(holdRecord.Endpoint)
		}
	}

	// No profile defaultHold and no own hold records - use AppView default
	return nr.defaultHoldDID
}

// isHoldReachable checks if a hold service is reachable
// Used in test mode to fallback to default hold when user's hold is unavailable
func (nr *NamespaceResolver) isHoldReachable(ctx context.Context, holdDID string) bool {
	// Try to fetch the DID document
	hostname := strings.TrimPrefix(holdDID, "did:web:")

	// Try HTTP first (local), then HTTPS
	for _, scheme := range []string{"http", "https"} {
		testURL := fmt.Sprintf("%s://%s/.well-known/did.json", scheme, hostname)
		client := atproto.NewClient("", "", "")
		_, err := client.FetchDIDDocument(ctx, testURL)
		if err == nil {
			return true
		}
	}

	return false
}