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

package oci

import (
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strconv"

	"atcr.io/pkg/atproto"
	"atcr.io/pkg/hold/pds"
	"atcr.io/pkg/s3"
	storagedriver "github.com/distribution/distribution/v3/registry/storage/driver"
	"github.com/go-chi/chi/v5"
)

// XRPCHandler handles OCI-specific XRPC endpoints for multipart uploads
type XRPCHandler struct {
	driver               storagedriver.StorageDriver
	disablePresignedURLs bool
	s3Service            s3.S3Service
	MultipartMgr         *MultipartManager // Exported for access in route handlers
	pds                  *pds.HoldPDS
	httpClient           pds.HTTPClient
	enableBlueskyPosts   bool
}

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

// 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.Put(atproto.HoldUploadPart, h.HandleUploadPart)
		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 := DecodeJSON(r, &req); err != nil {
		RespondError(w, http.StatusBadRequest, err.Error())
		return
	}

	if req.Digest == "" {
		RespondError(w, http.StatusBadRequest, "digest is required")
		return
	}

	uploadID, _, err := h.StartMultipartUploadWithManager(r.Context(), req.Digest)
	if err != nil {
		RespondError(w, http.StatusInternalServerError, fmt.Sprintf("failed to initiate upload: %v", err))
		return
	}

	RespondJSON(w, http.StatusOK, 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 := DecodeJSON(r, &req); err != nil {
		RespondError(w, http.StatusBadRequest, err.Error())
		return
	}

	if req.UploadID == "" || req.PartNumber == 0 {
		RespondError(w, http.StatusBadRequest, "uploadId and partNumber are required")
		return
	}

	uploadInfo, err := h.GetPartUploadURL(r.Context(), req.UploadID, req.PartNumber)
	if err != nil {
		RespondError(w, http.StatusInternalServerError, fmt.Sprintf("failed to get part upload URL: %v", err))
		return
	}

	RespondJSON(w, http.StatusOK, uploadInfo)
}

// HandleUploadPart handles direct buffered part uploads
// Moved from pds/xrpc.go - this is OCI-specific multipart upload logic
func (h *XRPCHandler) HandleUploadPart(w http.ResponseWriter, r *http.Request) {
	uploadID := r.Header.Get("X-Upload-Id")
	partNumberStr := r.Header.Get("X-Part-Number")

	if uploadID == "" || partNumberStr == "" {
		RespondError(w, http.StatusBadRequest, "X-Upload-Id and X-Part-Number headers are required")
		return
	}

	partNumber, err := strconv.Atoi(partNumberStr)
	if err != nil {
		RespondError(w, http.StatusBadRequest, fmt.Sprintf("invalid part number: %v", err))
		return
	}

	data, err := io.ReadAll(r.Body)
	if err != nil {
		RespondError(w, http.StatusInternalServerError, fmt.Sprintf("failed to read part data: %v", err))
		return
	}

	etag, err := h.HandleBufferedPartUpload(r.Context(), uploadID, partNumber, data)
	if err != nil {
		RespondError(w, http.StatusInternalServerError, fmt.Sprintf("failed to upload part: %v", err))
		return
	}

	RespondJSON(w, http.StatusOK, map[string]any{
		"etag": etag,
	})
}

// 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 := DecodeJSON(r, &req); err != nil {
		RespondError(w, http.StatusBadRequest, err.Error())
		return
	}

	if req.UploadID == "" || req.Digest == "" || len(req.Parts) == 0 {
		RespondError(w, http.StatusBadRequest, "uploadId, digest, and parts are required")
		return
	}

	err := h.CompleteMultipartUploadWithManager(r.Context(), req.UploadID, req.Digest, req.Parts)
	if err != nil {
		RespondError(w, http.StatusInternalServerError, fmt.Sprintf("failed to complete upload: %v", err))
		return
	}

	RespondJSON(w, http.StatusOK, 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 := DecodeJSON(r, &req); err != nil {
		RespondError(w, http.StatusBadRequest, err.Error())
		return
	}

	if req.UploadID == "" {
		RespondError(w, http.StatusBadRequest, "uploadId is required")
		return
	}

	err := h.AbortMultipartUploadWithManager(r.Context(), req.UploadID)
	if err != nil {
		RespondError(w, http.StatusInternalServerError, fmt.Sprintf("failed to abort upload: %v", err))
		return
	}

	RespondJSON(w, http.StatusOK, map[string]any{
		"status": "aborted",
	})
}

// HandleNotifyManifest handles manifest upload notifications from AppView
// Creates layer records and optionally posts to Bluesky
func (h *XRPCHandler) HandleNotifyManifest(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

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

	// Parse request
	var req struct {
		Repository string `json:"repository"`
		Tag        string `json:"tag"`
		UserDID    string `json:"userDid"`
		UserHandle string `json:"userHandle"`
		Manifest   struct {
			MediaType string `json:"mediaType"`
			Config    struct {
				Digest string `json:"digest"`
				Size   int64  `json:"size"`
			} `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"`
		} `json:"manifest"`
	}

	if err := DecodeJSON(r, &req); err != nil {
		RespondError(w, http.StatusBadRequest, err.Error())
		return
	}

	// Verify user DID matches token
	if req.UserDID != validatedUser.DID {
		RespondError(w, http.StatusForbidden, "user DID mismatch")
		return
	}

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

	// Create layer records for each blob
	layersCreated := 0
	for _, layer := range req.Manifest.Layers {
		record := atproto.NewLayerRecord(
			layer.Digest,
			layer.Size,
			layer.MediaType,
			req.Repository,
			req.UserDID,
			req.UserHandle,
		)

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

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

	// 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
	var platforms []string
	if isMultiArch {
		for _, m := range req.Manifest.Manifests {
			if m.Platform != nil {
				platforms = append(platforms, m.Platform.OS+"/"+m.Platform.Architecture)
			}
		}
	}

	// Create Bluesky post if enabled
	var postURI string
	postCreated := false
	if postsEnabled {
		// 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.driver,
			req.Repository,
			req.Tag,
			req.UserHandle,
			req.UserDID,
			manifestDigest,
			totalSize,
			platforms,
		)
		if err != nil {
			slog.Error("Failed to create manifest post", "error", err)
		} else {
			postCreated = true
		}
	}

	// Return response
	resp := map[string]any{
		"success":       layersCreated > 0 || postCreated,
		"layersCreated": layersCreated,
		"postCreated":   postCreated,
	}
	if postURI != "" {
		resp["postUri"] = postURI
	}
	if err != nil && layersCreated == 0 && !postCreated {
		resp["error"] = err.Error()
	}

	RespondJSON(w, http.StatusOK, 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)
	})
}