at-container-registry/pkg/hold/oci/xrpc.go

// Package oci provides OCI registry endpoints for the hold service.
package oci

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

	"atcr.io/pkg/atproto"
	"atcr.io/pkg/hold/pds"
	"atcr.io/pkg/hold/quota"
	"atcr.io/pkg/s3"
	"github.com/go-chi/chi/v5"
	"github.com/go-chi/render"
)

// XRPCHandler handles OCI-specific XRPC endpoints for multipart uploads
type XRPCHandler struct {
	s3Service          s3.S3Service
	MultipartMgr       *MultipartManager // Exported for access in route handlers
	pds                *pds.HoldPDS
	httpClient         pds.HTTPClient
	enableBlueskyPosts bool
	quotaMgr           *quota.Manager       // Quota manager for tier-based limits
	scanBroadcaster    *pds.ScanBroadcaster // Scan job dispatcher (nil = scanning disabled)
}

// NewXRPCHandler creates a new OCI XRPC handler
func NewXRPCHandler(holdPDS *pds.HoldPDS, s3Service s3.S3Service, enableBlueskyPosts bool, httpClient pds.HTTPClient, quotaMgr *quota.Manager) *XRPCHandler {
	return &XRPCHandler{
		MultipartMgr:       NewMultipartManager(),
		s3Service:          s3Service,
		pds:                holdPDS,
		httpClient:         httpClient,
		enableBlueskyPosts: enableBlueskyPosts,
		quotaMgr:           quotaMgr,
	}
}

// SetScanBroadcaster sets the scan broadcaster for triggering scans on push
func (h *XRPCHandler) SetScanBroadcaster(sb *pds.ScanBroadcaster) {
	h.scanBroadcaster = sb
}

// RegisterHandlers registers all OCI XRPC endpoints with the chi router
func (h *XRPCHandler) RegisterHandlers(r chi.Router) {
	// All multipart upload endpoints require blob:write permission
	r.Group(func(r chi.Router) {
		r.Use(h.requireBlobWriteAccess)

		r.Post(atproto.HoldInitiateUpload, h.HandleInitiateUpload)
		r.Post(atproto.HoldGetPartUploadURL, h.HandleGetPartUploadURL)
		r.Post(atproto.HoldCompleteUpload, h.HandleCompleteUpload)
		r.Post(atproto.HoldAbortUpload, h.HandleAbortUpload)
		r.Post(atproto.HoldNotifyManifest, h.HandleNotifyManifest)
	})
}

// HandleInitiateUpload starts a new multipart upload
// Replaces the old "action: start" pattern
func (h *XRPCHandler) HandleInitiateUpload(w http.ResponseWriter, r *http.Request) {
	var req struct {
		Digest string `json:"digest"`
	}

	if err := render.Decode(r, &req); err != nil {
		render.Status(r, http.StatusBadRequest)
		render.JSON(w, r, map[string]string{"error": err.Error()})
		return
	}

	if req.Digest == "" {
		render.Status(r, http.StatusBadRequest)
		render.JSON(w, r, map[string]string{"error": "digest is required"})
		return
	}

	uploadID, err := h.StartMultipartUploadWithManager(r.Context(), req.Digest)
	if err != nil {
		render.Status(r, http.StatusInternalServerError)
		render.JSON(w, r, map[string]string{"error": fmt.Sprintf("failed to initiate upload: %v", err)})
		return
	}

	render.JSON(w, r, map[string]any{
		"uploadId": uploadID,
	})
}

// HandleGetPartUploadURL returns a presigned URL or endpoint info for uploading a part
// Replaces the old "action: part" pattern
func (h *XRPCHandler) HandleGetPartUploadURL(w http.ResponseWriter, r *http.Request) {
	var req struct {
		UploadID   string `json:"uploadId"`
		PartNumber int    `json:"partNumber"`
	}

	if err := render.Decode(r, &req); err != nil {
		render.Status(r, http.StatusBadRequest)
		render.JSON(w, r, map[string]string{"error": err.Error()})
		return
	}

	if req.UploadID == "" || req.PartNumber == 0 {
		render.Status(r, http.StatusBadRequest)
		render.JSON(w, r, map[string]string{"error": "uploadId and partNumber are required"})
		return
	}

	uploadInfo, err := h.GetPartUploadURL(r.Context(), req.UploadID, req.PartNumber)
	if err != nil {
		render.Status(r, http.StatusInternalServerError)
		render.JSON(w, r, map[string]string{"error": fmt.Sprintf("failed to get part upload URL: %v", err)})
		return
	}

	render.JSON(w, r, uploadInfo)
}

// HandleCompleteUpload finalizes a multipart upload
// Replaces the old "action: complete" pattern
func (h *XRPCHandler) HandleCompleteUpload(w http.ResponseWriter, r *http.Request) {
	var req struct {
		UploadID string     `json:"uploadId"`
		Digest   string     `json:"digest"`
		Parts    []PartInfo `json:"parts"`
	}

	if err := render.Decode(r, &req); err != nil {
		render.Status(r, http.StatusBadRequest)
		render.JSON(w, r, map[string]string{"error": err.Error()})
		return
	}

	if req.UploadID == "" || req.Digest == "" || len(req.Parts) == 0 {
		render.Status(r, http.StatusBadRequest)
		render.JSON(w, r, map[string]string{"error": "uploadId, digest, and parts are required"})
		return
	}

	err := h.CompleteMultipartUploadWithManager(r.Context(), req.UploadID, req.Digest, req.Parts)
	if err != nil {
		render.Status(r, http.StatusInternalServerError)
		render.JSON(w, r, map[string]string{"error": fmt.Sprintf("failed to complete upload: %v", err)})
		return
	}

	render.JSON(w, r, map[string]any{
		"status": "completed",
		"digest": req.Digest,
	})
}

// HandleAbortUpload cancels a multipart upload
// Replaces the old "action: abort" pattern
func (h *XRPCHandler) HandleAbortUpload(w http.ResponseWriter, r *http.Request) {
	var req struct {
		UploadID string `json:"uploadId"`
	}

	if err := render.Decode(r, &req); err != nil {
		render.Status(r, http.StatusBadRequest)
		render.JSON(w, r, map[string]string{"error": err.Error()})
		return
	}

	if req.UploadID == "" {
		render.Status(r, http.StatusBadRequest)
		render.JSON(w, r, map[string]string{"error": "uploadId is required"})
		return
	}

	err := h.AbortMultipartUploadWithManager(r.Context(), req.UploadID)
	if err != nil {
		render.Status(r, http.StatusInternalServerError)
		render.JSON(w, r, map[string]string{"error": fmt.Sprintf("failed to abort upload: %v", err)})
		return
	}

	render.JSON(w, r, map[string]any{
		"status": "aborted",
	})
}

// notifyManifestInfo mirrors io.atcr.hold.notifyManifest#manifestInfo: the
// slice of the pushed OCI manifest the AppView forwards to the hold.
type notifyManifestInfo struct {
	MediaType string `json:"mediaType"`
	Config    struct {
		Digest    string `json:"digest"`
		Size      int64  `json:"size"`
		MediaType string `json:"mediaType"`
	} `json:"config"`
	Layers []struct {
		Digest    string `json:"digest"`
		Size      int64  `json:"size"`
		MediaType string `json:"mediaType"`
	} `json:"layers"`
	Manifests []struct {
		Digest    string `json:"digest"`
		Size      int64  `json:"size"`
		MediaType string `json:"mediaType"`
		Platform  *struct {
			OS           string `json:"os"`
			Architecture string `json:"architecture"`
		} `json:"platform"`
	} `json:"manifests"`
	// Subject is the descriptor an OCI referrer artifact points at. Only an
	// AppView new enough to send it populates this; an older one leaves it nil,
	// which reads as "ordinary image" exactly as it did before the field existed.
	Subject *struct {
		Digest    string `json:"digest"`
		Size      int64  `json:"size"`
		MediaType string `json:"mediaType"`
	} `json:"subject"`
}

// IsMultiArch reports whether this is a manifest list / image index, which has
// child manifests instead of layers of its own.
func (m *notifyManifestInfo) IsMultiArch() bool {
	return len(m.Manifests) > 0
}

// IsReferrer reports whether this manifest is an OCI referrer artifact
// (attestation, signature, SBOM) attached to another manifest via subject,
// rather than an image in its own right.
func (m *notifyManifestInfo) IsReferrer() bool {
	return m.Subject != nil
}

// HasScannableContent reports whether a vulnerability scan of this manifest
// could find anything. Manifest lists carry no layers, and referrer artifacts
// carry metadata rather than a filesystem, so neither is worth enqueueing.
func (m *notifyManifestInfo) HasScannableContent() bool {
	return !m.IsMultiArch() && !m.IsReferrer()
}

// HandleNotifyManifest handles manifest notifications from AppView
// For pushes: Creates layer records and optionally posts to Bluesky
// For pulls: Just increments stats (no layer records or posts)
// Always increments stats (pull or push counts)
func (h *XRPCHandler) HandleNotifyManifest(w http.ResponseWriter, r *http.Request) {
	// Detach from the request context for the PDS writes below. AppView sends
	// this notify as a best-effort fire-and-forget call (with its own 30s
	// timeout), and an upstream proxy may 504 around the 10s mark. If we wrote
	// layer/config/stats records on r.Context(), a client disconnect or proxy
	// timeout cancels them mid-flight — leaving the manifest with missing
	// records ("getb tx, context canceled"). Once we start writing to the
	// embedded PDS we must finish; cap with a timeout so it can't run unbounded.
	ctx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), 2*time.Minute)
	defer cancel()

	// Validate service token (same auth as blob:write endpoints)
	validatedUser, err := pds.ValidateBlobWriteAccess(r, h.pds, h.httpClient)
	if err != nil {
		render.Status(r, http.StatusForbidden)
		render.JSON(w, r, map[string]string{"error": fmt.Sprintf("authorization failed: %v", err)})
		return
	}

	// Parse request
	var req struct {
		Repository     string             `json:"repository"`
		Tag            string             `json:"tag"`
		UserDID        string             `json:"userDid"`
		ManifestDigest string             `json:"manifestDigest"` // For building layer record AT-URIs
		Operation      string             `json:"operation"`      // "push" or "pull", defaults to "push" for backward compatibility
		Manifest       notifyManifestInfo `json:"manifest"`
	}

	if err := render.Decode(r, &req); err != nil {
		render.Status(r, http.StatusBadRequest)
		render.JSON(w, r, map[string]string{"error": err.Error()})
		return
	}

	// Default operation to "push" for backward compatibility
	operation := req.Operation
	if operation == "" {
		operation = "push"
	}

	// Validate operation
	if operation != "push" && operation != "pull" {
		render.Status(r, http.StatusBadRequest)
		render.JSON(w, r, map[string]string{"error": fmt.Sprintf("invalid operation: %s (must be 'push' or 'pull')", operation)})
		return
	}

	// Verify user DID matches token - only for pushes
	// For pulls: userDID is the repo owner (for stats), but the token belongs to the puller
	// This allows anyone to pull from a public repo and have stats tracked under the owner
	if operation == "push" && req.UserDID != validatedUser.DID {
		render.Status(r, http.StatusForbidden)
		render.JSON(w, r, map[string]string{"error": "user DID mismatch"})
		return
	}

	var layersCreated int
	var postCreated bool
	var postURI string

	// Only create layer records and Bluesky posts for pushes
	if operation == "push" {
		// Quota enforcement happens at the auth phase (AppView /auth/token
		// via pkg/appview/authgate). By the time this notification fires
		// the blobs are in S3 and the manifest is in the user's PDS, so a
		// post-hoc check here cannot block the push. We still load stats
		// because the tier name feeds scan-on-push selection below.
		stats, _ := h.pds.GetQuotaForUserWithTier(ctx, req.UserDID, h.quotaMgr)

		// Check if manifest posts are enabled
		// Read from captain record (which is synced with HOLD_BLUESKY_POSTS_ENABLED env var)
		postsEnabled := false
		_, captain, err := h.pds.GetCaptainRecord(ctx)
		if err == nil {
			postsEnabled = captain.EnableBlueskyPosts
		} else {
			// Fallback to env var if captain record doesn't exist (shouldn't happen in normal operation)
			postsEnabled = h.enableBlueskyPosts
		}

		// Build manifest AT-URI for layer records
		manifestURI := atproto.BuildManifestURI(req.UserDID, req.ManifestDigest)

		// Skip layer record creation if records already exist for this manifest
		existingLayers, _ := h.pds.ListLayerRecordsForManifest(ctx, manifestURI)
		if len(existingLayers) > 0 {
			layersCreated = len(existingLayers)
			slog.Debug("Layer records already exist for manifest, skipping creation",
				"manifestURI", manifestURI, "existing", len(existingLayers))
		} else {
			// Create layer records for each blob
			for _, layer := range req.Manifest.Layers {
				record := atproto.NewLayerRecord(
					layer.Digest,
					layer.Size,
					layer.MediaType,
					req.UserDID,
					manifestURI,
				)

				_, _, err := h.pds.CreateLayerRecord(ctx, record)
				if err != nil {
					slog.Error("Failed to create layer record", "error", err)
					// Continue creating other records
				} else {
					layersCreated++
				}
			}
		}

		// Store OCI image config as a separate record (best-effort)
		if req.Manifest.Config.Digest != "" {
			configBytes, err := h.s3Service.GetBytes(ctx, s3.BlobPath(req.Manifest.Config.Digest))
			if err != nil {
				slog.Warn("Failed to fetch config blob for image config record", "error", err, "configDigest", req.Manifest.Config.Digest)
			} else {
				configRecord := atproto.NewImageConfigRecord(manifestURI, string(configBytes))
				if _, _, err := h.pds.CreateImageConfigRecord(ctx, configRecord, req.ManifestDigest); err != nil {
					slog.Warn("Failed to create image config record", "error", err)
				}
			}
		}

		// Check if this is a multi-arch image (has manifests instead of layers)
		isMultiArch := req.Manifest.IsMultiArch()

		// Calculate total size from all layers (for single-arch images)
		var totalSize int64
		for _, layer := range req.Manifest.Layers {
			totalSize += layer.Size
		}
		totalSize += req.Manifest.Config.Size // Add config blob size

		// Extract platforms for multi-arch images
		// Filter out attestation manifests which have unknown/unknown or empty platforms
		var platforms []string
		if isMultiArch {
			for _, m := range req.Manifest.Manifests {
				if m.Platform != nil &&
					m.Platform.OS != "" && m.Platform.OS != "unknown" &&
					m.Platform.Architecture != "" && m.Platform.Architecture != "unknown" {
					platforms = append(platforms, m.Platform.OS+"/"+m.Platform.Architecture)
				}
			}
		}

		// Detect artifact type from config media type
		artifactType := "container-image"
		if strings.Contains(req.Manifest.Config.MediaType, "helm.config") {
			artifactType = "helm-chart"
		}

		// Create Bluesky post if enabled and tag is present
		// Skip posts for tagless pushes (e.g., buildx platform manifests pushed by digest)
		if postsEnabled && req.Tag != "" {
			// Resolve handle from DID (cached, 24-hour TTL)
			_, userHandle, _, resolveErr := atproto.ResolveIdentity(ctx, req.UserDID)
			if resolveErr != nil {
				slog.Warn("Failed to resolve handle for user", "did", req.UserDID, "error", resolveErr)
				userHandle = req.UserDID // Fallback to DID if resolution fails
			}

			// Extract manifest digest from first layer (or use config digest as fallback)
			manifestDigest := req.Manifest.Config.Digest
			if len(req.Manifest.Layers) > 0 {
				manifestDigest = req.Manifest.Layers[0].Digest
			}

			postURI, err = h.pds.CreateManifestPost(
				ctx,
				&h.s3Service,
				req.Repository,
				req.Tag,
				userHandle,
				req.UserDID,
				manifestDigest,
				totalSize,
				platforms,
				artifactType,
			)
			if err != nil {
				slog.Error("Failed to create manifest post", "error", err)
			} else {
				postCreated = true
			}
		}

		// Enqueue scan job if scanner is connected (skip manifest lists and
		// referrer artifacts such as attestations — no scannable content).
		if h.scanBroadcaster != nil && req.Manifest.HasScannableContent() {
			tier := "deckhand"
			if stats != nil && stats.Tier != "" {
				tier = stats.Tier
			}

			// Check if this tier gets scan-on-push.
			// Captain ("owner") always gets scan-on-push.
			// When quotas are disabled, all pushes trigger scans (backwards compat).
			shouldScan := tier == "owner" ||
				h.quotaMgr == nil || !h.quotaMgr.IsEnabled() ||
				h.quotaMgr.ScanOnPush(tier)

			if shouldScan {
				configJSON, _ := json.Marshal(req.Manifest.Config)
				layersJSON, _ := json.Marshal(req.Manifest.Layers)

				// Resolve handle for scanner context
				_, userHandle, _, resolveErr := atproto.ResolveIdentity(ctx, req.UserDID)
				if resolveErr != nil {
					userHandle = req.UserDID
				}

				if err := h.scanBroadcaster.Enqueue(&pds.ScanJobEvent{
					ManifestDigest: req.ManifestDigest,
					Repository:     req.Repository,
					Tag:            req.Tag,
					UserDID:        req.UserDID,
					UserHandle:     userHandle,
					Tier:           tier,
					Config:         configJSON,
					Layers:         layersJSON,
				}); err != nil {
					slog.Error("Failed to enqueue scan job",
						"repository", req.Repository,
						"error", err)
				}
			} else {
				slog.Debug("Scan-on-push skipped for tier",
					"tier", tier,
					"repository", req.Repository,
					"userDid", req.UserDID)
			}
		}
	}

	// ALWAYS increment stats (even if Bluesky posts disabled, even for pulls)
	statsUpdated := false
	if err := h.pds.IncrementStats(ctx, req.UserDID, req.Repository, operation); err != nil {
		slog.Error("Failed to increment stats", "operation", operation, "error", err)
	} else {
		statsUpdated = true
	}

	// Also increment daily stats for trend tracking
	if err := h.pds.IncrementDailyStats(ctx, req.UserDID, req.Repository, operation); err != nil {
		slog.Warn("Failed to increment daily stats", "operation", operation, "error", err)
	}

	// Return response
	resp := map[string]any{
		"success":      statsUpdated || layersCreated > 0 || postCreated,
		"operation":    operation,
		"statsUpdated": statsUpdated,
	}

	// Only include push-specific fields for push operations
	if operation == "push" {
		resp["layersCreated"] = layersCreated
		resp["postCreated"] = postCreated
		if postURI != "" {
			resp["postUri"] = postURI
		}
	}

	render.JSON(w, r, resp)
}

// requireBlobWriteAccess middleware - validates DPoP + OAuth and checks for blob:write permission
func (h *XRPCHandler) requireBlobWriteAccess(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, err := pds.ValidateBlobWriteAccess(r, h.pds, h.httpClient)
		if err != nil {
			http.Error(w, fmt.Sprintf("authorization failed: %v", err), http.StatusForbidden)
			return
		}

		// Validation successful - user has blob:write permission
		// No need to store user in context since handlers don't need it
		next.ServeHTTP(w, r)
	})
}