at-container-registry/pkg/appview/storage/proxy_blob_store.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())

	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
	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 {
		// 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
}

// 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

	// mu guards everything above against the abandoned-upload sweeper, which is
	// the only thing that ever touches a writer concurrently with the request
	// that owns it. Distribution hands a given upload to one request at a time,
	// so Write, Commit and Cancel never contend with each other.
	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)
}

// budgetDelta is the extra budget the writer must acquire before an n byte
// write can land.
//
// What is charged is the buffer's backing array, not the bytes written: a
// writer holds exactly one buffer's worth at a time, whatever it has grown to.
// It is never 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.
//
// Callers hold w.mu.
func (w *ProxyBlobWriter) budgetDelta(n int) (int64, error) {
	want := 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. 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
}

// 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, and
	// Write holds this lock across a part upload, which can take minutes. Skip
	// it and look again on the next sweep rather than blocking every other
	// writer's reap behind one live upload.
	if !w.mu.TryLock() {
		return 0, false
	}
	defer w.mu.Unlock()

	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.
// Buffers data and flushes a part once the buffer reaches maxBufferSize.
func (w *ProxyBlobWriter) Write(p []byte) (int, error) {
	// Work out what this write costs, then wait for it without the lock held.
	// Blocking here is the point: it is backpressure on the Docker client. But
	// holding the writer's lock while blocked would make Cancel and the sweeper
	// queue behind a write that is waiting on memory nobody has yet returned.
	w.mu.Lock()
	if w.closed {
		w.mu.Unlock()
		return 0, fmt.Errorf("writer closed")
	}
	w.lastActivity = time.Now()
	delta, err := w.budgetDelta(len(p))
	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 or reaped 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 already did.
		uploadBudget.release(delta)
		return 0, fmt.Errorf("writer closed")
	}
	// Only Write moves w.charged upward, and distribution drives a given upload
	// from one request at a time, so nothing can have charged the buffer while
	// this write was waiting.
	w.charged += delta

	// 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, flushing, 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.
	written := 0
	for written < len(p) {
		chunk := p[written:]
		if room := maxBufferSize - w.buffer.Len(); len(chunk) > room {
			chunk = chunk[:room]
		}

		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)
		written += 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])

		// Flush once the buffer is full (S3 part size)
		if w.buffer.Len() >= maxBufferSize {
			if err := w.flushPart(); err != nil {
				return written, err
			}
		}
	}

	return written, 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())
}

// 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) {
	// 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()
	closed := w.closed
	w.mu.Unlock()
	if 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 {
	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 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{}, fmt.Errorf("writer closed")
	}
	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.
	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.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()

	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")
}