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

package oci

import (
	"context"
	"fmt"
	"log/slog"
	"sort"
	"strings"
	"sync"
	"time"

	"atcr.io/pkg/s3"
	awss3 "github.com/aws/aws-sdk-go-v2/service/s3"
	s3types "github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/google/uuid"
)

// PartInfo represents an uploaded part with its ETag
type PartInfo struct {
	PartNumber int    `json:"part_number"`
	ETag       string `json:"etag"`
}

// PartUploadInfo contains the presigned URL for uploading a part
type PartUploadInfo struct {
	URL string `json:"url"` // Presigned URL to PUT the part to
}

// MultipartSession tracks an in-progress multipart upload
type MultipartSession struct {
	UploadID     string    // Unique upload ID
	Digest       string    // Target digest path
	S3UploadID   string    // S3 upload ID
	CreatedAt    time.Time // When upload started
	LastActivity time.Time // Last part upload
}

// MultipartPart represents a single part in a multipart upload
type MultipartPart struct {
	PartNumber int       // Part number (1-indexed)
	ETag       string    // ETag from S3
	Size       int64     // Part size in bytes
	UploadedAt time.Time // When part was uploaded
}

// MultipartManager manages multipart upload sessions
type MultipartManager struct {
	sessions map[string]*MultipartSession // uploadID -> session
	mu       sync.RWMutex                 // Protects sessions map
}

// NewMultipartManager creates a new multipart manager
func NewMultipartManager() *MultipartManager {
	mgr := &MultipartManager{
		sessions: make(map[string]*MultipartSession),
	}

	// Start cleanup goroutine for abandoned uploads
	go mgr.cleanupLoop()

	return mgr
}

// cleanupLoop periodically cleans up expired sessions
func (m *MultipartManager) cleanupLoop() {
	ticker := time.NewTicker(15 * time.Minute)
	defer ticker.Stop()

	for range ticker.C {
		m.cleanupExpiredSessions()
	}
}

// cleanupExpiredSessions removes sessions inactive for >24 hours
func (m *MultipartManager) cleanupExpiredSessions() {
	m.mu.Lock()
	defer m.mu.Unlock()

	now := time.Now()
	for uploadID, session := range m.sessions {
		if now.Sub(session.LastActivity) > 24*time.Hour {
			slog.Debug("Cleaning up expired multipart session",
				"uploadID", uploadID,
				"age", now.Sub(session.CreatedAt))
			delete(m.sessions, uploadID)
		}
	}
}

// CreateSession creates a new multipart upload session
func (m *MultipartManager) CreateSession(digest string, s3UploadID string) *MultipartSession {
	uploadID := uuid.New().String()

	session := &MultipartSession{
		UploadID:     uploadID,
		Digest:       digest,
		S3UploadID:   s3UploadID,
		CreatedAt:    time.Now(),
		LastActivity: time.Now(),
	}

	m.mu.Lock()
	m.sessions[uploadID] = session
	m.mu.Unlock()

	slog.Debug("Created multipart session",
		"uploadID", uploadID,
		"digest", digest)
	return session
}

// GetSession retrieves a multipart session by upload ID
func (m *MultipartManager) GetSession(uploadID string) (*MultipartSession, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	session, ok := m.sessions[uploadID]
	if !ok {
		return nil, fmt.Errorf("multipart session not found: %s", uploadID)
	}

	return session, nil
}

// DeleteSession removes a multipart session
func (m *MultipartManager) DeleteSession(uploadID string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	delete(m.sessions, uploadID)
	slog.Debug("Deleted multipart session", "uploadID", uploadID)
}

// StartMultipartUploadWithManager initiates a multipart upload using the manager
// Returns the upload ID for tracking the session
func (h *XRPCHandler) StartMultipartUploadWithManager(ctx context.Context, digest string) (string, error) {
	if h.s3Service.Client == nil {
		return "", fmt.Errorf("S3 not configured - S3 is required for blob storage")
	}

	path := s3.BlobPath(digest)
	s3Key := strings.TrimPrefix(path, "/")
	if h.s3Service.PathPrefix != "" {
		s3Key = h.s3Service.PathPrefix + "/" + s3Key
	}

	result, err := h.s3Service.Client.CreateMultipartUpload(ctx, &awss3.CreateMultipartUploadInput{
		Bucket: &h.s3Service.Bucket,
		Key:    &s3Key,
	})
	if err != nil {
		return "", fmt.Errorf("failed to start S3 multipart upload: %w", err)
	}

	s3UploadID := *result.UploadId
	session := h.MultipartMgr.CreateSession(digest, s3UploadID)
	slog.Debug("Started S3 multipart upload",
		"digest", digest,
		"uploadID", session.UploadID,
		"s3UploadID", s3UploadID)
	return session.UploadID, nil
}

// GetPartUploadURL generates a presigned URL for uploading a part
func (h *XRPCHandler) GetPartUploadURL(ctx context.Context, uploadID string, partNumber int) (*PartUploadInfo, error) {
	session, err := h.MultipartMgr.GetSession(uploadID)
	if err != nil {
		return nil, err
	}

	if h.s3Service.Client == nil {
		return nil, fmt.Errorf("S3 not configured")
	}

	path := s3.BlobPath(session.Digest)
	s3Key := strings.TrimPrefix(path, "/")
	if h.s3Service.PathPrefix != "" {
		s3Key = h.s3Service.PathPrefix + "/" + s3Key
	}
	pnum := int32(partNumber)
	url, err := h.s3Service.Client.PresignUploadPart(ctx, &awss3.UploadPartInput{
		Bucket:     &h.s3Service.Bucket,
		Key:        &s3Key,
		UploadId:   &session.S3UploadID,
		PartNumber: &pnum,
	}, 15*time.Minute)
	if err != nil {
		return nil, err
	}

	slog.Debug("Generated part presigned URL",
		"digest", session.Digest,
		"uploadID", uploadID,
		"part", partNumber)

	return &PartUploadInfo{
		URL: url,
	}, nil
}

// CompleteMultipartUploadWithManager completes a multipart upload and moves to final location
// finalDigest is the real digest (e.g., "sha256:abc123...") for the final storage location
// session.Digest is the temp location (e.g., "uploads/temp-<uuid>")
func (h *XRPCHandler) CompleteMultipartUploadWithManager(ctx context.Context, uploadID string, finalDigest string, parts []PartInfo) error {
	session, err := h.MultipartMgr.GetSession(uploadID)
	defer h.MultipartMgr.DeleteSession(uploadID)
	if err != nil {
		return err
	}

	if h.s3Service.Client == nil {
		return fmt.Errorf("S3 not configured")
	}

	// Sort parts by part number (S3 requires ascending order)
	sort.Slice(parts, func(i, j int) bool {
		return parts[i].PartNumber < parts[j].PartNumber
	})

	// Convert to S3 CompletedPart format
	// IMPORTANT: S3 requires ETags to be quoted in the CompleteMultipartUpload XML
	s3Parts := make([]s3types.CompletedPart, len(parts))
	for i, p := range parts {
		etag := normalizeETag(p.ETag)
		pnum := int32(p.PartNumber)
		s3Parts[i] = s3types.CompletedPart{
			PartNumber: &pnum,
			ETag:       &etag,
		}
	}
	sourcePath := s3.BlobPath(session.Digest)
	s3Key := strings.TrimPrefix(sourcePath, "/")
	if h.s3Service.PathPrefix != "" {
		s3Key = h.s3Service.PathPrefix + "/" + s3Key
	}

	_, err = h.s3Service.Client.CompleteMultipartUpload(ctx, &awss3.CompleteMultipartUploadInput{
		Bucket:   &h.s3Service.Bucket,
		Key:      &s3Key,
		UploadId: &session.S3UploadID,
		MultipartUpload: &s3types.CompletedMultipartUpload{
			Parts: s3Parts,
		},
	})
	if err != nil {
		return fmt.Errorf("failed to complete multipart upload: digest=%s, uploadID=%s, err=%v", session.Digest, uploadID, err)
	}
	slog.Info("Completed S3 native multipart at temp location",
		"digest", session.Digest,
		"uploadID", session.UploadID,
		"parts", len(s3Parts))

	// Verify the blob exists at temp location before moving
	destPath := s3.BlobPath(finalDigest)
	slog.Debug("About to move blob",
		"source", sourcePath,
		"dest", destPath)

	if _, err := h.driver.Stat(ctx, sourcePath); err != nil {
		slog.Error("Source blob not found after multipart complete",
			"path", sourcePath,
			"error", err)
		return fmt.Errorf("source blob not found after multipart complete: %w", err)
	}
	slog.Debug("Source blob verified", "path", sourcePath)

	// Move from temp to final digest location using driver
	// Driver handles path management correctly (including S3 prefix)
	if err := h.driver.Move(ctx, sourcePath, destPath); err != nil {
		slog.Error("Failed to move blob",
			"source", sourcePath,
			"dest", destPath,
			"error", err)
		return fmt.Errorf("failed to move blob to final location: %w", err)
	}

	slog.Info("Moved blob to final location",
		"from", session.Digest,
		"to", finalDigest,
		"sourcePath", sourcePath,
		"destPath", destPath)
	return nil
}

// AbortMultipartUploadWithManager aborts a multipart upload
func (h *XRPCHandler) AbortMultipartUploadWithManager(ctx context.Context, uploadID string) error {
	session, err := h.MultipartMgr.GetSession(uploadID)
	defer h.MultipartMgr.DeleteSession(uploadID)
	if err != nil {
		return err
	}

	if h.s3Service.Client == nil {
		return fmt.Errorf("S3 not configured")
	}

	path := s3.BlobPath(session.Digest)
	s3Key := strings.TrimPrefix(path, "/")
	if h.s3Service.PathPrefix != "" {
		s3Key = h.s3Service.PathPrefix + "/" + s3Key
	}

	_, err = h.s3Service.Client.AbortMultipartUpload(ctx, &awss3.AbortMultipartUploadInput{
		Bucket:   &h.s3Service.Bucket,
		Key:      &s3Key,
		UploadId: &session.S3UploadID,
	})
	if err != nil {
		return fmt.Errorf("failed to abort multipart upload: digest=%s, uploadID=%s, err=%v", session.Digest, uploadID, err)
	}
	slog.Debug("Aborted S3 multipart",
		"digest", session.Digest,
		"uploadID", session.UploadID)
	return nil
}

// normalizeETag ensures an ETag has quotes (required by S3 CompleteMultipartUpload)
// S3 returns ETags with quotes, but HTTP clients may strip them
func normalizeETag(etag string) string {
	// Already has quotes
	if strings.HasPrefix(etag, "\"") && strings.HasSuffix(etag, "\"") {
		return etag
	}
	// Add quotes
	return fmt.Sprintf("\"%s\"", etag)
}