Files
at-container-registry/pkg/appview/storage/proxy_blob_store.go
T

773 lines
22 KiB
Go

package storage
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"sync"
"time"
"atcr.io/pkg/atproto"
"atcr.io/pkg/auth"
"github.com/distribution/distribution/v3"
"github.com/opencontainers/go-digest"
)
const (
// maxChunkSize is the maximum buffer size before flushing to hold service
// Matches S3's minimum multipart upload size
maxChunkSize = 10 * 1024 * 1024 // 10MB
)
// 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
database DatabaseMetrics
repository string
authorizer auth.HoldAuthorizer
holdDID string
}
// NewProxyBlobStore creates a new proxy blob store
func NewProxyBlobStore(storageEndpoint, did string, database DatabaseMetrics, repository string, authorizer auth.HoldAuthorizer) *ProxyBlobStore {
// Convert storage endpoint URL to did:web DID for authorization
holdDID := atproto.ResolveHoldDIDFromURL(storageEndpoint)
fmt.Printf("DEBUG [proxy_blob_store]: NewProxyBlobStore created with endpoint=%s, holdDID=%s, userDID=%s, repo=%s\n",
storageEndpoint, holdDID, did, repository)
return &ProxyBlobStore{
storageEndpoint: storageEndpoint,
httpClient: &http.Client{
Timeout: 5 * time.Minute, // Timeout for presigned URL requests and uploads
Transport: &http.Transport{
DisableKeepAlives: false, // Re-enable keep-alive
MaxIdleConns: 100,
MaxIdleConnsPerHost: 100,
MaxConnsPerHost: 0, // unlimited
IdleConnTimeout: 90 * time.Second,
},
},
did: did,
database: database,
repository: repository,
authorizer: authorizer,
holdDID: holdDID,
}
}
// checkReadAccess verifies the user has read access to the hold
func (p *ProxyBlobStore) checkReadAccess(ctx context.Context) error {
if p.authorizer == nil {
// No authorizer configured - allow access (backward compatibility)
return nil
}
hasAccess, err := p.authorizer.CheckReadAccess(ctx, p.holdDID, p.did)
if err != nil {
return fmt.Errorf("authorization check failed: %w", err)
}
if !hasAccess {
return distribution.ErrBlobUnknown // Return same error as missing blob for security
}
return nil
}
// checkWriteAccess verifies the user has write access to the hold
func (p *ProxyBlobStore) checkWriteAccess(ctx context.Context) error {
if p.authorizer == nil {
// No authorizer configured - allow access (backward compatibility)
return nil
}
hasAccess, err := p.authorizer.CheckWriteAccess(ctx, p.holdDID, p.did)
if err != nil {
return fmt.Errorf("authorization check failed: %w", err)
}
if !hasAccess {
return fmt.Errorf("write access denied to hold %s", p.holdDID)
}
return nil
}
// Stat returns the descriptor for a blob
func (p *ProxyBlobStore) Stat(ctx context.Context, dgst digest.Digest) (distribution.Descriptor, error) {
// Check read access
if err := p.checkReadAccess(ctx); err != nil {
return distribution.Descriptor{}, err
}
// Get presigned HEAD URL
url, err := p.getHeadURL(ctx, dgst)
if err != nil {
return distribution.Descriptor{}, distribution.ErrBlobUnknown
}
// Make HEAD request to presigned URL
req, err := http.NewRequestWithContext(ctx, "HEAD", url, nil)
if err != nil {
return distribution.Descriptor{}, distribution.ErrBlobUnknown
}
resp, err := p.httpClient.Do(req)
if err != nil {
return distribution.Descriptor{}, distribution.ErrBlobUnknown
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return distribution.Descriptor{}, distribution.ErrBlobUnknown
}
// Return a minimal descriptor with size from Content-Length if available
size := int64(0)
if contentLength := resp.Header.Get("Content-Length"); contentLength != "" {
fmt.Sscanf(contentLength, "%d", &size)
}
return distribution.Descriptor{
Digest: dgst,
Size: size,
MediaType: "application/octet-stream",
}, nil
}
// Get retrieves a blob
func (p *ProxyBlobStore) Get(ctx context.Context, dgst digest.Digest) ([]byte, error) {
// Check read access
if err := p.checkReadAccess(ctx); err != nil {
return nil, err
}
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) {
// Check read access
if err := p.checkReadAccess(ctx); err != nil {
return nil, err
}
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 using the multipart upload flow
// This ensures all uploads go through the same XRPC path
func (p *ProxyBlobStore) Put(ctx context.Context, mediaType string, content []byte) (distribution.Descriptor, error) {
// Check write access (fast-fail before starting multipart upload)
if err := p.checkWriteAccess(ctx); err != nil {
return distribution.Descriptor{}, err
}
// Calculate digest
dgst := digest.FromBytes(content)
// Use Create() flow for all uploads (goes through multipart XRPC endpoints)
writer, err := p.Create(ctx)
if err != nil {
fmt.Printf("[proxy_blob_store/Put] Failed to create writer: %v\n", err)
return distribution.Descriptor{}, err
}
// Write the content
if _, err := writer.Write(content); err != nil {
writer.Cancel(ctx)
fmt.Printf("[proxy_blob_store/Put] Failed to write content: %v\n", err)
return distribution.Descriptor{}, err
}
// Commit with the calculated digest
desc, err := writer.Commit(ctx, distribution.Descriptor{
Digest: dgst,
Size: int64(len(content)),
MediaType: mediaType,
})
if err != nil {
fmt.Printf("[proxy_blob_store/Put] Failed to commit: %v\n", err)
return distribution.Descriptor{}, err
}
fmt.Printf("[proxy_blob_store/Put] Upload successful: digest=%s, size=%d\n", dgst, len(content))
return desc, 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 {
// Check read access
if err := p.checkReadAccess(ctx); err != nil {
return err
}
// For HEAD requests, redirect to presigned HEAD URL
if r.Method == http.MethodHead {
url, err := p.getHeadURL(ctx, dgst)
if err != nil {
return err
}
// Redirect to presigned HEAD URL
http.Redirect(w, r, url, http.StatusTemporaryRedirect)
return nil
}
// For GET requests, redirect to presigned 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 using multipart upload
func (p *ProxyBlobStore) Create(ctx context.Context, options ...distribution.BlobCreateOption) (distribution.BlobWriter, error) {
// Check write access
if err := p.checkWriteAccess(ctx); err != nil {
return nil, err
}
// Parse options
var opts distribution.CreateOptions
for _, option := range options {
if err := option.Apply(&opts); err != nil {
return nil, err
}
}
// Generate unique writer ID
writerID := fmt.Sprintf("upload-%d", time.Now().UnixNano())
// Use temp digest for upload location (will be moved to final digest on commit)
tempDigest := fmt.Sprintf("uploads/temp-%s", writerID)
// Start multipart upload via hold service
uploadID, err := p.startMultipartUpload(ctx, tempDigest)
if err != nil {
return nil, fmt.Errorf("failed to start multipart upload: %w", err)
}
fmt.Printf(" Started multipart upload: uploadID=%s\n", uploadID)
writer := &ProxyBlobWriter{
store: p,
options: opts,
uploadID: uploadID,
parts: make([]CompletedPart, 0),
partNumber: 1,
buffer: bytes.NewBuffer(make([]byte, 0, maxChunkSize)),
id: writerID,
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
}
// Just return the writer - parts are buffered and flushed on demand
return writer, nil
}
// getDownloadURL returns the XRPC getBlob URL for downloading a blob
// The hold service will redirect to a presigned S3 URL
func (p *ProxyBlobStore) getDownloadURL(ctx context.Context, dgst digest.Digest) (string, error) {
// Use XRPC endpoint: GET /xrpc/com.atproto.sync.getBlob?did={holdDID}&cid={digest}
// Per migration doc: hold accepts OCI digest directly as cid parameter (checks for sha256: prefix)
url := fmt.Sprintf("%s/xrpc/com.atproto.sync.getBlob?did=%s&cid=%s",
p.storageEndpoint, p.holdDID, dgst.String())
return url, nil
}
// getHeadURL returns the XRPC getBlob URL for HEAD requests
// The hold service will redirect to a presigned S3 URL
func (p *ProxyBlobStore) getHeadURL(ctx context.Context, dgst digest.Digest) (string, error) {
// Same as GET - hold service handles HEAD method on getBlob endpoint
url := fmt.Sprintf("%s/xrpc/com.atproto.sync.getBlob?did=%s&cid=%s",
p.storageEndpoint, p.holdDID, dgst.String())
return url, nil
}
// getUploadURL is deprecated - single blob uploads should use Create() instead
// XRPC migration: No direct presigned upload URL endpoint, use multipart flow for all uploads
func (p *ProxyBlobStore) getUploadURL(ctx context.Context, dgst digest.Digest, size int64) (string, error) {
return "", fmt.Errorf("single blob upload via Put() not supported with XRPC endpoints - use Create() instead")
}
// startMultipartUpload initiates a multipart upload via XRPC uploadBlob endpoint
func (p *ProxyBlobStore) startMultipartUpload(ctx context.Context, digest string) (string, error) {
reqBody := map[string]any{
"action": "start",
"digest": digest,
}
body, err := json.Marshal(reqBody)
if err != nil {
return "", err
}
url := fmt.Sprintf("%s/xrpc/com.atproto.repo.uploadBlob", 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 {
bodyBytes, _ := io.ReadAll(resp.Body)
return "", fmt.Errorf("start multipart failed: status %d, body: %s", resp.StatusCode, string(bodyBytes))
}
var result struct {
UploadID string `json:"uploadId"`
Mode string `json:"mode"`
}
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
return "", err
}
return result.UploadID, nil
}
// PartUploadInfo contains structured information for uploading a part
type PartUploadInfo struct {
URL string `json:"url"`
Method string `json:"method,omitempty"`
Headers map[string]string `json:"headers,omitempty"`
}
// getPartUploadInfo gets structured upload info for uploading a specific part via XRPC
func (p *ProxyBlobStore) getPartUploadInfo(ctx context.Context, digest, uploadID string, partNumber int) (*PartUploadInfo, error) {
reqBody := map[string]any{
"action": "part",
"uploadId": uploadID,
"partNumber": partNumber,
"digest": digest,
}
body, err := json.Marshal(reqBody)
if err != nil {
return nil, err
}
url := fmt.Sprintf("%s/xrpc/com.atproto.repo.uploadBlob", p.storageEndpoint)
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
resp, err := p.httpClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
bodyBytes, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("get part URL failed: status %d, body: %s", resp.StatusCode, string(bodyBytes))
}
var uploadInfo PartUploadInfo
if err := json.NewDecoder(resp.Body).Decode(&uploadInfo); err != nil {
return nil, err
}
return &uploadInfo, nil
}
// completeMultipartUpload completes a multipart upload via XRPC uploadBlob endpoint
// The XRPC complete action handles the move from temp to final location internally
func (p *ProxyBlobStore) completeMultipartUpload(ctx context.Context, digest, uploadID string, parts []CompletedPart) error {
// Convert parts to XRPC format (partNumber instead of part_number)
xrpcParts := make([]map[string]any, len(parts))
for i, part := range parts {
xrpcParts[i] = map[string]any{
"partNumber": part.PartNumber,
"etag": part.ETag,
}
}
reqBody := map[string]any{
"action": "complete",
"uploadId": uploadID,
"digest": digest,
"parts": xrpcParts,
}
body, err := json.Marshal(reqBody)
if err != nil {
return err
}
url := fmt.Sprintf("%s/xrpc/com.atproto.repo.uploadBlob", 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 {
bodyBytes, _ := io.ReadAll(resp.Body)
return fmt.Errorf("complete multipart failed: status %d, body: %s", resp.StatusCode, string(bodyBytes))
}
return nil
}
// abortMultipartUpload aborts a multipart upload via XRPC uploadBlob endpoint
func (p *ProxyBlobStore) abortMultipartUpload(ctx context.Context, digest, uploadID string) error {
reqBody := map[string]any{
"action": "abort",
"uploadId": uploadID,
"digest": digest,
}
body, err := json.Marshal(reqBody)
if err != nil {
return err
}
url := fmt.Sprintf("%s/xrpc/com.atproto.repo.uploadBlob", 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 {
bodyBytes, _ := io.ReadAll(resp.Body)
return fmt.Errorf("abort multipart failed: status %d, body: %s", resp.StatusCode, string(bodyBytes))
}
return nil
}
// CompletedPart represents an uploaded part with its ETag
type CompletedPart struct {
PartNumber int `json:"part_number"`
ETag string `json:"etag"`
}
// ProxyBlobWriter implements distribution.BlobWriter for proxy uploads using multipart upload
type ProxyBlobWriter struct {
store *ProxyBlobStore
options distribution.CreateOptions
uploadID string // S3 multipart upload ID
parts []CompletedPart // Track uploaded parts with ETags
partNumber int // Current part number (starts at 1)
buffer *bytes.Buffer // Buffer for current part
size int64 // Total bytes written
closed bool
id string // Distribution's upload ID (for state)
startedAt time.Time
finalDigest string // Set on Commit
}
// 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
// Buffers data and flushes when buffer reaches 5MB
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)
// Flush if buffer reaches limit (S3 part size)
if w.buffer.Len() >= maxChunkSize {
if err := w.flushPart(); err != nil {
return n, err
}
}
return n, err
}
// flushPart uploads the current buffer as a part
func (w *ProxyBlobWriter) flushPart() error {
if w.buffer.Len() == 0 {
return nil
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
defer cancel()
// Get structured upload info for this part
tempDigest := fmt.Sprintf("uploads/temp-%s", w.id)
uploadInfo, err := w.store.getPartUploadInfo(ctx, tempDigest, w.uploadID, w.partNumber)
if err != nil {
return fmt.Errorf("failed to get part upload info: %w", err)
}
// Determine HTTP method (default to PUT)
method := uploadInfo.Method
if method == "" {
method = "PUT"
}
// Upload part (either to S3 presigned URL or back to XRPC with headers)
req, err := http.NewRequestWithContext(ctx, method, uploadInfo.URL, bytes.NewReader(w.buffer.Bytes()))
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/octet-stream")
// Apply any additional headers from the response (for buffered mode)
for key, value := range uploadInfo.Headers {
req.Header.Set(key, value)
}
resp, err := w.store.httpClient.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated {
bodyBytes, _ := io.ReadAll(resp.Body)
return fmt.Errorf("part upload failed: status %d, body: %s", resp.StatusCode, string(bodyBytes))
}
// Store ETag for completion
// For buffered mode, ETag might be in JSON response body
etag := resp.Header.Get("ETag")
if etag == "" {
// Try to parse JSON response for buffered mode
var result struct {
ETag string `json:"etag"`
}
if err := json.NewDecoder(resp.Body).Decode(&result); err == nil && result.ETag != "" {
etag = result.ETag
} else {
return fmt.Errorf("no ETag in response")
}
}
w.parts = append(w.parts, CompletedPart{
PartNumber: w.partNumber,
ETag: etag,
})
fmt.Printf("[flushPart] Part %d uploaded successfully: ETag=%s\n", w.partNumber, etag)
// Reset buffer and increment part number
w.buffer.Reset()
w.partNumber++
return nil
}
// ReadFrom reads from a reader
func (w *ProxyBlobWriter) ReadFrom(r io.Reader) (int64, error) {
if w.closed {
return 0, fmt.Errorf("writer closed")
}
// Read in chunks and flush when needed
buf := make([]byte, 32*1024) // 32KB read buffer
var total int64
for {
nr, err := r.Read(buf)
if nr > 0 {
nw, werr := w.Write(buf[:nr])
total += int64(nw)
if werr != nil {
return total, werr
}
}
if err == io.EOF {
break
}
if err != nil {
return total, err
}
}
return total, nil
}
// Size returns the current size
func (w *ProxyBlobWriter) Size() int64 {
return w.size
}
// Commit finalizes the upload by completing multipart upload and moving to final location
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()
// Flush any remaining buffered data
if w.buffer.Len() > 0 {
fmt.Printf("[Commit] Flushing final buffer: %d bytes\n", w.buffer.Len())
if err := w.flushPart(); err != nil {
// Try to abort multipart on error
tempDigest := fmt.Sprintf("uploads/temp-%s", w.id)
w.store.abortMultipartUpload(ctx, tempDigest, w.uploadID)
return distribution.Descriptor{}, fmt.Errorf("failed to flush final part: %w", err)
}
}
// Complete multipart upload - XRPC complete action handles move internally
tempDigest := fmt.Sprintf("uploads/temp-%s", w.id)
fmt.Printf("🔒 [Commit] Completing multipart upload: uploadID=%s, parts=%d\n", w.uploadID, len(w.parts))
if err := w.store.completeMultipartUpload(ctx, tempDigest, w.uploadID, w.parts); err != nil {
return distribution.Descriptor{}, fmt.Errorf("failed to complete multipart upload: %w", err)
}
fmt.Printf("[Commit] Upload completed successfully: digest=%s, size=%d, parts=%d\n", desc.Digest, w.size, len(w.parts))
return distribution.Descriptor{
Digest: desc.Digest,
Size: w.size,
MediaType: desc.MediaType,
}, nil
}
// Cancel cancels the upload by aborting the multipart upload
func (w *ProxyBlobWriter) Cancel(ctx context.Context) error {
w.closed = true
fmt.Printf("[Cancel] Cancelling upload: id=%s\n", w.id)
// Remove from global uploads map
globalUploadsMu.Lock()
delete(globalUploads, w.id)
globalUploadsMu.Unlock()
// Abort multipart upload
tempDigest := fmt.Sprintf("uploads/temp-%s", w.id)
if err := w.store.abortMultipartUpload(ctx, tempDigest, w.uploadID); err != nil {
fmt.Printf("⚠️ [Cancel] Failed to abort multipart upload: %v\n", err)
// Continue anyway - we want to mark upload as cancelled
}
fmt.Printf("[Cancel] Upload cancelled: id=%s\n", w.id)
return nil
}
// Close closes the writer
// Parts are flushed on demand, so this is a no-op
func (w *ProxyBlobWriter) Close() error {
// Don't set w.closed = true - allow resuming for next PATCH
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")
}