Files
at-container-registry/pkg/appview/storage/proxy_blob_store.go
T
Evan JarrettandClaude Fable 5.1 f4343d7956 appview: upload small blobs with one presigned PUT, and verify every digest
Every blob went through the multipart machinery: an S3 multipart started
on Docker's initial POST, a hold round trip per part, and a complete on
the hold that finished the multipart, HEADed the temp object, copied it
to its final key, and deleted the temp. For a 2KB config blob that was
three hold calls and six S3 operations. On production data 86% of
distinct layers and every config blob fit in a 16MB buffer, and 49% of
image manifests have no layer larger than that.

The writer now buffers up to 16MB (also the multipart part size) and
makes no hold call until it has to. A blob that never overflows the
buffer is written at Commit with a single presigned PUT to its final
key, via the hold's existing method=PUT presign; the multipart only
starts on the first flush. The hold's completeUpload does nothing the
direct path skips: quota, layer records, stats and scan dispatch all
hang off notifyManifest, which is unchanged.

The buffer starts empty and grows on demand, with the doubling capped so
capacity never overshoots 16MB: a config blob costs kilobytes, and only
layers that approach the threshold fill it.

Bytes are hashed as they arrive. Commit compares the computed sha256 to
the digest the client claimed before any network call, and returns
DIGEST_INVALID on mismatch, aborting a multipart if one was started.
Previously nothing verified the content, so a pusher could store wrong
bytes under a digest in the shared content-addressed space.

Tests observe request counts on a fake hold and fake S3 rather than
return values. The growth test streams in 24KB chunks because
power-of-two chunks land on 16MB by luck and hid an earlier weaker guard.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018Yf1ZVA7sXYhQNb9tCo1m5
2026-09-09 09:50:08 -05:00

1021 lines
34 KiB
Go

package storage
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
"strconv"
"sync"
"time"
"atcr.io/pkg/atproto"
"github.com/distribution/distribution/v3"
"github.com/distribution/distribution/v3/registry/api/errcode"
"github.com/opencontainers/go-digest"
)
const (
// maxBufferSize is the writer's in-memory buffer limit, and it plays two
// roles. It is the S3 multipart part size for blobs big enough to need
// multipart at all, and it is the cutoff below which a blob never touches
// the multipart machinery: everything still buffered at Commit goes to its
// final key with a single presigned PUT.
//
// 16MB rather than 10MB because of what the production data says: at 16MB,
// 86% of distinct layers and every config blob fit entirely in the buffer,
// so for the overwhelming majority of blobs the multipart path (3 hold
// calls and 6 S3 operations, one of them a full server side copy) was pure
// overhead. S3's 5MB multipart minimum is still satisfied for the parts of
// uploads that do go multipart.
maxBufferSize = 16 * 1024 * 1024 // 16MB
)
// 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
)
// The transport and client below are package-level on purpose. RoutingRepository
// (and therefore ProxyBlobStore) is built fresh on every registry request, so a
// per-instance transport was thrown away after a single request: nothing was ever
// reused, every XRPC call to the hold and every presigned S3 request paid for a
// fresh TCP + TLS handshake, the MaxIdleConns settings below were dead config, and
// each discarded transport still sat on its idle sockets for the full
// IdleConnTimeout. Sharing one transport process-wide is what makes the idle pool
// and keep-alives mean anything.
//
// ForceAttemptHTTP2 is explicit to document intent; no custom DialContext or
// TLSClientConfig is set, so Go's automatic HTTP/2 negotiation over ALPN stays on
// (see 416ba4a, where the load balancer started speaking HTTP/2 on the frontend).
// The LB talks HTTP/1.1 to the hold backend, so multiplexing stops at the LB. The
// win here is not multiplexing: it is removing the handshake and socket churn.
var (
sharedTransport = &http.Transport{
DisableKeepAlives: false,
MaxIdleConns: 100,
MaxIdleConnsPerHost: 100,
MaxConnsPerHost: 0, // unlimited
IdleConnTimeout: 90 * time.Second,
ForceAttemptHTTP2: true,
}
sharedHTTPClient = &http.Client{
Timeout: 5 * time.Minute, // Timeout for presigned URL requests and uploads
Transport: sharedTransport,
}
)
// ProxyBlobStore proxies blob requests to an external storage service
type ProxyBlobStore struct {
ctx *RegistryContext // All context and services
holdURL string // Resolved HTTP URL for XRPC requests
httpClient *http.Client
}
// NewProxyBlobStore creates a new proxy blob store
func NewProxyBlobStore(ctx *RegistryContext) *ProxyBlobStore {
// Use pre-resolved URL from RegistryContext (resolved in Registry.Repository())
holdURL := ctx.HoldURL
slog.Debug("NewProxyBlobStore created", "component", "proxy_blob_store", "hold_did", ctx.HoldDID, "hold_url", holdURL, "user_did", ctx.DID, "repo", ctx.Repository)
return &ProxyBlobStore{
ctx: ctx,
holdURL: holdURL,
// Field stays per-instance so tests can substitute a client; the default
// points at the process-wide client so connections are actually reused.
httpClient: sharedHTTPClient,
}
}
// doAuthenticatedRequest performs an HTTP request to the hold service, attaching
// the service token when one is present. An empty service token means an
// anonymous pull: the request is sent without an Authorization header and the
// hold authorizes it per captain.Public. Write call sites (multipart upload) are
// push-only and always carry a service token, so they never go out unauthenticated.
func (p *ProxyBlobStore) doAuthenticatedRequest(ctx context.Context, req *http.Request) (*http.Response, error) {
// Service token was validated and cached by middleware (which fails fast with
// HTTP 401 if the OAuth session is invalid). Anonymous reads have none.
if p.ctx.ServiceToken != "" {
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", p.ctx.ServiceToken))
}
return p.httpClient.Do(req)
}
// checkReadAccess validates that the user has read access to blobs in this hold
func (p *ProxyBlobStore) checkReadAccess(ctx context.Context) error {
if p.ctx.Authorizer == nil {
return nil // No authorization check if authorizer not configured
}
// Authorize the *requester*, not the repository owner. p.ctx.DID is the
// owner whose namespace is being read; passing it here asked "may the owner
// read their own hold", which is true for every private hold (any non-empty
// DID satisfies CheckReadAccessWithCaptain), so an anonymous request sailed
// through this gate and the Anonymous branch below was unreachable. An
// anonymous request has no identity, so it must be judged as one: only
// captain.Public can admit it.
requesterDID := p.ctx.DID
if p.ctx.Anonymous {
requesterDID = ""
}
allowed, err := p.ctx.Authorizer.CheckReadAccess(ctx, p.ctx.HoldDID, requesterDID)
if err != nil {
return fmt.Errorf("authorization check failed: %w", err)
}
if !allowed {
if p.ctx.Anonymous {
// Anonymous request to a private hold: surface a 401 so the Docker
// client prompts for credentials rather than treating it as a hard
// 403. The BearerChallenge middleware attaches WWW-Authenticate.
return errcode.ErrorCodeUnauthorized.WithMessage("authentication required")
}
// Authenticated but unauthorized: 403 Forbidden instead of masquerading
// as a missing blob, and without bouncing the user back to re-auth.
return errcode.ErrorCodeDenied.WithMessage("read access denied")
}
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
}
method := "HEAD"
blob, err := p.getPresignedURL(ctx, method, dgst)
if err != nil {
// Preserve an authorization verdict. distribution calls Stat before
// ServeBlob on both GET and HEAD, so flattening everything to
// ErrBlobUnknown here turns the hold's "private, authenticate first" into
// a 404 and leaves BearerChallenge with no 401 to annotate — the client
// is told the blob doesn't exist instead of being asked for credentials.
var ecErr errcode.Error
if errors.As(err, &ecErr) {
return distribution.Descriptor{}, err
}
return distribution.Descriptor{}, distribution.ErrBlobUnknown
}
if blob.Size != nil {
// The hold reported the size, so the descriptor is complete and the
// blob's bytes never have to be touched. This is the whole point of the
// field: Stat is called before every GET and HEAD of a blob, and the
// round trip below was buying nothing but Content-Length.
return distribution.Descriptor{
Digest: dgst,
Size: *blob.Size,
MediaType: "application/octet-stream",
}, nil
}
// No size in the response: the hold predates the field. Fall back to the
// original behaviour and read Content-Length off the presigned URL, so a
// new AppView keeps working against a hold that has not been upgraded.
req, err := http.NewRequestWithContext(ctx, method, blob.URL, nil)
if err != nil {
return distribution.Descriptor{}, distribution.ErrBlobUnknown
}
// Go directly to the presigned URL, no need to authenticate
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 != "" {
if parsed, err := strconv.ParseInt(contentLength, 10, 64); err == nil {
size = parsed
}
}
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
}
method := "GET"
blob, err := p.getPresignedURL(ctx, method, dgst)
if err != nil {
return nil, err
}
// Download the blob from presigned URL
req, err := http.NewRequestWithContext(ctx, method, blob.URL, nil)
if err != nil {
return nil, err
}
// Go directly to the presigned URL, no need to authenticate
resp, err := p.httpClient.Do(req)
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
}
method := "GET"
blob, err := p.getPresignedURL(ctx, method, dgst)
if err != nil {
return nil, err
}
// Download the blob from presigned URL
req, err := http.NewRequestWithContext(ctx, method, blob.URL, nil)
if err != nil {
return nil, err
}
// Go directly to the presigned URL, no need to authenticate
resp, err := p.httpClient.Do(req)
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 through the Create/Write/Commit writer, so it takes the
// same path a client push does: a single direct PUT for anything under
// maxBufferSize (which is every blob the AppView puts itself), multipart above
// it. Routing it through the writer also means Put's content is digest-verified
// by the same check.
//
// Write authorization is gated at /auth/token (pkg/appview/authgate); the
// JWT carries the resolved authorization for its lifetime. Hold-side
// requireBlobWriteAccess is the final defense (validates the service-token
// audience). No re-check needed here.
func (p *ProxyBlobStore) Put(ctx context.Context, mediaType string, content []byte) (distribution.Descriptor, error) {
// Calculate digest
dgst := digest.FromBytes(content)
// Use the Create() flow for all uploads so every blob takes one code path
writer, err := p.Create(ctx)
if err != nil {
slog.Error("Failed to create writer", "component", "proxy_blob_store/Put", "error", err)
return distribution.Descriptor{}, err
}
// Write the content
if _, err := writer.Write(content); err != nil {
writer.Cancel(ctx)
slog.Error("Failed to write content", "component", "proxy_blob_store/Put", "error", 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 {
slog.Error("Failed to commit", "component", "proxy_blob_store/Put", "error", err)
return distribution.Descriptor{}, err
}
slog.Debug("Upload successful", "component", "proxy_blob_store/Put", "digest", dgst, "size", len(content))
return desc, nil
}
// Delete removes a blob.
//
// Blob deletion is not offered on the client-facing OCI path: layer bytes live
// in the hold's S3 and are reclaimed by the hold's reference-counted GC once no
// manifest references them (see PurgeOnHold on manifest delete). Returning the
// distribution.ErrUnsupported sentinel makes the (always-registered) blob DELETE
// route respond with a clean OCI UNSUPPORTED error instead of a generic 500.
func (p *ProxyBlobStore) Delete(ctx context.Context, dgst digest.Digest) error {
return distribution.ErrUnsupported
}
// ServeBlob serves a blob via HTTP redirect or proxied response
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
}
blob, err := p.getPresignedURL(ctx, r.Method, dgst)
if err != nil {
return err
}
// Redirect to presigned URL
http.Redirect(w, r, blob.URL, http.StatusTemporaryRedirect)
return nil
}
// Create returns a blob writer for uploading.
//
// No hold call is made here. Docker opens an upload with a POST before it
// knows anything about the blob, and most blobs turn out to fit entirely in
// the writer's buffer, so starting an S3 multipart upload at this point meant
// opening (and then moving and deleting) a temp object for uploads that never
// needed one. The multipart upload is started lazily, on the first flush.
//
// The buffer starts empty and grows on demand for the same reason: a 2KB image
// config should not reserve 16MB.
//
// Write authorization is gated at /auth/token; see ProxyBlobStore.Put for
// the rationale on why we don't re-check here.
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
}
}
// Generate unique writer ID
writerID := fmt.Sprintf("upload-%d", time.Now().UnixNano())
writer := &ProxyBlobWriter{
store: p,
options: opts,
parts: make([]CompletedPart, 0),
partNumber: 1,
buffer: &bytes.Buffer{},
digester: digest.Canonical.Digester(),
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
}
// presignedBlob is the hold's answer to a getBlob presign request.
type presignedBlob struct {
// URL is the presigned S3 URL for the requested operation.
URL string
// Size is the blob's byte size as reported by the hold, or nil when the
// hold did not report one. A pointer rather than a plain int64 so that
// "the hold said nothing" stays distinguishable from "the hold said zero":
// a hold older than the size field reports nothing, and callers must fall
// back rather than believe in a zero-length blob.
Size *int64
}
// getPresignedURL asks the hold for a presigned URL for a blob operation, and
// for reads gets the blob's size back with it.
func (p *ProxyBlobStore) getPresignedURL(ctx context.Context, operation string, dgst digest.Digest) (presignedBlob, error) {
// Use XRPC endpoint: /xrpc/com.atproto.sync.getBlob?did={userDID}&cid={digest}
// The 'did' parameter is the USER's DID (whose blob we're fetching), not the hold service DID
// Per migration doc: hold accepts OCI digest directly as cid parameter (checks for sha256: prefix)
xrpcURL := fmt.Sprintf("%s%s?did=%s&cid=%s&method=%s",
p.holdURL, atproto.SyncGetBlob, p.ctx.DID, dgst.String(), operation)
req, err := http.NewRequestWithContext(ctx, "GET", xrpcURL, nil)
if err != nil {
return presignedBlob{}, fmt.Errorf("failed to create request: %w", err)
}
resp, err := p.doAuthenticatedRequest(ctx, req)
if err != nil {
// Don't wrap errcode errors - return them directly
if _, ok := err.(errcode.Error); ok {
return presignedBlob{}, err
}
return presignedBlob{}, fmt.Errorf("failed to get presigned URL: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode == http.StatusForbidden && p.ctx.Anonymous {
// Stale local captain cache let an anonymous request through, but the
// hold says private. Surface a 401 so the client re-authenticates.
return presignedBlob{}, errcode.ErrorCodeUnauthorized.WithMessage("authentication required")
}
if resp.StatusCode == http.StatusNotFound {
// The hold checked its storage and the blob is not there. Return the
// sentinel rather than a generic failure: Stat passes it through as
// blob-unknown, and Get and Open hand their callers the error the
// distribution interface documents for a missing blob.
return presignedBlob{}, distribution.ErrBlobUnknown
}
if resp.StatusCode != http.StatusOK {
bodyBytes, _ := io.ReadAll(resp.Body)
return presignedBlob{}, fmt.Errorf("hold service returned error: status %d, body: %s", resp.StatusCode, string(bodyBytes))
}
// Parse JSON response to get the presigned URL, and the size when the hold
// reports one. Size is a pointer so an older hold, which sends no size
// field at all, is not read as a zero-length blob.
var result struct {
URL string `json:"url"`
Size *int64 `json:"size"`
}
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
return presignedBlob{}, fmt.Errorf("failed to parse hold service response: %w", err)
}
if result.URL == "" {
return presignedBlob{}, fmt.Errorf("hold service returned empty URL")
}
slog.Debug("Got presigned URL from hold service", "component", "proxy_blob_store", "url", result.URL, "size_reported", result.Size != nil)
return presignedBlob{URL: result.URL, Size: result.Size}, nil
}
// startMultipartUpload initiates a multipart upload via XRPC initiateUpload endpoint
func (p *ProxyBlobStore) startMultipartUpload(ctx context.Context, digest string) (string, error) {
reqBody := map[string]any{
"digest": digest,
}
body, err := json.Marshal(reqBody)
if err != nil {
return "", err
}
url := fmt.Sprintf("%s%s", p.holdURL, atproto.HoldInitiateUpload)
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
if err != nil {
return "", err
}
req.Header.Set("Content-Type", "application/json")
// Use authenticated request (OAuth with DPoP)
resp, err := p.doAuthenticatedRequest(ctx, 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"`
}
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
return "", err
}
return result.UploadID, nil
}
// PartUploadInfo contains the presigned URL for uploading a part
type PartUploadInfo struct {
URL string `json:"url"` // Presigned URL to PUT the part to
}
// 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{
"uploadId": uploadID,
"partNumber": partNumber,
}
body, err := json.Marshal(reqBody)
if err != nil {
return nil, err
}
url := fmt.Sprintf("%s%s", p.holdURL, atproto.HoldGetPartUploadURL)
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
// Use authenticated request (OAuth with DPoP)
resp, err := p.doAuthenticatedRequest(ctx, 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 completeUpload 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
xrpcParts := make([]map[string]any, len(parts))
for i, part := range parts {
xrpcParts[i] = map[string]any{
"part_number": part.PartNumber,
"etag": part.ETag,
}
}
reqBody := map[string]any{
"uploadId": uploadID,
"digest": digest,
"parts": xrpcParts,
}
body, err := json.Marshal(reqBody)
if err != nil {
return err
}
url := fmt.Sprintf("%s%s", p.holdURL, atproto.HoldCompleteUpload)
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/json")
// Use authenticated request (OAuth with DPoP)
resp, err := p.doAuthenticatedRequest(ctx, 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 abortUpload endpoint
func (p *ProxyBlobStore) abortMultipartUpload(ctx context.Context, uploadID string) error {
reqBody := map[string]any{
"uploadId": uploadID,
}
body, err := json.Marshal(reqBody)
if err != nil {
return err
}
url := fmt.Sprintf("%s%s", p.holdURL, atproto.HoldAbortUpload)
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/json")
// Use authenticated request (OAuth with DPoP)
resp, err := p.doAuthenticatedRequest(ctx, 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.
//
// Small blobs (everything that still fits in buffer at Commit) are PUT once to
// their final content-addressed key. Larger ones fall back to an S3 multipart
// upload, started on the first flush.
type ProxyBlobWriter struct {
store *ProxyBlobStore
options distribution.CreateOptions
uploadID string // S3 multipart upload ID; empty until the first flush starts one
parts []CompletedPart // Track uploaded parts with ETags
partNumber int // Current part number (starts at 1)
buffer *bytes.Buffer // Buffer for current part
digester digest.Digester // Hashes every byte written, for verification at Commit
size int64 // Total bytes written
closed bool
id string // Distribution's upload ID (for state)
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.
// Buffers data and flushes a part once the buffer reaches maxBufferSize.
func (w *ProxyBlobWriter) Write(p []byte) (int, error) {
if w.closed {
return 0, fmt.Errorf("writer closed")
}
w.growBuffer(len(p))
n, err := w.buffer.Write(p)
w.size += int64(n)
if n > 0 {
// Hash as we go. Nothing else in this path ever looked at the bytes:
// Commit took the digest in the client's final PUT on trust, which made
// the content address of a blob whatever the client claimed it was.
w.digester.Hash().Write(p[:n])
}
// Flush if buffer reaches limit (S3 part size)
if w.buffer.Len() >= maxBufferSize {
if err := w.flushPart(); err != nil {
return n, err
}
}
return n, err
}
// growBuffer sizes the buffer's backing array ahead of an n byte write so that
// bytes.Buffer's doubling never overshoots maxBufferSize.
//
// Doubling is the right strategy while the buffer is small: a config blob ends
// up with a few KB of backing array instead of the full 16MB the writer used
// to reserve up front. But bytes.Buffer grows to max(needed, 2*cap), so a
// capacity that is anywhere past half the threshold doubles clean past it, and
// everything above maxBufferSize is wasted: the buffer is flushed and reset the
// moment it reaches the threshold. It is not a rounding error either. A layer
// streamed in 24KB chunks walks its capacity to 12MB and then doubles to 24MB,
// half of which is never used.
//
// So doubling is allowed only while the capacity it would land on still leaves
// room to double again. Past that, grow to exactly the threshold and stop,
// which is safe because the capacity at that point is at most half of it.
//
// After a flush, Reset keeps the capacity, so a large upload allocates its
// 16MB once and reuses it for every part.
func (w *ProxyBlobWriter) growBuffer(n int) {
c := w.buffer.Cap()
if c >= maxBufferSize {
return // Already at full size, nothing to do
}
// What bytes.Buffer would grow to on its own if this write does not fit.
projected := max(w.buffer.Len()+n, 2*c)
if projected <= maxBufferSize/2 {
return // Still room for another doubling afterwards
}
w.buffer.Grow(maxBufferSize - w.buffer.Len())
}
// 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()
tempDigest := fmt.Sprintf("uploads/temp-%s", w.id)
// Start the multipart upload on the first flush rather than in Create. A
// blob that never fills the buffer is committed with a single direct PUT
// and needs no multipart session, no temp object and no server side copy.
if w.uploadID == "" {
uploadID, err := w.store.startMultipartUpload(ctx, tempDigest)
if err != nil {
return fmt.Errorf("failed to start multipart upload: %w", err)
}
w.uploadID = uploadID
}
// Get structured upload info for this part
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)
}
// Upload part to S3 presigned URL
req, err := http.NewRequestWithContext(ctx, "PUT", uploadInfo.URL, bytes.NewReader(w.buffer.Bytes()))
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/octet-stream")
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,
})
slog.Debug("Part uploaded successfully", "component", "proxy_blob_store/flushPart", "part_number", w.partNumber, "etag", 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.
//
// The digest the client sent is verified against the bytes actually received
// before anything else happens, then the blob is finalized: a direct PUT to
// the final key if it is all still buffered, otherwise a final part plus the
// hold's completeUpload.
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()
// Verify before any network call, so a bad digest costs nothing and lands
// nothing in storage. The digest is the blob's address in a shared,
// content-addressed bucket, so a client that names its bytes wrongly would
// otherwise overwrite or shadow someone else's layer.
if desc.Digest.Algorithm() != digest.Canonical {
w.abortIfStarted(ctx)
slog.Warn("Rejected blob with unsupported digest algorithm", "component", "proxy_blob_store/Commit", "algorithm", desc.Digest.Algorithm(), "id", w.id)
return distribution.Descriptor{}, distribution.ErrBlobDigestUnsupported
}
if computed := w.digester.Digest(); computed != desc.Digest {
w.abortIfStarted(ctx)
slog.Warn("Rejected blob whose content does not match its digest", "component", "proxy_blob_store/Commit", "claimed", desc.Digest, "computed", computed, "size", w.size)
return distribution.Descriptor{}, distribution.ErrBlobInvalidDigest{
Digest: desc.Digest,
Reason: fmt.Errorf("content digest is %s", computed),
}
}
// Nothing was ever flushed, so the whole blob is in memory and can go
// straight to its final content-addressed key. This is the common case:
// every config blob and the large majority of layers land here.
if w.uploadID == "" {
if err := w.putDirect(ctx, desc.Digest); err != nil {
return distribution.Descriptor{}, err
}
slog.Info("Upload completed successfully", "component", "proxy_blob_store/Commit", "digest", desc.Digest, "size", w.size, "mode", "direct")
return distribution.Descriptor{
Digest: desc.Digest,
Size: w.size,
MediaType: desc.MediaType,
}, nil
}
// Flush any remaining buffered data
if w.buffer.Len() > 0 {
slog.Debug("Flushing final buffer", "component", "proxy_blob_store/Commit", "bytes", w.buffer.Len())
if err := w.flushPart(); err != nil {
// Try to abort multipart on error
w.abortIfStarted(ctx)
return distribution.Descriptor{}, fmt.Errorf("failed to flush final part: %w", err)
}
}
// Complete multipart upload - XRPC complete action handles move internally
// Send the real digest (not tempDigest) so hold can move temp → final location
slog.Info("Completing multipart upload", "component", "proxy_blob_store/Commit", "upload_id", w.uploadID, "parts", len(w.parts), "digest", desc.Digest)
if err := w.store.completeMultipartUpload(ctx, desc.Digest.String(), w.uploadID, w.parts); err != nil {
return distribution.Descriptor{}, fmt.Errorf("failed to complete multipart upload: %w", err)
}
slog.Info("Upload completed successfully", "component", "proxy_blob_store/Commit", "digest", desc.Digest, "size", w.size, "parts", len(w.parts), "mode", "multipart")
return distribution.Descriptor{
Digest: desc.Digest,
Size: w.size,
MediaType: desc.MediaType,
}, nil
}
// putDirect uploads the fully buffered blob to its final content-addressed key
// with a single presigned PUT, skipping the multipart dance entirely.
func (w *ProxyBlobWriter) putDirect(ctx context.Context, dgst digest.Digest) error {
// Same hold endpoint the read path uses, asked for a write capability.
// The hold gates method=PUT on blob write access and skips its size lookup,
// since the object being presigned does not exist yet.
blob, err := w.store.getPresignedURL(ctx, http.MethodPut, dgst)
if err != nil {
return fmt.Errorf("failed to get presigned upload URL: %w", err)
}
body := w.buffer.Bytes()
req, err := http.NewRequestWithContext(ctx, http.MethodPut, blob.URL, bytes.NewReader(body))
if err != nil {
return err
}
// The hold signs the PUT with ContentType "application/octet-stream"
// (GetPresignedURL in pkg/hold/pds/xrpc.go), and a signed header that the
// request does not carry fails S3's signature check. Nothing else is baked
// into the signature, in particular no content length.
req.Header.Set("Content-Type", "application/octet-stream")
req.ContentLength = int64(len(body))
resp, err := w.store.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to upload blob: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated {
bodyBytes, _ := io.ReadAll(resp.Body)
return fmt.Errorf("blob upload failed: status %d, body: %s", resp.StatusCode, string(bodyBytes))
}
slog.Debug("Blob uploaded directly to final location", "component", "proxy_blob_store/putDirect", "digest", dgst, "size", len(body))
return nil
}
// abortIfStarted aborts the multipart upload, if one was ever started. A writer
// whose blob stayed inside the buffer has no session to abort.
func (w *ProxyBlobWriter) abortIfStarted(ctx context.Context) {
if w.uploadID == "" {
return
}
if err := w.store.abortMultipartUpload(ctx, w.uploadID); err != nil {
slog.Warn("Failed to abort multipart upload", "component", "proxy_blob_store", "error", err)
// Continue anyway - we want to mark upload as cancelled
}
}
// Cancel cancels the upload by aborting the multipart upload
func (w *ProxyBlobWriter) Cancel(ctx context.Context) error {
w.closed = true
slog.Debug("Cancelling upload", "component", "proxy_blob_store/Cancel", "id", w.id)
// Remove from global uploads map
globalUploadsMu.Lock()
delete(globalUploads, w.id)
globalUploadsMu.Unlock()
w.abortIfStarted(ctx)
slog.Debug("Upload cancelled", "component", "proxy_blob_store/Cancel", "id", 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")
}