at-container-registry/pkg/hold/s3.go

package hold

import (
	"context"
	"fmt"
	"log"
	"sort"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

// initS3Client initializes the S3 client for presigned URL generation
// Returns nil error if S3 client is successfully initialized
// Returns error if storage is not S3 or if initialization fails (service will fall back to proxy mode)
func (s *HoldService) initS3Client() error {
	// Check if presigned URLs are explicitly disabled
	if s.config.Server.DisablePresignedURLs {
		log.Printf("⚠️  S3 presigned URLs DISABLED by config (DISABLE_PRESIGNED_URLS=true)")
		log.Printf("   All uploads will use buffered mode (parts buffered in hold service)")
		return nil // Not an error - just using buffered mode
	}

	// Check if storage driver is S3
	if s.config.Storage.Type() != "s3" {
		log.Printf("Storage driver is %s (not S3), presigned URLs disabled", s.config.Storage.Type())
		return nil // Not an error - just using different driver
	}

	// Extract S3 configuration from storage parameters
	params := s.config.Storage.Parameters()

	// Extract required S3 configuration
	region, _ := params["region"].(string)
	if region == "" {
		region = "us-east-1" // Default region
	}

	accessKey, _ := params["accesskey"].(string)
	secretKey, _ := params["secretkey"].(string)
	bucket, _ := params["bucket"].(string)

	if bucket == "" {
		return fmt.Errorf("S3 bucket not configured")
	}

	// Build AWS config
	awsConfig := &aws.Config{
		Region: aws.String(region),
	}

	// Add credentials if provided (allow IAM role auth if not provided)
	if accessKey != "" && secretKey != "" {
		awsConfig.Credentials = credentials.NewStaticCredentials(accessKey, secretKey, "")
	}

	// Add custom endpoint for S3-compatible services (Storj, MinIO, R2, etc.)
	if endpoint, ok := params["regionendpoint"].(string); ok && endpoint != "" {
		awsConfig.Endpoint = aws.String(endpoint)
		awsConfig.S3ForcePathStyle = aws.Bool(true) // Required for MinIO, Storj
	}

	// Create AWS session
	sess, err := session.NewSession(awsConfig)
	if err != nil {
		return fmt.Errorf("failed to create AWS session: %w", err)
	}

	// Create S3 client
	s.s3Client = s3.New(sess)
	s.bucket = bucket

	// Extract path prefix if configured (rootdirectory in S3 params)
	if rootDir, ok := params["rootdirectory"].(string); ok && rootDir != "" {
		s.s3PathPrefix = strings.TrimPrefix(rootDir, "/")
	}

	log.Printf("✅ S3 presigned URLs enabled")

	return nil
}

// startMultipartUpload initiates a multipart upload and returns upload ID
func (s *HoldService) startMultipartUpload(ctx context.Context, digest string) (string, error) {
	if s.s3Client == nil {
		return "", fmt.Errorf("S3 not configured")
	}

	path := blobPath(digest)
	s3Key := strings.TrimPrefix(path, "/")
	if s.s3PathPrefix != "" {
		s3Key = s.s3PathPrefix + "/" + s3Key
	}

	result, err := s.s3Client.CreateMultipartUploadWithContext(ctx, &s3.CreateMultipartUploadInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(s3Key),
	})
	if err != nil {
		return "", err
	}

	log.Printf("Started multipart upload: digest=%s, uploadID=%s", digest, *result.UploadId)
	return *result.UploadId, nil
}

// getPartPresignedURL generates presigned URL for a specific part
func (s *HoldService) getPartPresignedURL(ctx context.Context, digest, uploadID string, partNumber int) (string, error) {
	if s.s3Client == nil {
		return "", fmt.Errorf("S3 not configured")
	}

	path := blobPath(digest)
	s3Key := strings.TrimPrefix(path, "/")
	if s.s3PathPrefix != "" {
		s3Key = s.s3PathPrefix + "/" + s3Key
	}

	req, _ := s.s3Client.UploadPartRequest(&s3.UploadPartInput{
		Bucket:     aws.String(s.bucket),
		Key:        aws.String(s3Key),
		UploadId:   aws.String(uploadID),
		PartNumber: aws.Int64(int64(partNumber)),
	})

	url, err := req.Presign(15 * time.Minute)
	if err != nil {
		return "", err
	}

	log.Printf("Generated part presigned URL: digest=%s, uploadID=%s, part=%d", digest, uploadID, partNumber)
	return url, 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)
}

// completeMultipartUpload finalizes the multipart upload
func (s *HoldService) completeMultipartUpload(ctx context.Context, digest, uploadID string, parts []CompletedPart) error {
	if s.s3Client == nil {
		return fmt.Errorf("S3 not configured")
	}

	path := blobPath(digest)
	s3Key := strings.TrimPrefix(path, "/")
	if s.s3PathPrefix != "" {
		s3Key = s.s3PathPrefix + "/" + s3Key
	}

	// 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([]*s3.CompletedPart, len(parts))
	for i, p := range parts {
		etag := normalizeETag(p.ETag)
		s3Parts[i] = &s3.CompletedPart{
			PartNumber: aws.Int64(int64(p.PartNumber)),
			ETag:       aws.String(etag),
		}
	}

	_, err := s.s3Client.CompleteMultipartUploadWithContext(ctx, &s3.CompleteMultipartUploadInput{
		Bucket:   aws.String(s.bucket),
		Key:      aws.String(s3Key),
		UploadId: aws.String(uploadID),
		MultipartUpload: &s3.CompletedMultipartUpload{
			Parts: s3Parts,
		},
	})

	if err != nil {
		log.Printf("Failed to complete multipart upload: digest=%s, uploadID=%s, err=%v", digest, uploadID, err)
		return err
	}

	log.Printf("Completed multipart upload: digest=%s, uploadID=%s, parts=%d", digest, uploadID, len(parts))
	return nil
}

// abortMultipartUpload aborts an in-progress multipart upload
func (s *HoldService) abortMultipartUpload(ctx context.Context, digest, uploadID string) error {
	if s.s3Client == nil {
		return fmt.Errorf("S3 not configured")
	}

	path := blobPath(digest)
	s3Key := strings.TrimPrefix(path, "/")
	if s.s3PathPrefix != "" {
		s3Key = s.s3PathPrefix + "/" + s3Key
	}

	_, err := s.s3Client.AbortMultipartUploadWithContext(ctx, &s3.AbortMultipartUploadInput{
		Bucket:   aws.String(s.bucket),
		Key:      aws.String(s3Key),
		UploadId: aws.String(uploadID),
	})

	if err != nil {
		log.Printf("Failed to abort multipart upload: digest=%s, uploadID=%s, err=%v", digest, uploadID, err)
		return err
	}

	log.Printf("Aborted multipart upload: digest=%s, uploadID=%s", digest, uploadID)
	return nil
}