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

package middleware

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

	"github.com/bluesky-social/indigo/atproto/identity"
	"github.com/bluesky-social/indigo/atproto/syntax"
	"github.com/distribution/distribution/v3"
	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"
)

// Global refresher instance (set by main.go)
var globalRefresher *oauth.Refresher

// Global database instance (set by main.go for pull tracking)
var globalDatabase interface {
	IncrementPullCount(did, repository string) error
	IncrementPushCount(did, repository string) error
}

// SetGlobalRefresher sets the global OAuth refresher instance
func SetGlobalRefresher(refresher *oauth.Refresher) {
	globalRefresher = refresher
}

// SetGlobalDatabase sets the global database instance for metrics tracking
func SetGlobalDatabase(database interface {
	IncrementPullCount(did, repository string) error
	IncrementPushCount(did, repository string) error
}) {
	globalDatabase = database
}

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

// NamespaceResolver wraps a namespace and resolves names
type NamespaceResolver struct {
	distribution.Namespace
	directory              identity.Directory
	defaultStorageEndpoint string
	repositories           sync.Map // Cache of RoutingRepository instances by key (did:reponame)
}

// initATProtoResolver initializes the name resolution middleware
func initATProtoResolver(ctx context.Context, ns distribution.Namespace, _ driver.StorageDriver, options map[string]any) (distribution.Namespace, error) {
	// Use indigo's default directory (includes caching)
	directory := identity.DefaultDirectory()

	// Get default storage endpoint from config (optional)
	defaultStorageEndpoint := ""
	if endpoint, ok := options["default_storage_endpoint"].(string); ok {
		defaultStorageEndpoint = endpoint
	}

	return &NamespaceResolver{
		Namespace:              ns,
		directory:              directory,
		defaultStorageEndpoint: defaultStorageEndpoint,
	}, nil
}

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

	// Parse identity (handle or DID)
	atID, err := syntax.ParseAtIdentifier(identityStr)
	if err != nil {
		return nil, fmt.Errorf("invalid identity %s: %w", identityStr, err)
	}

	// Resolve identity to DID and PDS using indigo's directory
	ident, err := nr.directory.Lookup(ctx, *atID)
	if err != nil {
		return nil, fmt.Errorf("failed to resolve identity %s: %w", identityStr, err)
	}

	did := ident.DID.String()
	pdsEndpoint := ident.PDSEndpoint()
	if pdsEndpoint == "" {
		return nil, fmt.Errorf("no PDS endpoint found for %s", identityStr)
	}

	fmt.Printf("DEBUG [registry/middleware]: Resolved identity: did=%s, pds=%s, handle=%s\n", did, pdsEndpoint, ident.Handle.String())

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

	// 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 globalRefresher != nil {
		// Try OAuth flow first
		session, err := globalRefresher.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 {
			fmt.Printf("DEBUG [registry/middleware]: OAuth refresh failed for DID=%s: %v, falling back to Basic Auth\n", did, err)
		}
	}

	// Fall back to Basic Auth token cache if OAuth not available
	if atprotoClient == nil {
		accessToken, ok := auth.GetGlobalTokenCache().Get(did)
		if !ok {
			fmt.Printf("DEBUG [registry/middleware]: No cached access token found for DID=%s (neither OAuth nor Basic Auth)\n", did)
			accessToken = "" // Will fail on manifest push, but let it try
		} else {
			fmt.Printf("DEBUG [registry/middleware]: Using Basic Auth access token for DID=%s (length=%d)\n", did, 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
	if cached, ok := nr.repositories.Load(cacheKey); ok {
		return cached.(*storage.RoutingRepository), nil
	}

	// Create routing repository - routes manifests to ATProto, blobs to hold service
	// The registry is stateless - no local storage is used
	// Pass storage endpoint and DID as parameters (can't use context as it gets lost)
	routingRepo := storage.NewRoutingRepository(repo, atprotoClient, repositoryName, storageEndpoint, did, globalDatabase)

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

// findStorageEndpoint determines which hold endpoint 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 endpoint
// Returns the storage endpoint URL, or empty string if none configured
func (nr *NamespaceResolver) findStorageEndpoint(ctx context.Context, did, pdsEndpoint string) string {
	// Create ATProto client (without auth - reading public records)
	client := atproto.NewClient(pdsEndpoint, did, "")

	// 1. Check for sailor profile
	profile, err := atproto.GetProfile(ctx, client)
	if err != nil {
		// Error reading profile (not a 404) - log and continue
		fmt.Printf("WARNING: failed to read profile for %s: %v\n", did, err)
	}

	if profile != nil && profile.DefaultHold != "" {
		// Profile exists with defaultHold set - use it
		return profile.DefaultHold
	}

	// 2. 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.defaultStorageEndpoint
	}

	// 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
		if holdRecord.Endpoint != "" {
			return holdRecord.Endpoint
		}
	}

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