at-container-registry/pkg/storage/proxy_blob_store.go

package storage

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/distribution/distribution/v3"
	"github.com/opencontainers/go-digest"
)

// Global upload tracking (shared across all ProxyBlobStore instances)
// This is necessary because distribution creates new repository/blob store instances per request
var (
	globalUploads   = make(map[string]*ProxyBlobWriter)
	globalUploadsMu sync.RWMutex
)

// ProxyBlobStore proxies blob requests to an external storage service
type ProxyBlobStore struct {
	storageEndpoint string
	httpClient      *http.Client
	did             string
}

// NewProxyBlobStore creates a new proxy blob store
func NewProxyBlobStore(storageEndpoint, did string) *ProxyBlobStore {
	fmt.Printf("DEBUG [proxy_blob_store]: NewProxyBlobStore created with endpoint=%s, did=%s\n", storageEndpoint, did)
	return &ProxyBlobStore{
		storageEndpoint: storageEndpoint,
		httpClient: &http.Client{
			Timeout: 5 * time.Minute, // Timeout for presigned URL requests and uploads
		},
		did: did,
	}
}

// Stat returns the descriptor for a blob
func (p *ProxyBlobStore) Stat(ctx context.Context, dgst digest.Digest) (distribution.Descriptor, error) {
	// For simplicity, we'll just check if we can get a download URL
	// In production, you'd want a dedicated stat endpoint
	url, err := p.getDownloadURL(ctx, dgst)
	if err != nil {
		return distribution.Descriptor{}, distribution.ErrBlobUnknown
	}

	// We don't have size info from the storage service
	// Return a minimal descriptor
	return distribution.Descriptor{
		Digest:    dgst,
		MediaType: "application/octet-stream",
		URLs:      []string{url},
	}, nil
}

// Get retrieves a blob
func (p *ProxyBlobStore) Get(ctx context.Context, dgst digest.Digest) ([]byte, error) {
	url, err := p.getDownloadURL(ctx, dgst)
	if err != nil {
		return nil, err
	}

	// Download the blob
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, distribution.ErrBlobUnknown
	}

	return io.ReadAll(resp.Body)
}

// Open returns a reader for a blob
func (p *ProxyBlobStore) Open(ctx context.Context, dgst digest.Digest) (io.ReadSeekCloser, error) {
	url, err := p.getDownloadURL(ctx, dgst)
	if err != nil {
		return nil, err
	}

	// Download the blob
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		resp.Body.Close()
		return nil, distribution.ErrBlobUnknown
	}

	// Wrap in a ReadSeekCloser
	return &readSeekCloser{
		ReadCloser: resp.Body,
	}, nil
}

// Put stores a blob
func (p *ProxyBlobStore) Put(ctx context.Context, mediaType string, content []byte) (distribution.Descriptor, error) {
	// Calculate digest
	dgst := digest.FromBytes(content)

	// Get upload URL
	url, err := p.getUploadURL(ctx, dgst, int64(len(content)))
	if err != nil {
		return distribution.Descriptor{}, err
	}

	// Upload the blob
	req, err := http.NewRequestWithContext(ctx, "PUT", url, bytes.NewReader(content))
	if err != nil {
		return distribution.Descriptor{}, err
	}
	req.Header.Set("Content-Type", "application/octet-stream")

	resp, err := p.httpClient.Do(req)
	if err != nil {
		return distribution.Descriptor{}, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated {
		return distribution.Descriptor{}, fmt.Errorf("upload failed with status %d", resp.StatusCode)
	}

	return distribution.Descriptor{
		Digest:    dgst,
		Size:      int64(len(content)),
		MediaType: mediaType,
	}, nil
}

// Delete removes a blob
func (p *ProxyBlobStore) Delete(ctx context.Context, dgst digest.Digest) error {
	// Not implemented - storage service would need a delete endpoint
	return fmt.Errorf("delete not supported for proxy blob store")
}

// ServeBlob serves a blob via HTTP redirect
func (p *ProxyBlobStore) ServeBlob(ctx context.Context, w http.ResponseWriter, r *http.Request, dgst digest.Digest) error {
	// Get presigned download URL
	url, err := p.getDownloadURL(ctx, dgst)
	if err != nil {
		return err
	}

	// Redirect to presigned URL
	http.Redirect(w, r, url, http.StatusTemporaryRedirect)
	return nil
}

// Create returns a blob writer for uploading
func (p *ProxyBlobStore) Create(ctx context.Context, options ...distribution.BlobCreateOption) (distribution.BlobWriter, error) {
	// Parse options
	var opts distribution.CreateOptions
	for _, option := range options {
		if err := option.Apply(&opts); err != nil {
			return nil, err
		}
	}

	// Create proxy blob writer
	writer := &ProxyBlobWriter{
		store:     p,
		ctx:       ctx,
		options:   opts,
		id:        fmt.Sprintf("upload-%d", time.Now().UnixNano()),
		startedAt: time.Now(),
	}

	// Store in global uploads map for resume support
	globalUploadsMu.Lock()
	globalUploads[writer.id] = writer
	globalUploadsMu.Unlock()

	return writer, nil
}

// Resume returns a blob writer for resuming an upload
func (p *ProxyBlobStore) Resume(ctx context.Context, id string) (distribution.BlobWriter, error) {
	// Retrieve upload from global map
	globalUploadsMu.RLock()
	writer, ok := globalUploads[id]
	globalUploadsMu.RUnlock()

	if !ok {
		return nil, distribution.ErrBlobUploadUnknown
	}

	return writer, nil
}

// getDownloadURL requests a presigned download URL from the storage service
func (p *ProxyBlobStore) getDownloadURL(ctx context.Context, dgst digest.Digest) (string, error) {
	reqBody := map[string]any{
		"did":    p.did,
		"digest": dgst.String(),
	}

	body, err := json.Marshal(reqBody)
	if err != nil {
		return "", err
	}

	url := fmt.Sprintf("%s/get-presigned-url", p.storageEndpoint)
	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := p.httpClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("failed to get download URL: status %d", resp.StatusCode)
	}

	var result struct {
		URL string `json:"url"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", err
	}

	return result.URL, nil
}

// getUploadURL requests a presigned upload URL from the storage service
func (p *ProxyBlobStore) getUploadURL(ctx context.Context, dgst digest.Digest, size int64) (string, error) {
	fmt.Printf("DEBUG [proxy_blob_store/getUploadURL]: storageEndpoint=%s, digest=%s\n", p.storageEndpoint, dgst)

	reqBody := map[string]any{
		"did":    p.did,
		"digest": dgst.String(),
		"size":   size,
	}

	body, err := json.Marshal(reqBody)
	if err != nil {
		return "", err
	}

	url := fmt.Sprintf("%s/put-presigned-url", p.storageEndpoint)
	fmt.Printf("DEBUG [proxy_blob_store/getUploadURL]: Calling %s\n", url)
	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := p.httpClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("failed to get upload URL: status %d", resp.StatusCode)
	}

	var result struct {
		URL string `json:"url"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", err
	}

	fmt.Printf("DEBUG [proxy_blob_store/getUploadURL]: Got presigned URL=%s\n", result.URL)
	return result.URL, nil
}

// ProxyBlobWriter implements distribution.BlobWriter for proxy uploads
type ProxyBlobWriter struct {
	store     *ProxyBlobStore
	ctx       context.Context
	options   distribution.CreateOptions
	buffer    bytes.Buffer
	size      int64
	closed    bool
	id        string
	startedAt time.Time
}

// ID returns the upload ID
func (w *ProxyBlobWriter) ID() string {
	return w.id
}

// StartedAt returns when the upload started
func (w *ProxyBlobWriter) StartedAt() time.Time {
	return w.startedAt
}

// Write writes data to the upload
func (w *ProxyBlobWriter) Write(p []byte) (int, error) {
	if w.closed {
		return 0, fmt.Errorf("writer closed")
	}
	n, err := w.buffer.Write(p)
	w.size += int64(n)
	return n, err
}

// ReadFrom reads from a reader
func (w *ProxyBlobWriter) ReadFrom(r io.Reader) (int64, error) {
	if w.closed {
		return 0, fmt.Errorf("writer closed")
	}
	n, err := w.buffer.ReadFrom(r)
	w.size += n
	return n, err
}

// Size returns the current size
func (w *ProxyBlobWriter) Size() int64 {
	return w.size
}

// Commit finalizes the upload
func (w *ProxyBlobWriter) Commit(ctx context.Context, desc distribution.Descriptor) (distribution.Descriptor, error) {
	if w.closed {
		return distribution.Descriptor{}, fmt.Errorf("writer closed")
	}
	w.closed = true

	// Remove from global uploads map
	globalUploadsMu.Lock()
	delete(globalUploads, w.id)
	globalUploadsMu.Unlock()

	// Upload the buffered content
	content := w.buffer.Bytes()
	dgst := digest.FromBytes(content)

	// Verify digest matches
	if desc.Digest != "" && dgst != desc.Digest {
		return distribution.Descriptor{}, fmt.Errorf("digest mismatch")
	}

	// Get upload URL
	url, err := w.store.getUploadURL(ctx, dgst, int64(len(content)))
	if err != nil {
		return distribution.Descriptor{}, err
	}

	// Upload
	req, err := http.NewRequestWithContext(ctx, "PUT", url, bytes.NewReader(content))
	if err != nil {
		return distribution.Descriptor{}, err
	}
	req.Header.Set("Content-Type", "application/octet-stream")

	resp, err := w.store.httpClient.Do(req)
	if err != nil {
		return distribution.Descriptor{}, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated {
		return distribution.Descriptor{}, fmt.Errorf("upload failed: status %d", resp.StatusCode)
	}

	return distribution.Descriptor{
		Digest:    dgst,
		Size:      int64(len(content)),
		MediaType: desc.MediaType,
	}, nil
}

// Cancel cancels the upload
func (w *ProxyBlobWriter) Cancel(ctx context.Context) error {
	w.closed = true

	// Remove from global uploads map
	globalUploadsMu.Lock()
	delete(globalUploads, w.id)
	globalUploadsMu.Unlock()

	return nil
}

// Close closes the writer
// NOTE: For resumable uploads, we don't mark as closed here
// Distribution calls Close() after each PATCH, but the upload may continue
// Only Commit() and Cancel() actually finalize the upload
func (w *ProxyBlobWriter) Close() error {
	// Don't set w.closed = true here - allow resuming
	return nil
}

// readSeekCloser wraps an io.ReadCloser to implement ReadSeekCloser
type readSeekCloser struct {
	io.ReadCloser
}

func (r *readSeekCloser) Seek(offset int64, whence int) (int64, error) {
	// Not implemented - would need buffering or re-downloading
	return 0, fmt.Errorf("seek not supported")
}