Files
at-container-registry/pkg/appview/storage/proxy_blob_store.go
T
Evan JarrettandClaude Fable 5.1 9dbc53b670 appview: stop serving service tokens past their expiry, and challenge the client when the hold rejects one
Seen in production on 2026-09-11: three cold pulls of a 22-layer image failed
with BLOB_UNKNOWN for layers that exist. The hold had answered 403 "service
token authentication failed: token has expired", and the same blobs served
fine a minute later.

Three things lined up. The registry middleware's validation cache kept a
fetched service token for a flat 45 seconds regardless of its real remaining
life, so a token fetched with 12 seconds left was still handed to the hold
half a minute after it died. The registry JWT is stamped from the auth cache's
expiry, which trailed the real exp by only 10 seconds, while distribution
accepts a JWT for 60 seconds past its exp, so a client could hold an accepted
JWT for most of a minute after the credential behind it was gone. And the
hold's 403 was flattened to BLOB_UNKNOWN, so the client failed instead of
re-authenticating.

Now the validation cache bounds an entry by the token's exp minus a shared
ServiceTokenSafetyMargin of 60 seconds, the same margin the auth cache and the
JWT stamp use, chosen to equal distribution's leeway so the last instant a JWT
is accepted is the service token's real exp. A PDS that grants less than the
margin gets half its remaining life instead of an already-past deadline. When
the hold rejects the service token as expired or missing, the appview drops
both cached copies and returns a 401 challenge so Docker and crane re-run the
token dance and retry; a genuine permission denial stays a 403, and a hold
that is down still maps to blob unknown.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EvFJr4Dwz8p2NDAeXmgmBt
2026-09-11 19:26:51 -05:00

1868 lines
69 KiB
Go

package storage
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
"strconv"
"strings"
"sync"
"time"
"atcr.io/pkg/atproto"
"github.com/distribution/distribution/v3"
"github.com/distribution/distribution/v3/registry/api/errcode"
"github.com/google/uuid"
"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
)
// newWriterID names a blob upload. A variable so tests can force a collision
// against the guard in Create.
//
// This used to be fmt.Sprintf("upload-%d", time.Now().UnixNano()), which is not
// unique: Go's wall clock is coarser than the spacing between goroutines, so
// two uploads opened at the same instant read the same nanosecond often enough
// to matter (roughly one release in five, measured on a tsc clocksource box).
// Docker and crane POST the config blob and several layers concurrently, so
// every multi blob push rolled those dice. The loser's writer replaced the
// winner's in globalUploads, both clients' PATCHes then resumed the same
// writer, and distribution rejected the second with a 416 "upload resumed at
// wrong offset", which the client reports as RANGE_INVALID.
//
// The UUID is hex and hyphens, so the ID stays URL safe: it travels in the
// upload URL and inside distribution's _state token.
var newWriterID = func() string {
return "upload-" + uuid.NewString()
}
// 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
// presigns memoizes the hold's presign answers for the life of this store,
// which is the life of one registry request (RoutingRepository.Blobs builds
// the store once per request under a sync.Once, so Stat and ServeBlob for
// the same HTTP request share this instance).
//
// One blob GET used to cost three identical round trips to the hold.
// distribution's blob handler calls Stat and then ServeBlob, and the
// notifications listener — installed unconditionally — calls Stat a third
// time after ServeBlob returns to build its BlobPulled event. A 22 layer
// pull was 66 com.atproto.sync.getBlob calls for 22 redirects.
//
// Keyed by method as well as digest because S3 SigV4 signs the HTTP verb: a
// URL presigned for HEAD answers 403 to a GET. Stat therefore presigns for
// the verb the client is actually going to send (see Stat), which is what
// makes ServeBlob's lookup a hit rather than a second call.
//
// Only successful presigns are remembered, so nothing here can change the
// error semantics of a later call.
presignMu sync.Mutex
presigns map[presignKey]presignedBlob
}
// presignKey identifies a memoized presign. The method is part of the identity,
// not an attribute: presigned URLs are verb-bound.
type presignKey struct {
digest digest.Digest
method string
}
// 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
}
// Presign for the verb the client is about to send, not always for HEAD.
// distribution calls Stat immediately before ServeBlob, so on a blob GET the
// two want the same URL — but only if they asked for the same verb, because
// SigV4 signs it. Asking for GET here is what lets ServeBlob be free.
//
// Anything that is not a read (a push's existence check arrives as a PUT,
// POST or PATCH, and manifest verification stats blobs mid-PUT) falls back
// to HEAD. That is not just a default: the hold refuses method=PUT on the
// read path unless the caller has write access, and skips the size lookup
// for it, so Stat must never ask for a write verb.
method := http.MethodHead
if m, ok := ctx.Value(HTTPRequestMethod).(string); ok && (m == http.MethodGet || m == http.MethodHead) {
method = m
}
blob, err := p.presign(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.
//
// The probe is always a HEAD, and when the URL above was signed for GET it
// cannot be used for it: sending the HEAD to a GET-signed URL is a 403, and
// sending a GET instead would download the entire blob just to read a
// header. So ask the hold for a HEAD-signed URL — the one extra call is
// paid only by the old-hold path, and only for the first Stat of a digest.
probe := blob
if method != http.MethodHead {
probe, err = p.presign(ctx, http.MethodHead, dgst)
if err != nil {
return distribution.Descriptor{}, distribution.ErrBlobUnknown
}
if probe.Size != nil {
return distribution.Descriptor{
Digest: dgst,
Size: *probe.Size,
MediaType: "application/octet-stream",
}, nil
}
}
req, err := http.NewRequestWithContext(ctx, http.MethodHead, probe.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"
// Through the memo: on a blob GET, Stat already presigned for GET, so this
// is usually free too.
blob, err := p.presign(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"
// Through the memo: on a blob GET, Stat already presigned for GET, so this
// is usually free too.
blob, err := p.presign(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
}
// Normally free: Stat ran moments ago on this same store and presigned for
// this same verb, so this is a memo hit and the redirect costs no hold call.
blob, err := p.presign(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
}
}
writerID := newWriterID()
now := time.Now()
writer := &ProxyBlobWriter{
store: p,
options: opts,
parts: make([]CompletedPart, 0),
partNumber: 1,
buffer: &bytes.Buffer{},
digester: digest.Canonical.Digester(),
id: writerID,
startedAt: now,
lastActivity: now,
requestCtx: ctx,
}
// Store in global uploads map for resume support.
//
// Refuse to overwrite: with a UUID an occupied key cannot be chance, so
// something is badly wrong and silently dropping the sitting writer would
// be worse than failing this Create. The sitting writer keeps its buffer
// budget and its hold side multipart, and its own client keeps resuming it.
globalUploadsMu.Lock()
if _, taken := globalUploads[writer.id]; taken {
globalUploadsMu.Unlock()
return nil, fmt.Errorf("upload id %s is already in use", writer.id)
}
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 {
// Also what a client sees after the sweeper reaped an abandoned upload:
// distribution turns this into BLOB_UPLOAD_UNKNOWN and the client
// starts the layer over.
return nil, distribution.ErrBlobUploadUnknown
}
// This request now owns the writer: hand it the context to block on, and
// count the resume as activity so a long push made of many PATCHes is never
// mistaken for an abandoned one.
writer.adopt(ctx)
// 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
}
// presign is getPresignedURL through the per-request memo: the first caller for
// a (digest, method) pair pays the hold round trip and every later caller in the
// same request is answered from memory. See the presigns field for why one
// request asks the same question up to three times.
//
// Failures are not remembered. A presign fails for reasons that belong to the
// call that made it (authorization, a missing blob, a hold that is down), and
// caching those would mean one call's verdict silently deciding another's.
func (p *ProxyBlobStore) presign(ctx context.Context, method string, dgst digest.Digest) (presignedBlob, error) {
key := presignKey{digest: dgst, method: method}
p.presignMu.Lock()
blob, hit := p.presigns[key]
p.presignMu.Unlock()
if hit {
return blob, nil
}
blob, err := p.getPresignedURL(ctx, method, dgst)
if err != nil {
return presignedBlob{}, err
}
p.presignMu.Lock()
if p.presigns == nil {
p.presigns = make(map[presignKey]presignedBlob)
}
p.presigns[key] = blob
p.presignMu.Unlock()
return blob, nil
}
// maxHoldErrorBody bounds how much of a hold error response is read for
// classification. The hold's is a single http.Error line.
const maxHoldErrorBody = 4 << 10
// holdAuthFailureMarkers are the substrings that mark a hold 401/403 as "we
// could not authenticate you" rather than "you may not read this".
//
// The hold answers every blob authorization failure the same way
// (pkg/hold/pds/xrpc.go): http.Error with 403 and a body of
// "authorization failed: " + err.Error(). Two shapes come out of that, and they
// need opposite handling:
//
// authorization failed: service token authentication failed: token has expired
// authorization failed: access denied for blob:read: crew member lacks permission (required: blob:read or blob:write)
//
// The first says our own credential is stale, which a retry with a fresh one
// fixes, so it has to reach the client as a 401 challenge. The second is the
// hold's verdict on the user and must stay a 403, or the client would be sent
// to re-authenticate in a loop it can never win.
//
// Anything unrecognised is treated as a permission denial. That is the
// conservative direction: a misread denial costs one failed request, whereas a
// misread challenge costs a re-auth loop.
var holdAuthFailureMarkers = []string{
// "service token authentication failed", "DPoP authentication failed"
"authentication failed",
// pds.ErrTokenExpired, in case the hold ever reports it unwrapped
"token has expired",
"missing authorization header",
"invalid authorization header format",
"invalid authorization scheme",
"missing token",
}
// isHoldAuthenticationFailure reports whether a hold rejection is about the
// credential we presented rather than about what the caller is allowed to do.
func isHoldAuthenticationFailure(status int, body []byte) bool {
// A 401 is by definition a challenge: the hold is asking for credentials,
// not refusing a request it understood.
if status == http.StatusUnauthorized {
return true
}
lower := strings.ToLower(string(body))
for _, marker := range holdAuthFailureMarkers {
if strings.Contains(lower, marker) {
return true
}
}
return false
}
// 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.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
if 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")
}
// The body is a short http.Error line; cap the read anyway so a
// misbehaving hold cannot make us buffer a response of its choosing.
bodyBytes, _ := io.ReadAll(io.LimitReader(resp.Body, maxHoldErrorBody))
if isHoldAuthenticationFailure(resp.StatusCode, bodyBytes) {
// Our credential is stale, not the user's permissions. Drop the
// cached token so the retry mints a new one, and answer with a 401
// so BearerChallenge attaches WWW-Authenticate and the client runs
// the token dance again instead of failing the pull.
slog.Warn("Hold rejected the service token, challenging client to re-authenticate",
"component", "proxy_blob_store",
"hold_did", p.ctx.HoldDID,
"puller_did", p.ctx.PullerDID,
"status", resp.StatusCode,
"body", strings.TrimSpace(string(bodyBytes)))
if p.ctx.InvalidateServiceToken != nil {
p.ctx.InvalidateServiceToken()
}
return presignedBlob{}, errcode.ErrorCodeUnauthorized.WithMessage("service token expired, re-authenticate")
}
// A real permission decision: the puller authenticated fine and may
// not read this hold. Sending them back to re-authenticate would just
// loop, so keep it a 403.
return presignedBlob{}, errcode.ErrorCodeDenied.WithMessage("read access denied")
}
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"`
}
// errWriterClosed is what every entry point reports once the writer is done
// with, whether it was committed, cancelled, reaped, or closed by a part
// upload that failed. Where there is a more specific cause (a failed part) the
// writer reports that instead; see closedErr.
var errWriterClosed = errors.New("writer closed")
// inFlightPart is the one part upload a writer may have running in the
// background. The writer hands it a full buffer and carries straight on
// filling a second one.
//
// Ownership of buf passes to the goroutine at hand-off and comes back to the
// writer when done is closed. Nothing else may touch buf in between: the
// goroutine is reading straight out of it, with no copy.
type inFlightPart struct {
// number is the S3 part number this upload was assigned. Assigned under
// w.mu at hand-off, so part numbers follow the order the bytes arrived in.
number int
// buf is the full buffer being uploaded.
buf *bytes.Buffer
// done is closed once the goroutine has recorded its outcome under w.mu
// and handed buf back. A reader that takes w.mu after observing this close
// sees everything the goroutine wrote.
done chan struct{}
}
// 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.
//
// A large blob keeps one part in flight while the next buffer fills. The
// upload used to be strictly serial: while a part was being PUT to S3 the
// Docker PATCH body was not being drained, and while the buffer filled S3 sat
// idle, so the wall clock for a layer was receive time plus send time. The
// price is a second buffer, so a large upload's peak memory is
// 2 * maxBufferSize (32MB), not one buffer's worth. That second buffer is
// taken only when the budget can spare it without waiting; a writer that
// cannot have one falls back to the serial upload, which is slower but needs
// no memory the writer does not already hold.
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 // Next part number to hand out (starts at 1)
buffer *bytes.Buffer // Buffer currently being filled by Write
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
// flight is the part upload running in the background, or nil when none
// is. At most one, ever: when the buffer fills and this is not nil, the
// write waits for it before handing off the next part. Keeping it to one
// is what makes part numbers and ETags ordered for free.
flight *inFlightPart
// spare is the other buffer while it is idle: handed back by a finished
// part upload with its capacity intact, waiting to be filled again. nil
// until a part has landed, so a blob that never flushes never allocates a
// second buffer, and so does a writer whose parts go up one at a time.
spare *bytes.Buffer
// flightErr is the failure of a background part upload, kept so it can be
// reported by whatever touches the writer next: a Write, the wait for the
// in-flight slot, or Commit. Sticky, so the cause stays visible instead of
// degrading into a bare "writer closed" one call later.
flightErr error
// aborted records that the hold-side multipart session has already been
// aborted. A failing part aborts from its own goroutine, and Cancel,
// Commit and the sweeper may all arrive afterwards; the session must only
// be aborted once.
aborted bool
// mu guards everything above against the abandoned-upload sweeper and
// against the writer's own background part upload. Distribution hands a
// given upload to one request at a time, so Write, Commit and Cancel never
// contend with each other.
//
// Nothing that talks to the network on the happy path is done under this
// lock: the hold calls and the S3 PUT of a part run in runPart with the
// lock dropped, which is the entire point of the change (a held lock there
// would serialise the upload again through Cancel and the sweeper, and
// would pin the writer for minutes). The lock is taken only to swap
// buffers and update state.
mu sync.Mutex
// charged is the buffer budget this writer currently holds, in bytes. It is
// the single source of truth for release: every exit path releases exactly
// this and zeroes it, so nothing can be released twice.
charged int64
// lastActivity is when the writer last did something on the client's
// behalf. startedAt is the wrong signal for the sweeper: a slow but healthy
// push can run for hours, while an abandoned one goes quiet immediately.
lastActivity time.Time
// requestCtx is the context of the request currently driving this writer,
// set by Create and by every Resume. Write has no context of its own (it is
// an io.Writer), and it can block waiting for budget, so it needs one:
// without it a client that has already hung up keeps a slot in the queue.
requestCtx context.Context
}
// 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
}
// adopt records the request now driving this writer and counts as activity.
// Called on Create and on every Resume, so a client that keeps coming back for
// the next PATCH is never mistaken for an abandoned one.
func (w *ProxyBlobWriter) adopt(ctx context.Context) {
w.mu.Lock()
defer w.mu.Unlock()
w.requestCtx = ctx
w.lastActivity = time.Now()
}
// waitContext is the context a blocking budget acquire should honour: the
// request currently driving the writer, or the background if there is none.
// Callers hold w.mu.
func (w *ProxyBlobWriter) waitContext() context.Context {
if w.requestCtx != nil {
return w.requestCtx
}
return context.Background()
}
// projectedCap is the backing array size the buffer will hold once an n byte
// write has landed. It mirrors growBuffer plus bytes.Buffer's own doubling: the
// array either already fits the write, grows to the doubled size, or jumps to
// the threshold. Nothing is ever sized past the threshold: Write splits a
// larger slice into threshold-sized parts.
func (w *ProxyBlobWriter) projectedCap(n int) int64 {
have := int64(w.buffer.Cap())
// Write fills the buffer to the threshold and flushes before taking more,
// so no write, however large, needs the buffer to grow past it.
need := min(int64(w.buffer.Len())+int64(n), maxBufferSize)
if need <= have {
return have
}
return min(max(need, 2*have), maxBufferSize)
}
// otherBufferCap is the backing array held by the writer's second buffer: the
// one currently being uploaded, or the one a finished upload handed back. Zero
// whenever the writer has only one buffer, which is every blob that never
// flushes and every writer the budget could not spare a second buffer for.
//
// Reading Cap on a buffer that is in flight is safe: runPart only reads out of
// it (Bytes), and the Reset that hands it back is done under w.mu.
//
// Callers hold w.mu.
func (w *ProxyBlobWriter) otherBufferCap() int64 {
if w.flight != nil {
return int64(w.flight.buf.Cap())
}
if w.spare != nil {
return int64(w.spare.Cap())
}
return 0
}
// budgetDelta is the extra budget the writer must acquire before an n byte
// write can land.
//
// What is charged is the backing arrays, not the bytes written. This path only
// ever charges for the buffer being filled, and only while that is the writer's
// first one: a second buffer is charged in full at hand-off (see
// tryTakeSecondBuffer), so once one exists the two capacities together are
// already covered by w.charged and the delta here is zero for the rest of the
// upload.
//
// That split is deliberate. This is the blocking charge, and a write must only
// ever block for memory the writer genuinely cannot proceed without. One full
// buffer is enough to finish any blob, so waiting for the first one is honest
// backpressure; waiting for the second was a wedge, because every writer
// mid-growth held memory that only a Commit could release and no writer could
// reach Commit.
//
// A blob that never fills a buffer never causes a hand-off, never allocates a
// second buffer, and is charged only for what its one buffer grew to.
//
// Nothing is released between parts, because bytes.Buffer.Reset keeps the
// array: releasing there would report memory as free while the writer still
// holds every byte of it. Everything goes back at once on Commit, Cancel or a
// reap.
//
// Callers hold w.mu.
func (w *ProxyBlobWriter) budgetDelta(n int) (int64, error) {
want := w.otherBufferCap() + w.projectedCap(n)
if want <= w.charged {
return 0, nil
}
if want > uploadBudget.limit {
// Only reachable for a single write larger than the whole budget, which
// means a caller handing the writer an entire oversized blob in one
// call rather than streaming it. The semaphore would never grant this,
// so say so now instead of waiting out the acquire's deadline.
return 0, fmt.Errorf("upload buffer of %d bytes exceeds the process-wide budget of %d", want, uploadBudget.limit)
}
return want - w.charged, nil
}
// releaseBudget returns everything this writer holds, both buffers included.
// Safe to call more than once: the second call has nothing to release. Callers
// hold w.mu.
func (w *ProxyBlobWriter) releaseBudget() {
uploadBudget.release(w.charged)
w.charged = 0
}
// closedErr is what a closed writer reports. A writer closed by a part upload
// that failed reports that failure instead of the generic sentinel, so the
// cause is not lost behind the closure it caused. Callers hold w.mu.
func (w *ProxyBlobWriter) closedErr() error {
if w.flightErr != nil {
return w.flightErr
}
return errWriterClosed
}
// reapIfIdle cancels the writer if it has been inactive for at least ttl,
// reporting how idle it was and whether it was reaped. The hold-side multipart
// is aborted on a detached context: there is no request left to borrow one from,
// and the request that opened this upload is long gone.
func (w *ProxyBlobWriter) reapIfIdle(now time.Time, ttl time.Duration) (time.Duration, bool) {
// A writer someone is actively inside is by definition not abandoned. Now
// that the part upload runs off the lock this is a short window (a buffer
// swap, some bookkeeping) rather than the length of a PUT, but skipping and
// looking again on the next sweep is still the right answer.
if !w.mu.TryLock() {
return 0, false
}
defer w.mu.Unlock()
// A part on its way to S3 is a live upload, and it is liveness the clock
// cannot see: lastActivity is only bumped when a part lands, so a writer
// whose last Write was longer ago than the timeout may still be mid-PUT.
// Reaping it would abort the multipart session out from under a goroutine
// that is still feeding it, and the abandoned-upload sweep exists to
// reclaim clients that hung up, not uploads that are working.
//
// This cannot hide a genuinely abandoned writer forever: runPart's context
// is bounded by uploadPartTimeout, so the part always lands or fails, and
// the next sweep sees a writer with no flight and a stale lastActivity.
if w.flight != nil {
return 0, false
}
idle := now.Sub(w.lastActivity)
if w.closed || idle < ttl {
return 0, false
}
w.closed = true
ctx, cancel := context.WithTimeout(context.Background(), uploadAbortTimeout)
defer cancel()
w.abortIfStarted(ctx)
w.releaseBudget()
return idle, true
}
// Write writes data to the upload.
//
// Bytes are buffered until the buffer reaches maxBufferSize, at which point it
// is handed to a background part upload and this call carries straight on into
// a second buffer. At most one part is ever in flight: when the second buffer
// fills too, the write blocks until the running part lands, takes its buffer
// back, and hands the second one off. A writer the budget could not give a
// second buffer to uploads the part where it stands and carries on in the same
// buffer.
//
// Never let the buffer pass the threshold. Appending a whole chunk and
// checking afterwards let the last chunk before a flush land a few bytes past
// 16MB, which does not fit the 16MB backing array, so bytes.Buffer doubled it
// to 32MB and Reset kept that for the rest of the upload. Only chunk sizes that
// tile 16MB exactly avoided it, and the network read loop does not promise
// those. Filling to exactly the threshold, handing off, and continuing with the
// remainder pins capacity at 16MB for any chunk size, makes every part exactly
// one buffer, and means a single oversized Write streams through as parts
// instead of buffering whole.
func (w *ProxyBlobWriter) Write(p []byte) (int, error) {
written := 0
for written < len(p) {
// Each pass takes at most what fits in the current buffer, so a hand-off
// always happens at a buffer boundary. The buffer is empty again on the
// next pass, so a pass can never accept zero bytes and spin.
n, err := w.writeChunk(p[written:])
written += n
if err != nil {
return written, err
}
}
return written, nil
}
// writeChunk is one pass of Write: it takes as much of p as fits in the
// current buffer, charges the budget for it, appends it, and hands the buffer
// off to a background part upload if that filled it.
//
// The budget acquire happens with the lock dropped. Blocking there is the
// point (it is backpressure on the Docker client), but blocking with w.mu held
// would make Cancel, the sweeper and the writer's own part goroutine queue
// behind a write that is waiting on memory nobody has yet returned.
func (w *ProxyBlobWriter) writeChunk(p []byte) (int, error) {
w.mu.Lock()
if w.closed {
// Includes the one-part-late case: a background part that failed closed
// the writer, and closedErr reports why rather than just that it is shut.
err := w.closedErr()
w.mu.Unlock()
return 0, err
}
w.lastActivity = time.Now()
chunk := p
if room := maxBufferSize - w.buffer.Len(); len(chunk) > room {
chunk = chunk[:room]
}
delta, err := w.budgetDelta(len(chunk))
waitCtx := w.waitContext()
w.mu.Unlock()
if err != nil {
return 0, err
}
if err := uploadBudget.acquire(waitCtx, delta); err != nil {
return 0, err
}
w.mu.Lock()
defer w.mu.Unlock()
if w.closed {
// Cancelled, reaped, or closed by a failing part while this write waited
// for its budget. The delta was never folded into w.charged, so releasing
// it here cannot collide with the release that closing the writer did.
uploadBudget.release(delta)
return 0, w.closedErr()
}
// Only Write moves w.charged upward, and distribution drives a given upload
// from one request at a time, so nothing can have charged the buffers while
// this write was waiting.
w.charged += delta
w.growBuffer(len(chunk))
// bytes.Buffer.Write only fails by panicking on allocation, never by
// returning an error, so n is always len(chunk).
n, _ := w.buffer.Write(chunk)
w.size += int64(n)
// 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(chunk[:n])
if w.buffer.Len() >= maxBufferSize {
if err := w.handOffBuffer(); err != nil {
return n, err
}
}
return n, nil
}
// 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())
}
// handOffBuffer gives the full buffer to a part upload and leaves the writer an
// empty one to keep filling.
//
// The part goes up in the background whenever the writer has somewhere to carry
// on writing: the buffer a finished part handed back, or a fresh one when the
// budget for it is free right now. When it is not, the part goes up on this
// goroutine and the same buffer comes back, which is slower but always
// possible. That asymmetry is the safety property the budget rests on: a writer
// holding one full buffer can always finish the blob, so it can always reach
// the Commit that returns its budget.
//
// Only one part is ever in flight, so if a previous one is still running this
// waits for it first and reuses the buffer it hands back. That wait is where a
// client outrunning S3 gets its backpressure, and it is also where a part that
// failed a buffer ago is finally reported.
//
// Callers hold w.mu. The lock is dropped while waiting and held again on
// return.
func (w *ProxyBlobWriter) handOffBuffer() error {
if w.buffer.Len() == 0 {
return nil
}
if err := w.drainFlight(); err != nil {
return err
}
if w.closed {
// Cancelled or reaped while this call waited for the previous part.
return w.closedErr()
}
// The buffer the finished part handed back, capacity intact
// (bytes.Buffer.Reset keeps the backing array), so a large upload allocates
// its two buffers once and then just swaps them. nil only while the writer
// has never had a second buffer: a blob that never fills a buffer never
// gets here, and one whose second buffer the budget could not afford
// uploads its parts one at a time out of the first.
next := w.spare
w.spare = nil
if next == nil {
next = w.tryTakeSecondBuffer()
}
f := &inFlightPart{
number: w.partNumber,
buf: w.buffer,
done: make(chan struct{}),
}
w.partNumber++
w.flight = f
if next == nil {
// No budget to spare, so there is nothing to overlap the upload with:
// run it here and take the same buffer back. The next hand-off tries
// again, so a writer starts pipelining as soon as memory frees up.
return w.runPartInline(f)
}
w.buffer = next
go w.runPart(f)
return nil
}
// tryTakeSecondBuffer returns a second buffer for the writer to carry on in, or
// nil when the budget cannot spare one right now. Never waits: this is the
// charge the writer must be able to do without.
//
// The full maxBufferSize is charged and allocated up front rather than grown
// into. A writer only reaches here having already filled one whole buffer, so
// this one will be filled too unless the blob ends within the next part, and
// charging it whole keeps w.charged exactly equal to the backing arrays the
// writer holds. It also means no later write in this upload has to ask the
// budget for anything.
//
// A buffer's worth of budget is left free on purpose. Pipelining is a bonus,
// and spending the last of the budget on it would leave an arriving writer
// unable to fill even its first buffer, which is the one charge that has to be
// waited for.
//
// Callers hold w.mu.
func (w *ProxyBlobWriter) tryTakeSecondBuffer() *bytes.Buffer {
if !uploadBudget.tryAcquire(maxBufferSize, maxBufferSize) {
return nil
}
w.charged += maxBufferSize
next := &bytes.Buffer{}
next.Grow(maxBufferSize)
return next
}
// runPartInline uploads the part on this goroutine and takes its buffer back,
// which is what a writer with only one buffer has to do.
//
// The writer is left exactly as a background part leaves it, minus the overlap:
// the part is recorded (or its failure is), and the buffer comes back empty
// with its capacity intact to be filled again.
//
// Callers hold w.mu. The lock is dropped for the upload, which it must be:
// runPart takes it to record the outcome and hand the buffer back. w.flight is
// set for the whole of it, so the sweeper leaves the writer alone and a Cancel
// waits for the part rather than aborting the session underneath it.
func (w *ProxyBlobWriter) runPartInline(f *inFlightPart) error {
w.mu.Unlock()
w.runPart(f)
w.mu.Lock()
// runPart handed the buffer back as the spare. It is the only one this
// writer has, so it goes straight back to being the one being filled.
w.buffer = w.spare
w.spare = nil
if w.closed {
// The part failed, or a Cancel arrived while it was going up.
return w.closedErr()
}
return nil
}
// drainFlight waits for the in-flight part, if there is one, and reports
// whether it failed. The failure is sticky, so every later caller sees the
// cause rather than just a closed writer.
//
// Callers hold w.mu. The lock is dropped for the wait, which it has to be:
// runPart needs it to record its result and hand the buffer back.
func (w *ProxyBlobWriter) drainFlight() error {
if f := w.flight; f != nil {
w.mu.Unlock()
<-f.done
w.mu.Lock()
}
return w.flightErr
}
// abandonFlight waits for the in-flight part the way Cancel needs to: bounded,
// so a wedged PUT cannot pin a client's DELETE, and reporting whether the
// goroutine actually finished.
//
// Cancel must not abort the multipart session while a part PUT is still
// running. An abort that overtakes the PUT can arrive before the hold has even
// issued the upload ID, and then there is nothing left that knows the session
// exists: it leaks in S3 until the bucket's own multipart expiry catches it.
//
// Callers hold w.mu. The lock is dropped for the wait.
func (w *ProxyBlobWriter) abandonFlight() bool {
f := w.flight
if f == nil {
return true
}
timer := time.NewTimer(uploadFlightAbandonTimeout)
defer timer.Stop()
w.mu.Unlock()
finished := false
select {
case <-f.done:
finished = true
case <-timer.C:
}
w.mu.Lock()
return finished
}
// runPart uploads one part in the background and hands the buffer back.
//
// The context is detached and bounded rather than the request's: the PATCH
// that filled this buffer has usually been answered by the time the PUT
// finishes, and cancelling an upload because the request that produced its
// bytes ended is exactly wrong.
func (w *ProxyBlobWriter) runPart(f *inFlightPart) {
ctx, cancel := context.WithTimeout(context.Background(), uploadPartTimeout)
defer cancel()
etag, err := w.uploadPart(ctx, f.number, f.buf.Bytes())
w.mu.Lock()
if err != nil {
// Nobody was watching this goroutine, so the failure is reported one
// part late: by the next Write, by the next hand-off's wait, or by
// Commit. Close the writer here so no further bytes are accepted.
w.flightErr = fmt.Errorf("part %d upload failed: %w", f.number, err)
w.closed = true
} else {
// Ordered by construction: part number+1 is not handed off until this
// append has happened, because hand-off waits on f.done first.
w.parts = append(w.parts, CompletedPart{PartNumber: f.number, ETag: etag})
slog.Debug("Part uploaded successfully", "component", "proxy_blob_store/runPart", "part_number", f.number, "etag", etag)
}
// The buffer comes back to the writer with its capacity intact, ready to be
// the next part's. No budget is released: the memory is still held.
f.buf.Reset()
w.spare = f.buf
w.flight = nil
// A part landing is the writer doing work on the client's behalf, so it
// counts as activity. Without this a writer whose only remaining job was a
// slow PUT would keep ageing towards the sweeper's idle timeout.
w.lastActivity = time.Now()
if err != nil {
// Abort from here rather than leaving it to whoever notices the error:
// the upload ID is known now, this goroutine is the last thing touching
// the session, and the writer may never be touched again. abortIfStarted
// is idempotent, so a later Cancel or Commit does not double-abort.
abortCtx, abortCancel := context.WithTimeout(context.Background(), uploadAbortTimeout)
w.abortIfStarted(abortCtx)
abortCancel()
}
w.mu.Unlock()
close(f.done)
}
// uploadPart starts the multipart session if this is the first part, asks the
// hold for a presigned part URL, and PUTs the bytes to S3.
//
// No lock is held for any of it. That is what lets the next buffer fill while
// this one is going up, and it is why body is passed in rather than read off
// the writer: the caller owns those bytes for the duration.
func (w *ProxyBlobWriter) uploadPart(ctx context.Context, partNumber int, body []byte) (string, error) {
tempDigest := fmt.Sprintf("uploads/temp-%s", w.id)
w.mu.Lock()
uploadID := w.uploadID
w.mu.Unlock()
// Start the multipart upload on the first part 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.
// Only the first part can find this empty: parts are strictly serialised.
if uploadID == "" {
id, err := w.store.startMultipartUpload(ctx, tempDigest)
if err != nil {
return "", fmt.Errorf("failed to start multipart upload: %w", err)
}
uploadID = id
w.mu.Lock()
w.uploadID = uploadID
w.mu.Unlock()
}
// Get structured upload info for this part
uploadInfo, err := w.store.getPartUploadInfo(ctx, tempDigest, uploadID, 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(body))
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")
}
}
return etag, nil
}
// flushFinalPart uploads whatever is left in the buffer as the last part.
//
// Synchronous, unlike every other part: Commit has nothing left to overlap it
// with, and the ETag has to be in w.parts before completeUpload names them.
// Callers hold w.mu; the lock is dropped for the upload itself so a final PUT
// does not pin the writer. That is safe because Commit has already marked the
// writer closed and taken it out of globalUploads, so nothing else will accept
// bytes for it or reap it while this runs.
func (w *ProxyBlobWriter) flushFinalPart(ctx context.Context) error {
if w.buffer.Len() == 0 {
return nil
}
number := w.partNumber
w.partNumber++
body := w.buffer.Bytes()
w.mu.Unlock()
etag, err := w.uploadPart(ctx, number, body)
w.mu.Lock()
if err != nil {
return err
}
w.parts = append(w.parts, CompletedPart{PartNumber: number, ETag: etag})
w.buffer.Reset()
return nil
}
// ReadFrom reads from a reader
func (w *ProxyBlobWriter) ReadFrom(r io.Reader) (int64, error) {
// The lock is taken for this check and released again: every byte below
// goes through Write, which takes it per chunk. Holding it across the whole
// copy would pin the writer for the length of an upload and starve the
// sweeper's TryLock for just as long.
w.mu.Lock()
var closedErr error
if w.closed {
closedErr = w.closedErr()
}
w.mu.Unlock()
if closedErr != nil {
return 0, closedErr
}
// 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 {
w.mu.Lock()
defer w.mu.Unlock()
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 part still going up (if any) is waited
// for, 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) {
w.mu.Lock()
defer w.mu.Unlock()
if w.closed {
return distribution.Descriptor{}, w.closedErr()
}
w.closed = true
w.lastActivity = time.Now()
// Every path out of Commit is terminal, so the budget goes back here and
// nowhere else. The buffer is still needed below (the direct PUT reads
// straight out of it), which is why this is deferred rather than done now.
defer w.releaseBudget()
// 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.
//
// This runs before the wait for the in-flight part on purpose: the hash
// covers every byte Write accepted, including the ones still going up, so
// waiting would buy nothing and would only delay refusing a bad blob.
// Aborting, on the other hand, has to wait (see drainFlight and
// abandonFlight): an abort that overtakes a running PUT can leak the
// session.
if desc.Digest.Algorithm() != digest.Canonical {
// The abort has to wait for the part, even though the verdict did not.
_ = w.drainFlight()
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 {
// The abort has to wait for the part, even though the verdict did not.
_ = w.drainFlight()
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),
}
}
// Wait for the part still going up. Its ETag has to be in w.parts before
// completeUpload names them, and a failure it hit has to be reported here
// rather than swallowed. This is also what decides whether there is a
// multipart session at all: the first part starts one, and it may not have
// got that far yet.
if err := w.drainFlight(); err != nil {
w.abortIfStarted(ctx)
return distribution.Descriptor{}, err
}
// 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.flushFinalPart(ctx); 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 and has
// not been aborted already. A writer whose blob stayed inside the buffer has no
// session to abort.
//
// Exactly once, because there are now several ways to arrive here: a failing
// part aborts from its own goroutine, and Cancel, Commit and the sweeper may
// all follow it. Callers hold w.mu.
func (w *ProxyBlobWriter) abortIfStarted(ctx context.Context) {
if w.uploadID == "" || w.aborted {
return
}
w.aborted = true
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.mu.Lock()
defer w.mu.Unlock()
w.closed = true
defer w.releaseBudget()
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()
// Closed above, so no further bytes are taken; now let the part already on
// its way to S3 finish before aborting the session it belongs to. Bounded:
// a wedged PUT must not pin the client's DELETE, and if the wait does time
// out the abort still goes out, since a session that may exist is worth one
// best-effort abort.
if !w.abandonFlight() {
slog.Warn("Cancelling an upload whose part is still in flight",
"component", "proxy_blob_store/Cancel", "id", w.id, "waited", uploadFlightAbandonTimeout)
}
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")
}