diff --git a/CLAUDE.md b/CLAUDE.md index 7798730..40a4321 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -88,7 +88,7 @@ ATCR uses **distribution/distribution** as a library, extending it via middlewar ### Request Flow Summary -**Push:** Client pushes to `atcr.io//:`. Registry middleware resolves identity → DID → PDS, discovers hold DID (from sailor profile `defaultHold` → AppView default). Blobs go to hold via XRPC multipart upload (presigned S3 URLs). Manifests stored in user's PDS as `io.atcr.manifest` records with `holdDid` reference. +**Push:** Client pushes to `atcr.io//:`. Registry middleware resolves identity → DID → PDS, discovers hold DID (from sailor profile `defaultHold` → AppView default). Blobs go to hold as presigned S3 uploads: anything under the appview's 16MB buffer is a single presigned PUT straight to the blob's final key, and only larger blobs use the XRPC multipart endpoints. Either way the appview hashes the bytes as it receives them and rejects a blob whose content does not match the digest the client claimed. Manifests stored in user's PDS as `io.atcr.manifest` records with `holdDid` reference. **Pull:** AppView fetches manifest from user's PDS. The manifest's `holdDid` field tells where blobs were stored. Blobs fetched from that hold via presigned download URLs. Pull always uses the historical hold from the manifest, even if the user changed their default since pushing. diff --git a/README.md b/README.md index 72d1ff3..1dd31f2 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,7 @@ atcr.io/did:plc:xyz123/myapp:latest **Storage model:** - Manifests → ATProto records in user's PDS (small JSON, includes `holdDid` reference) -- Blobs → Hold services via XRPC multipart upload (large binaries, stored in S3/etc.) +- Blobs → Hold services via presigned S3 uploads, direct PUT for small blobs and XRPC multipart for large ones (stored in S3/etc.) - AppView uses service tokens to communicate with holds on behalf of users ## Features diff --git a/docs/BYOS.md b/docs/BYOS.md index 290495a..afe4008 100644 --- a/docs/BYOS.md +++ b/docs/BYOS.md @@ -222,23 +222,37 @@ fly secrets set HOLD_REGISTRATION_OWNER_DID=did:plc:your-did-here GET /xrpc/com.atproto.server.getServiceAuth?aud=did:web:alice-storage.fly.dev Response: { "token": "eyJ..." } -5. AppView initiates multipart upload to hold: +5. AppView buffers the blob (16MB limit) and verifies the bytes it received + against the digest the client claimed. A mismatch is rejected here, before + anything reaches storage. + +5a. Small blob (fits in the buffer, which is every config blob and most layers): + AppView asks for one write capability and PUTs the whole blob to its final + content-addressed key. No multipart session, no temp object, no copy. + GET /xrpc/com.atproto.sync.getBlob?did=...&cid=sha256:abc...&method=PUT + Authorization: Bearer {serviceToken} + Response: { "url": "https://s3.../presigned" } + AppView: PUT that URL with Content-Type: application/octet-stream + +5b. Large blob (outgrew the buffer): multipart, as below. + +6. AppView initiates multipart upload to hold, on the first flush: POST https://alice-storage.fly.dev/xrpc/io.atcr.hold.initiateUpload Authorization: Bearer {serviceToken} Body: { "digest": "sha256:abc..." } Response: { "uploadId": "xyz" } -6. For each part: +7. For each part: - AppView: POST /xrpc/io.atcr.hold.getPartUploadUrl - Hold validates service token, checks crew membership - Hold returns: { "url": "https://s3.../presigned" } - - Client uploads directly to S3 presigned URL + - AppView uploads the part to the S3 presigned URL -7. AppView completes upload: +8. AppView completes upload: POST /xrpc/io.atcr.hold.completeUpload Body: { "uploadId": "xyz", "digest": "sha256:abc...", "parts": [...] } -8. Manifest stored in alice's PDS: +9. Manifest stored in alice's PDS: - holdDid: "did:web:alice-storage.fly.dev" - holdEndpoint: "https://alice-storage.fly.dev" (backward compat) ``` diff --git a/docs/HOLD_PUSH_OFFLOAD.md b/docs/HOLD_PUSH_OFFLOAD.md index 5655714..8ca3df5 100644 --- a/docs/HOLD_PUSH_OFFLOAD.md +++ b/docs/HOLD_PUSH_OFFLOAD.md @@ -20,7 +20,7 @@ property matters most. The **write path** inverts this. Today a push streams: ``` -client --PATCH/PUT--> AppView (buffers 10MB chunks in RAM) --presigned PUT--> S3 +client --PATCH/PUT--> AppView (buffers 16MB chunks in RAM) --presigned PUT--> S3 ``` AppView ingests every layer and re-uploads it to S3 (`proxy_blob_store.go:586-665` @@ -107,9 +107,9 @@ on the BYOS owner's box. | Concern | Today | |---|---| -| Push handshake | AppView `POST .../blobs/uploads/` -> distribution lib calls `ProxyBlobStore.Create` -> XRPC `io.atcr.hold.initiateUpload` (`proxy_blob_store.go:287-325`) | -| Push bytes | Client -> AppView RAM (10MB chunks) -> presigned S3 PUT (`proxy_blob_store.go:605-665`) | -| Push finalize | `ProxyBlobWriter.Commit` -> XRPC `io.atcr.hold.completeUpload` (`proxy_blob_store.go:703-741`) | +| Push handshake | AppView `POST .../blobs/uploads/` -> distribution lib calls `ProxyBlobStore.Create`, which makes no hold call. `io.atcr.hold.initiateUpload` is only reached if the blob outgrows the 16MB buffer. | +| Push bytes | Client -> AppView RAM (16MB buffer) -> presigned S3 PUT | +| Push finalize | `ProxyBlobWriter.Commit` verifies the received bytes against the client's digest, then either PUTs the whole buffered blob to its final key (a `com.atproto.sync.getBlob` PUT presign) or, for a blob that went multipart, calls XRPC `io.atcr.hold.completeUpload` | | Hold upload API | Custom XRPC only: `initiateUpload`, `getPartUploadUrl`, `completeUpload`, `abortUpload`, `notifyManifest` (`pkg/hold/oci/xrpc.go:50-61`). No standard OCI `/v2` upload surface. | | Upload `Location` | AppView-relative `/v2//blobs/uploads/`, generated by the distribution library from `BlobWriter.ID()`, not by ATCR code. | | Hold auth | Service token (Bearer, `aud`=hold DID, signed by user's PDS) or DPoP. Validated on every op (`pkg/hold/pds/auth.go:377-429` `ValidateBlobWriteAccess`, `:507-608` `ValidateServiceToken`). The hold **cannot** validate AppView's registry JWT. | diff --git a/docs/HOLD_XRPC_ENDPOINTS.md b/docs/HOLD_XRPC_ENDPOINTS.md index a7d1564..c1480b4 100644 --- a/docs/HOLD_XRPC_ENDPOINTS.md +++ b/docs/HOLD_XRPC_ENDPOINTS.md @@ -57,6 +57,8 @@ The endpoint routes on the `cid` parameter and answers differently per branch. The field is additive. A client that reads only `url` behaves exactly as before, and an AppView that gets no `size` falls back to HEADing the presigned URL. `size` is never sent for a PUT presign: the object does not exist yet. +A `method=PUT` presign is a write capability (gated on `blob:write`, same as the multipart endpoints) and is how the AppView uploads a blob small enough to fit in its 16MB buffer: one presigned PUT to the blob's final `sha256:` key, instead of initiateUpload plus part URLs plus completeUpload plus the server side copy out of the temp key. The URL is signed with `Content-Type: application/octet-stream`, so the PUT has to carry that header or S3 rejects the signature. + On GET and HEAD, a blob that is not in storage is answered with **404** and a JSON error body instead of a presigned URL: ```json diff --git a/pkg/appview/storage/proxy_blob_store.go b/pkg/appview/storage/proxy_blob_store.go index d66c8b8..6da2473 100644 --- a/pkg/appview/storage/proxy_blob_store.go +++ b/pkg/appview/storage/proxy_blob_store.go @@ -20,9 +20,19 @@ import ( ) const ( - // maxChunkSize is the maximum buffer size before flushing to hold service - // Matches S3's minimum multipart upload size - maxChunkSize = 10 * 1024 * 1024 // 10MB + // 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) @@ -275,8 +285,11 @@ func (p *ProxyBlobStore) Open(ctx context.Context, dgst digest.Digest) (io.ReadS }, nil } -// Put stores a blob using the multipart upload flow -// This ensures all uploads go through the same XRPC path +// 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 @@ -286,7 +299,7 @@ func (p *ProxyBlobStore) Put(ctx context.Context, mediaType string, content []by // Calculate digest dgst := digest.FromBytes(content) - // Use Create() flow for all uploads (goes through multipart XRPC endpoints) + // 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) @@ -343,7 +356,16 @@ func (p *ProxyBlobStore) ServeBlob(ctx context.Context, w http.ResponseWriter, r return nil } -// Create returns a blob writer for uploading using multipart upload. +// 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. @@ -359,22 +381,13 @@ func (p *ProxyBlobStore) Create(ctx context.Context, options ...distribution.Blo // Generate unique writer ID writerID := fmt.Sprintf("upload-%d", time.Now().UnixNano()) - // Use temp digest for upload location (will be moved to final digest on commit) - tempDigest := fmt.Sprintf("uploads/temp-%s", writerID) - - // Start multipart upload via hold service - uploadID, err := p.startMultipartUpload(ctx, tempDigest) - if err != nil { - return nil, err - } - writer := &ProxyBlobWriter{ store: p, options: opts, - uploadID: uploadID, parts: make([]CompletedPart, 0), partNumber: 1, - buffer: bytes.NewBuffer(make([]byte, 0, maxChunkSize)), + buffer: &bytes.Buffer{}, + digester: digest.Canonical.Digester(), id: writerID, startedAt: time.Now(), } @@ -644,14 +657,19 @@ type CompletedPart struct { ETag string `json:"etag"` } -// ProxyBlobWriter implements distribution.BlobWriter for proxy uploads using multipart upload +// 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 + 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) @@ -668,18 +686,26 @@ func (w *ProxyBlobWriter) StartedAt() time.Time { return w.startedAt } -// Write writes data to the upload -// Buffers data and flushes when buffer reaches 5MB +// Write writes data to the upload. +// Buffers data and flushes a part once the buffer reaches maxBufferSize. func (w *ProxyBlobWriter) Write(p []byte) (int, error) { if w.closed { return 0, fmt.Errorf("writer closed") } + w.growBuffer(len(p)) + n, err := w.buffer.Write(p) w.size += int64(n) + if n > 0 { + // Hash as we go. Nothing else in this path ever looked at the bytes: + // Commit took the digest in the client's final PUT on trust, which made + // the content address of a blob whatever the client claimed it was. + w.digester.Hash().Write(p[:n]) + } // Flush if buffer reaches limit (S3 part size) - if w.buffer.Len() >= maxChunkSize { + if w.buffer.Len() >= maxBufferSize { if err := w.flushPart(); err != nil { return n, err } @@ -688,6 +714,39 @@ func (w *ProxyBlobWriter) Write(p []byte) (int, error) { return n, err } +// growBuffer sizes the buffer's backing array ahead of an n byte write so that +// bytes.Buffer's doubling never overshoots maxBufferSize. +// +// Doubling is the right strategy while the buffer is small: a config blob ends +// up with a few KB of backing array instead of the full 16MB the writer used +// to reserve up front. But bytes.Buffer grows to max(needed, 2*cap), so a +// capacity that is anywhere past half the threshold doubles clean past it, and +// everything above maxBufferSize is wasted: the buffer is flushed and reset the +// moment it reaches the threshold. It is not a rounding error either. A layer +// streamed in 24KB chunks walks its capacity to 12MB and then doubles to 24MB, +// half of which is never used. +// +// So doubling is allowed only while the capacity it would land on still leaves +// room to double again. Past that, grow to exactly the threshold and stop, +// which is safe because the capacity at that point is at most half of it. +// +// After a flush, Reset keeps the capacity, so a large upload allocates its +// 16MB once and reuses it for every part. +func (w *ProxyBlobWriter) growBuffer(n int) { + c := w.buffer.Cap() + if c >= maxBufferSize { + return // Already at full size, nothing to do + } + + // What bytes.Buffer would grow to on its own if this write does not fit. + projected := max(w.buffer.Len()+n, 2*c) + if projected <= maxBufferSize/2 { + return // Still room for another doubling afterwards + } + + w.buffer.Grow(maxBufferSize - w.buffer.Len()) +} + // flushPart uploads the current buffer as a part func (w *ProxyBlobWriter) flushPart() error { if w.buffer.Len() == 0 { @@ -697,8 +756,20 @@ func (w *ProxyBlobWriter) flushPart() error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) defer cancel() - // Get structured upload info for this part tempDigest := fmt.Sprintf("uploads/temp-%s", w.id) + + // 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) @@ -786,7 +857,12 @@ func (w *ProxyBlobWriter) Size() int64 { return w.size } -// Commit finalizes the upload by completing multipart upload and moving to final location +// Commit finalizes the upload. +// +// The digest the client sent is verified against the bytes actually received +// before anything else happens, then the blob is finalized: a direct PUT to +// the final key if it is all still buffered, otherwise a final part plus the +// hold's completeUpload. func (w *ProxyBlobWriter) Commit(ctx context.Context, desc distribution.Descriptor) (distribution.Descriptor, error) { if w.closed { return distribution.Descriptor{}, fmt.Errorf("writer closed") @@ -798,15 +874,47 @@ func (w *ProxyBlobWriter) Commit(ctx context.Context, desc distribution.Descript 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 - if err := w.store.abortMultipartUpload(ctx, w.uploadID); err != nil { - slog.Warn("Failed to abort multipart upload", "component", "proxy_blob_store/Cancel", "error", err) - // Continue anyway - we want to mark upload as cancelled - } + w.abortIfStarted(ctx) return distribution.Descriptor{}, fmt.Errorf("failed to flush final part: %w", err) } } @@ -818,7 +926,7 @@ func (w *ProxyBlobWriter) Commit(ctx context.Context, desc distribution.Descript 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)) + 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, @@ -827,6 +935,56 @@ func (w *ProxyBlobWriter) Commit(ctx context.Context, desc distribution.Descript }, nil } +// putDirect uploads the fully buffered blob to its final content-addressed key +// with a single presigned PUT, skipping the multipart dance entirely. +func (w *ProxyBlobWriter) putDirect(ctx context.Context, dgst digest.Digest) error { + // Same hold endpoint the read path uses, asked for a write capability. + // The hold gates method=PUT on blob write access and skips its size lookup, + // since the object being presigned does not exist yet. + blob, err := w.store.getPresignedURL(ctx, http.MethodPut, dgst) + if err != nil { + return fmt.Errorf("failed to get presigned upload URL: %w", err) + } + + body := w.buffer.Bytes() + req, err := http.NewRequestWithContext(ctx, http.MethodPut, blob.URL, bytes.NewReader(body)) + if err != nil { + return err + } + // The hold signs the PUT with ContentType "application/octet-stream" + // (GetPresignedURL in pkg/hold/pds/xrpc.go), and a signed header that the + // request does not carry fails S3's signature check. Nothing else is baked + // into the signature, in particular no content length. + req.Header.Set("Content-Type", "application/octet-stream") + req.ContentLength = int64(len(body)) + + resp, err := w.store.httpClient.Do(req) + if err != nil { + return fmt.Errorf("failed to upload blob: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated { + bodyBytes, _ := io.ReadAll(resp.Body) + return fmt.Errorf("blob upload failed: status %d, body: %s", resp.StatusCode, string(bodyBytes)) + } + + slog.Debug("Blob uploaded directly to final location", "component", "proxy_blob_store/putDirect", "digest", dgst, "size", len(body)) + return nil +} + +// abortIfStarted aborts the multipart upload, if one was ever started. A writer +// whose blob stayed inside the buffer has no session to abort. +func (w *ProxyBlobWriter) abortIfStarted(ctx context.Context) { + if w.uploadID == "" { + return + } + if err := w.store.abortMultipartUpload(ctx, w.uploadID); err != nil { + slog.Warn("Failed to abort multipart upload", "component", "proxy_blob_store", "error", err) + // Continue anyway - we want to mark upload as cancelled + } +} + // Cancel cancels the upload by aborting the multipart upload func (w *ProxyBlobWriter) Cancel(ctx context.Context) error { w.closed = true @@ -838,11 +996,7 @@ func (w *ProxyBlobWriter) Cancel(ctx context.Context) error { delete(globalUploads, w.id) globalUploadsMu.Unlock() - // Abort multipart upload - if err := w.store.abortMultipartUpload(ctx, w.uploadID); err != nil { - slog.Warn("Failed to abort multipart upload", "component", "proxy_blob_store/Cancel", "error", err) - // Continue anyway - we want to mark upload as cancelled - } + w.abortIfStarted(ctx) slog.Debug("Upload cancelled", "component", "proxy_blob_store/Cancel", "id", w.id) return nil diff --git a/pkg/appview/storage/proxy_blob_store_test.go b/pkg/appview/storage/proxy_blob_store_test.go index 4607225..9514633 100644 --- a/pkg/appview/storage/proxy_blob_store_test.go +++ b/pkg/appview/storage/proxy_blob_store_test.go @@ -660,12 +660,19 @@ type mockHoldServer struct { PartURLCalls []mockPartURLCall CompleteCalls []mockCompleteCall AbortCalls []mockAbortCall + PresignCalls []mockPresignCall + + // TotalCalls counts every request the hold received, so a test can pin + // "Create talks to the hold zero times" rather than only checking that a + // particular endpoint went unused. + TotalCalls int // Error injection InitiateError error PartURLError error CompleteError error AbortError error + PresignError error // Response customization UploadID string @@ -690,6 +697,13 @@ type mockAbortCall struct { UploadID string } +// mockPresignCall records a com.atproto.sync.getBlob presign request, which is +// what the direct PUT path asks for in place of the multipart endpoints. +type mockPresignCall struct { + Method string + CID string +} + // mockS3Server mocks S3 presigned URL uploads type mockS3Server struct { *httptest.Server @@ -698,6 +712,11 @@ type mockS3Server struct { mu sync.Mutex Parts map[int][]byte + // DirectPuts records whole-blob PUTs to the final key (the non-multipart + // path), separately from multipart parts. + DirectPuts [][]byte + DirectContentTypes []string + // Error injection UploadError error @@ -715,9 +734,29 @@ func newMockHoldServer(t *testing.T, s3URL string) *mockHoldServer { m.mu.Lock() defer m.mu.Unlock() + m.TotalCalls++ + w.Header().Set("Content-Type", "application/json") switch { + case strings.Contains(r.URL.Path, atproto.SyncGetBlob): + if m.PresignError != nil { + w.WriteHeader(http.StatusInternalServerError) + fmt.Fprintf(w, `{"error":"%s"}`, m.PresignError.Error()) + return + } + + m.PresignCalls = append(m.PresignCalls, mockPresignCall{ + Method: r.URL.Query().Get("method"), + CID: r.URL.Query().Get("cid"), + }) + + // A PUT presign points at the blob's final key, with no part number. + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(map[string]string{ + "url": fmt.Sprintf("%s/blob?cid=%s", s3URL, r.URL.Query().Get("cid")), + }) + case strings.Contains(r.URL.Path, atproto.HoldInitiateUpload): if m.InitiateError != nil { w.WriteHeader(http.StatusInternalServerError) @@ -822,11 +861,19 @@ func newMockS3Server(t *testing.T, etagInHeader bool) *mockS3Server { return } - // Parse part number from URL - partNum, _ := strconv.Atoi(r.URL.Query().Get("partNumber")) - // Read body body, _ := io.ReadAll(r.Body) + + // A whole-blob PUT lands on /blob and carries no part number. + if r.URL.Path == "/blob" { + m.DirectPuts = append(m.DirectPuts, body) + m.DirectContentTypes = append(m.DirectContentTypes, r.Header.Get("Content-Type")) + w.WriteHeader(http.StatusOK) + return + } + + // Parse part number from URL + partNum, _ := strconv.Atoi(r.URL.Query().Get("partNumber")) m.Parts[partNum] = body // Generate ETag @@ -870,7 +917,10 @@ func generateTestData(n int) []byte { return data } -// TestCreate_Success tests that Create() successfully initiates multipart upload +// TestCreate_Success tests that Create() returns a usable writer without +// talking to the hold at all. Docker POSTs to open an upload before it has +// said anything about the blob, and most blobs never need a multipart upload, +// so Create must not start one. func TestCreate_Success(t *testing.T) { s3Server := newMockS3Server(t, true) defer s3Server.Close() @@ -890,12 +940,13 @@ func TestCreate_Success(t *testing.T) { t.Fatal("Expected non-nil writer") } - // Verify initiate was called + // Verify Create made no hold call whatsoever holdServer.mu.Lock() - if len(holdServer.InitiateCalls) != 1 { - t.Errorf("Expected 1 initiate call, got %d", len(holdServer.InitiateCalls)) - } + totalCalls := holdServer.TotalCalls holdServer.mu.Unlock() + if totalCalls != 0 { + t.Errorf("Expected 0 hold calls from Create(), got %d", totalCalls) + } // Verify writer ID if writer.ID() == "" { @@ -914,8 +965,10 @@ func TestCreate_Success(t *testing.T) { writer.Cancel(context.Background()) } -// TestCreate_HoldError tests that Create() returns error when hold service fails -func TestCreate_HoldError(t *testing.T) { +// TestInitiate_HoldErrorSurfacesAtFirstFlush pins where a failing initiateUpload +// is now reported. Create no longer calls the hold, so the failure shows up on +// the first write that fills the buffer. +func TestInitiate_HoldErrorSurfacesAtFirstFlush(t *testing.T) { s3Server := newMockS3Server(t, true) defer s3Server.Close() @@ -926,16 +979,22 @@ func TestCreate_HoldError(t *testing.T) { store := createTestProxyBlobStore(t, holdServer.URL) writer, err := store.Create(context.Background()) - if err == nil { - t.Fatal("Expected error from Create()") + if err != nil { + t.Fatalf("Create() should not fail when the hold is down: %v", err) } - if writer != nil { - t.Error("Expected nil writer when error occurs") + defer writer.Cancel(context.Background()) + + _, err = writer.Write(generateTestData(maxBufferSize)) + if err == nil { + t.Fatal("Expected error from the flush that starts the multipart upload") } if !strings.Contains(err.Error(), "hold service unavailable") { t.Errorf("Expected hold error message, got: %v", err) } + if !strings.Contains(err.Error(), "failed to start multipart upload") { + t.Errorf("Expected start multipart error, got: %v", err) + } } // TestWrite_BasicBuffer tests that small writes are buffered @@ -980,7 +1039,8 @@ func TestWrite_BasicBuffer(t *testing.T) { } } -// TestWrite_TriggerFlush tests that writing 10MB triggers flush +// TestWrite_TriggerFlush tests that filling the buffer triggers a flush, and +// that the multipart upload is initiated then rather than at Create func TestWrite_TriggerFlush(t *testing.T) { s3Server := newMockS3Server(t, true) defer s3Server.Close() @@ -996,8 +1056,16 @@ func TestWrite_TriggerFlush(t *testing.T) { } defer writer.Cancel(context.Background()) - // Write exactly 10MB (the threshold) - data := generateTestData(10 * 1024 * 1024) + // Nothing has been written yet, so the hold has heard nothing + holdServer.mu.Lock() + initiatesBeforeWrite := len(holdServer.InitiateCalls) + holdServer.mu.Unlock() + if initiatesBeforeWrite != 0 { + t.Errorf("Expected 0 initiate calls before any write, got %d", initiatesBeforeWrite) + } + + // Write exactly the threshold + data := generateTestData(maxBufferSize) _, err = writer.Write(data) if err != nil { t.Fatalf("Write() failed: %v", err) @@ -1010,15 +1078,23 @@ func TestWrite_TriggerFlush(t *testing.T) { s3Server.mu.Unlock() if partCount != 1 { - t.Errorf("Expected 1 part uploaded after 10MB write, got %d", partCount) + t.Errorf("Expected 1 part uploaded after a full-buffer write, got %d", partCount) } - if len(uploadedData) != 10*1024*1024 { - t.Errorf("Expected uploaded part to be 10MB, got %d", len(uploadedData)) + if len(uploadedData) != maxBufferSize { + t.Errorf("Expected uploaded part to be %d bytes, got %d", maxBufferSize, len(uploadedData)) + } + + // The multipart upload was initiated by the flush, not by Create + holdServer.mu.Lock() + initiatesAfterWrite := len(holdServer.InitiateCalls) + holdServer.mu.Unlock() + if initiatesAfterWrite != 1 { + t.Errorf("Expected 1 initiate call after the first flush, got %d", initiatesAfterWrite) } } -// TestWrite_MultipleFlushes tests that writing 25MB triggers 2 flushes +// TestWrite_MultipleFlushes tests that writing 2.5x the threshold triggers 2 flushes func TestWrite_MultipleFlushes(t *testing.T) { s3Server := newMockS3Server(t, true) defer s3Server.Close() @@ -1034,8 +1110,8 @@ func TestWrite_MultipleFlushes(t *testing.T) { } defer writer.Cancel(context.Background()) - // Write 25MB in chunks (simulating Docker layer upload) - totalSize := 25 * 1024 * 1024 + // Write 2.5 buffers in chunks (simulating Docker layer upload) + totalSize := maxBufferSize * 5 / 2 chunkSize := 64 * 1024 // 64KB chunks data := generateTestData(chunkSize) @@ -1046,13 +1122,13 @@ func TestWrite_MultipleFlushes(t *testing.T) { } } - // Verify 2 flushes occurred (10MB + 10MB), 5MB remains in buffer + // Verify 2 flushes occurred, half a buffer remains s3Server.mu.Lock() partCount := len(s3Server.Parts) s3Server.mu.Unlock() if partCount != 2 { - t.Errorf("Expected 2 parts uploaded after 25MB write, got %d", partCount) + t.Errorf("Expected 2 parts uploaded after a 2.5x buffer write, got %d", partCount) } // Verify size tracking @@ -1117,7 +1193,7 @@ func TestFlushPart_Success(t *testing.T) { defer writer.Cancel(context.Background()) // Write enough to trigger flush - data := generateTestData(10 * 1024 * 1024) + data := generateTestData(maxBufferSize) _, err = writer.Write(data) if err != nil { t.Fatalf("Write() failed: %v", err) @@ -1155,7 +1231,7 @@ func TestFlushPart_ETagInJSON(t *testing.T) { defer writer.Cancel(context.Background()) // Write enough to trigger flush - data := generateTestData(10 * 1024 * 1024) + data := generateTestData(maxBufferSize) _, err = writer.Write(data) if err != nil { t.Fatalf("Write() failed: %v", err) @@ -1190,7 +1266,7 @@ func TestFlushPart_HoldError(t *testing.T) { defer writer.Cancel(context.Background()) // Write enough to trigger flush - data := generateTestData(10 * 1024 * 1024) + data := generateTestData(maxBufferSize) _, err = writer.Write(data) if err == nil { @@ -1220,7 +1296,7 @@ func TestFlushPart_S3Error(t *testing.T) { defer writer.Cancel(context.Background()) // Write enough to trigger flush - data := generateTestData(10 * 1024 * 1024) + data := generateTestData(maxBufferSize) _, err = writer.Write(data) if err == nil { @@ -1254,7 +1330,7 @@ func TestFlushPart_NoETag(t *testing.T) { defer writer.Cancel(context.Background()) // Write enough to trigger flush - data := generateTestData(10 * 1024 * 1024) + data := generateTestData(maxBufferSize) _, err = writer.Write(data) if err == nil { @@ -1321,8 +1397,8 @@ func TestReadFrom_LargeFile(t *testing.T) { } defer writer.Cancel(context.Background()) - // Stream 25MB through ReadFrom - data := generateTestData(25 * 1024 * 1024) + // Stream 2.5 buffers through ReadFrom + data := generateTestData(maxBufferSize * 5 / 2) reader := bytes.NewReader(data) n, err := writer.ReadFrom(reader) @@ -1340,7 +1416,7 @@ func TestReadFrom_LargeFile(t *testing.T) { s3Server.mu.Unlock() if partCount != 2 { - t.Errorf("Expected 2 parts (2x 10MB), got %d", partCount) + t.Errorf("Expected 2 parts (2 full buffers), got %d", partCount) } } @@ -1376,7 +1452,7 @@ func TestReadFrom_ClosedWriter(t *testing.T) { } } -// TestCommit_Success tests successful commit +// TestCommit_Success tests a successful multipart commit func TestCommit_Success(t *testing.T) { s3Server := newMockS3Server(t, true) defer s3Server.Close() @@ -1391,8 +1467,8 @@ func TestCommit_Success(t *testing.T) { t.Fatalf("Create() failed: %v", err) } - // Write some data - data := generateTestData(5 * 1024 * 1024) + // Write past the threshold so this takes the multipart path + data := generateTestData(maxBufferSize * 3 / 2) _, err = writer.Write(data) if err != nil { t.Fatalf("Write() failed: %v", err) @@ -1451,9 +1527,9 @@ func TestCommit_WithRemainingBuffer(t *testing.T) { t.Fatalf("Create() failed: %v", err) } - // Write 15MB in chunks (simulating realistic upload) - // This ensures: 1 flush at 10MB threshold + 5MB remaining in buffer - totalSize := 15 * 1024 * 1024 + // Write 1.5 buffers in chunks (simulating realistic upload) + // This ensures: 1 flush at the threshold + half a buffer remaining + totalSize := maxBufferSize * 3 / 2 chunkSize := 64 * 1024 // 64KB chunks chunk := generateTestData(chunkSize) @@ -1466,7 +1542,7 @@ func TestCommit_WithRemainingBuffer(t *testing.T) { allData = append(allData, chunk...) } - // At this point, 1 part should be uploaded (10MB), 5MB in buffer + // At this point, 1 full part should be uploaded, half a buffer remains s3Server.mu.Lock() partsBeforeCommit := len(s3Server.Parts) s3Server.mu.Unlock() @@ -1523,11 +1599,14 @@ func TestCommit_FlushError(t *testing.T) { t.Fatalf("Create() failed: %v", err) } - // Write some data - data := generateTestData(5 * 1024 * 1024) - _, err = writer.Write(data) - if err != nil { - t.Fatalf("Write() failed: %v", err) + // Write past the threshold in chunks so a multipart upload is started and a + // final part is still buffered when Commit runs + data := generateTestData(maxBufferSize + 1024) + for off := 0; off < len(data); off += 64 * 1024 { + end := min(off+64*1024, len(data)) + if _, err := writer.Write(data[off:end]); err != nil { + t.Fatalf("Write() failed: %v", err) + } } // Inject error for final flush @@ -1576,8 +1655,8 @@ func TestCommit_CompleteError(t *testing.T) { t.Fatalf("Create() failed: %v", err) } - // Write some data - data := generateTestData(1 * 1024 * 1024) + // Write past the threshold: completeUpload is only reached on the multipart path + data := generateTestData(maxBufferSize * 3 / 2) _, err = writer.Write(data) if err != nil { t.Fatalf("Write() failed: %v", err) @@ -1633,7 +1712,8 @@ func TestCommit_ClosedWriter(t *testing.T) { } } -// TestCancel_Success tests successful cancel +// TestCancel_Success tests that cancelling an upload that started a multipart +// aborts it func TestCancel_Success(t *testing.T) { s3Server := newMockS3Server(t, true) defer s3Server.Close() @@ -1648,6 +1728,11 @@ func TestCancel_Success(t *testing.T) { t.Fatalf("Create() failed: %v", err) } + // Fill the buffer so a multipart upload actually exists to abort + if _, err := writer.Write(generateTestData(maxBufferSize)); err != nil { + t.Fatalf("Write() failed: %v", err) + } + writerID := writer.ID() // Cancel @@ -1691,6 +1776,11 @@ func TestCancel_AbortError(t *testing.T) { t.Fatalf("Create() failed: %v", err) } + // Fill the buffer so there is a multipart upload to fail to abort + if _, err := writer.Write(generateTestData(maxBufferSize)); err != nil { + t.Fatalf("Write() failed: %v", err) + } + writerID := writer.ID() // Cancel should still return nil (graceful) @@ -1796,8 +1886,9 @@ func TestResume_NotFound(t *testing.T) { } } -// TestFullUploadFlow_25MB tests the complete upload flow with a 25MB file -func TestFullUploadFlow_25MB(t *testing.T) { +// TestFullUploadFlow_Multipart tests the complete upload flow with a blob big +// enough to need three parts +func TestFullUploadFlow_Multipart(t *testing.T) { s3Server := newMockS3Server(t, true) defer s3Server.Close() @@ -1812,8 +1903,10 @@ func TestFullUploadFlow_25MB(t *testing.T) { t.Fatalf("Create() failed: %v", err) } - // Write 25MB in chunks (simulating Docker layer upload) - totalSize := 25 * 1024 * 1024 + // Write 2 full buffers plus a partial one, in chunks (simulating a Docker + // layer upload) + tailSize := 5 * 1024 * 1024 + totalSize := 2*maxBufferSize + tailSize chunkSize := 64 * 1024 // 64KB chunks (realistic for Docker) allData := generateTestData(totalSize) @@ -1825,7 +1918,7 @@ func TestFullUploadFlow_25MB(t *testing.T) { } } - // Verify 2 parts uploaded during write (10MB + 10MB) + // Verify 2 parts uploaded during write (two full buffers) s3Server.mu.Lock() partsBeforeCommit := len(s3Server.Parts) s3Server.mu.Unlock() @@ -1844,7 +1937,7 @@ func TestFullUploadFlow_25MB(t *testing.T) { t.Fatalf("Commit() failed: %v", err) } - // Verify 3 total parts (10MB + 10MB + 5MB final) + // Verify 3 total parts (two full buffers plus the final remainder) s3Server.mu.Lock() partsAfterCommit := len(s3Server.Parts) s3Server.mu.Unlock() @@ -1880,9 +1973,9 @@ func TestFullUploadFlow_25MB(t *testing.T) { part3Size := len(s3Server.Parts[3]) s3Server.mu.Unlock() - expectedPart1 := 10 * 1024 * 1024 - expectedPart2 := 10 * 1024 * 1024 - expectedPart3 := 5 * 1024 * 1024 + expectedPart1 := maxBufferSize + expectedPart2 := maxBufferSize + expectedPart3 := tailSize if part1Size != expectedPart1 { t.Errorf("Part 1 expected %d bytes, got %d", expectedPart1, part1Size) @@ -2253,3 +2346,499 @@ func TestStat_HoldNotFoundIsBlobUnknown(t *testing.T) { t.Errorf("Expected ErrBlobUnknown from Get for a hold 404, got %v", err) } } + +// TestCommit_SmallBlobUsesDirectPut pins the fast path: a blob that never +// filled the buffer costs exactly one hold call (the PUT presign) and one S3 +// PUT, with no multipart session opened, no part URL handed out, no +// completeUpload, and therefore no temp object to copy and delete. +func TestCommit_SmallBlobUsesDirectPut(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + store := createTestProxyBlobStore(t, holdServer.URL) + + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + + // Create must not have talked to the hold at all + holdServer.mu.Lock() + callsAfterCreate := holdServer.TotalCalls + holdServer.mu.Unlock() + if callsAfterCreate != 0 { + t.Errorf("Expected 0 hold calls after Create(), got %d", callsAfterCreate) + } + + data := generateTestData(2048) + if _, err := writer.Write(data); err != nil { + t.Fatalf("Write() failed: %v", err) + } + + holdServer.mu.Lock() + callsAfterWrite := holdServer.TotalCalls + holdServer.mu.Unlock() + if callsAfterWrite != 0 { + t.Errorf("Expected 0 hold calls after a sub-threshold write, got %d", callsAfterWrite) + } + + dgst := digest.FromBytes(data) + desc, err := writer.Commit(context.Background(), distribution.Descriptor{ + Digest: dgst, + Size: int64(len(data)), + MediaType: "application/octet-stream", + }) + if err != nil { + t.Fatalf("Commit() failed: %v", err) + } + + if desc.Digest != dgst { + t.Errorf("Expected digest %s, got %s", dgst, desc.Digest) + } + if desc.Size != int64(len(data)) { + t.Errorf("Expected size %d, got %d", len(data), desc.Size) + } + + holdServer.mu.Lock() + totalCalls := holdServer.TotalCalls + presignCalls := append([]mockPresignCall(nil), holdServer.PresignCalls...) + initiateCount := len(holdServer.InitiateCalls) + partURLCount := len(holdServer.PartURLCalls) + completeCount := len(holdServer.CompleteCalls) + abortCount := len(holdServer.AbortCalls) + holdServer.mu.Unlock() + + if totalCalls != 1 { + t.Errorf("Expected exactly 1 hold call for a small blob, got %d", totalCalls) + } + if len(presignCalls) != 1 { + t.Fatalf("Expected 1 presign call, got %d", len(presignCalls)) + } + if presignCalls[0].Method != http.MethodPut { + t.Errorf("Expected a PUT presign, got method %q", presignCalls[0].Method) + } + if presignCalls[0].CID != dgst.String() { + t.Errorf("Expected presign for the final digest %s, got %s", dgst, presignCalls[0].CID) + } + if initiateCount != 0 || partURLCount != 0 || completeCount != 0 || abortCount != 0 { + t.Errorf("Expected no multipart traffic, got initiate=%d part=%d complete=%d abort=%d", + initiateCount, partURLCount, completeCount, abortCount) + } + + s3Server.mu.Lock() + directPuts := append([][]byte(nil), s3Server.DirectPuts...) + contentTypes := append([]string(nil), s3Server.DirectContentTypes...) + partCount := len(s3Server.Parts) + s3Server.mu.Unlock() + + if len(directPuts) != 1 { + t.Fatalf("Expected 1 direct S3 PUT, got %d", len(directPuts)) + } + if !bytes.Equal(directPuts[0], data) { + t.Error("Direct PUT body does not match the written bytes") + } + // The hold signs the PUT with ContentType application/octet-stream, so the + // request has to carry the same value or S3 rejects the signature. + if contentTypes[0] != "application/octet-stream" { + t.Errorf("Expected Content-Type application/octet-stream, got %q", contentTypes[0]) + } + if partCount != 0 { + t.Errorf("Expected 0 multipart parts, got %d", partCount) + } +} + +// TestPut_UsesDirectPut confirms ProxyBlobStore.Put inherits the fast path, +// since it goes through Create/Write/Commit. +func TestPut_UsesDirectPut(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + store := createTestProxyBlobStore(t, holdServer.URL) + + content := generateTestData(4096) + desc, err := store.Put(context.Background(), "application/vnd.oci.image.config.v1+json", content) + if err != nil { + t.Fatalf("Put() failed: %v", err) + } + + if desc.Digest != digest.FromBytes(content) { + t.Errorf("Expected digest %s, got %s", digest.FromBytes(content), desc.Digest) + } + + holdServer.mu.Lock() + totalCalls := holdServer.TotalCalls + initiateCount := len(holdServer.InitiateCalls) + holdServer.mu.Unlock() + + if totalCalls != 1 || initiateCount != 0 { + t.Errorf("Expected 1 hold call and no initiate from Put(), got total=%d initiate=%d", totalCalls, initiateCount) + } + + s3Server.mu.Lock() + directPuts := len(s3Server.DirectPuts) + s3Server.mu.Unlock() + if directPuts != 1 { + t.Errorf("Expected 1 direct S3 PUT from Put(), got %d", directPuts) + } +} + +// TestCommit_LargeBlobUsesMultipart pins the slow path: initiate happens on the +// first flush (not in Create), parts go up as before, and completeUpload closes +// it out with the right digest and size. +func TestCommit_LargeBlobUsesMultipart(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + store := createTestProxyBlobStore(t, holdServer.URL) + + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + + holdServer.mu.Lock() + callsAfterCreate := holdServer.TotalCalls + holdServer.mu.Unlock() + if callsAfterCreate != 0 { + t.Errorf("Expected 0 hold calls after Create(), got %d", callsAfterCreate) + } + + // Chunked, the way distribution streams a layer, so the tail is still + // buffered when Commit runs + data := generateTestData(maxBufferSize * 3 / 2) + for off := 0; off < len(data); off += 64 * 1024 { + end := min(off+64*1024, len(data)) + if _, err := writer.Write(data[off:end]); err != nil { + t.Fatalf("Write() failed: %v", err) + } + } + + // The first flush is what starts the multipart upload + holdServer.mu.Lock() + initiatesAfterWrite := len(holdServer.InitiateCalls) + partURLsAfterWrite := len(holdServer.PartURLCalls) + holdServer.mu.Unlock() + if initiatesAfterWrite != 1 { + t.Errorf("Expected 1 initiate call after the first flush, got %d", initiatesAfterWrite) + } + if partURLsAfterWrite != 1 { + t.Errorf("Expected 1 part URL call after the first flush, got %d", partURLsAfterWrite) + } + + dgst := digest.FromBytes(data) + desc, err := writer.Commit(context.Background(), distribution.Descriptor{ + Digest: dgst, + Size: int64(len(data)), + MediaType: "application/vnd.oci.image.layer.v1.tar+gzip", + }) + if err != nil { + t.Fatalf("Commit() failed: %v", err) + } + + if desc.Digest != dgst { + t.Errorf("Expected digest %s, got %s", dgst, desc.Digest) + } + if desc.Size != int64(len(data)) { + t.Errorf("Expected size %d, got %d", len(data), desc.Size) + } + + holdServer.mu.Lock() + initiateCount := len(holdServer.InitiateCalls) + completeCalls := append([]mockCompleteCall(nil), holdServer.CompleteCalls...) + presignCount := len(holdServer.PresignCalls) + holdServer.mu.Unlock() + + if initiateCount != 1 { + t.Errorf("Expected exactly 1 initiate call, got %d", initiateCount) + } + if len(completeCalls) != 1 { + t.Fatalf("Expected 1 complete call, got %d", len(completeCalls)) + } + if completeCalls[0].Digest != dgst.String() { + t.Errorf("Expected complete digest %s, got %s", dgst, completeCalls[0].Digest) + } + if len(completeCalls[0].Parts) != 2 { + t.Errorf("Expected 2 parts in the complete call, got %d", len(completeCalls[0].Parts)) + } + if presignCount != 0 { + t.Errorf("Expected no PUT presign on the multipart path, got %d", presignCount) + } + + s3Server.mu.Lock() + directPuts := len(s3Server.DirectPuts) + uploaded := len(s3Server.Parts[1]) + len(s3Server.Parts[2]) + s3Server.mu.Unlock() + + if directPuts != 0 { + t.Errorf("Expected no direct PUT on the multipart path, got %d", directPuts) + } + if uploaded != len(data) { + t.Errorf("Expected %d bytes uploaded across parts, got %d", len(data), uploaded) + } +} + +// TestCommit_DigestMismatchDirect pins that a client which misnames its bytes +// gets DIGEST_INVALID and puts nothing in the bucket. The digest is the blob's +// address in a shared content-addressed space, so this must be caught before +// any byte reaches S3. +func TestCommit_DigestMismatchDirect(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + store := createTestProxyBlobStore(t, holdServer.URL) + + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + + data := generateTestData(4096) + if _, err := writer.Write(data); err != nil { + t.Fatalf("Write() failed: %v", err) + } + + claimed := digest.FromString("something else entirely") + _, err = writer.Commit(context.Background(), distribution.Descriptor{ + Digest: claimed, + Size: int64(len(data)), + }) + if err == nil { + t.Fatal("Expected Commit() to reject a digest that does not match the content") + } + + var invalid distribution.ErrBlobInvalidDigest + if !errors.As(err, &invalid) { + t.Fatalf("Expected ErrBlobInvalidDigest, got %T: %v", err, err) + } + if invalid.Digest != claimed { + t.Errorf("Expected the claimed digest %s in the error, got %s", claimed, invalid.Digest) + } + + s3Server.mu.Lock() + directPuts := len(s3Server.DirectPuts) + partCount := len(s3Server.Parts) + s3Server.mu.Unlock() + if directPuts != 0 { + t.Errorf("Expected no S3 PUT for a mismatched digest, got %d", directPuts) + } + if partCount != 0 { + t.Errorf("Expected no S3 parts for a mismatched digest, got %d", partCount) + } + + holdServer.mu.Lock() + totalCalls := holdServer.TotalCalls + holdServer.mu.Unlock() + if totalCalls != 0 { + t.Errorf("Expected no hold calls for a mismatched digest, got %d", totalCalls) + } +} + +// TestCommit_DigestMismatchMultipart pins that a mismatch on the multipart path +// aborts the upload it already started instead of leaving a temp object behind. +func TestCommit_DigestMismatchMultipart(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + store := createTestProxyBlobStore(t, holdServer.URL) + + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + + data := generateTestData(maxBufferSize * 3 / 2) + if _, err := writer.Write(data); err != nil { + t.Fatalf("Write() failed: %v", err) + } + + _, err = writer.Commit(context.Background(), distribution.Descriptor{ + Digest: digest.FromString("wrong"), + Size: int64(len(data)), + }) + if err == nil { + t.Fatal("Expected Commit() to reject a digest that does not match the content") + } + + var invalid distribution.ErrBlobInvalidDigest + if !errors.As(err, &invalid) { + t.Fatalf("Expected ErrBlobInvalidDigest, got %T: %v", err, err) + } + + holdServer.mu.Lock() + abortCount := len(holdServer.AbortCalls) + completeCount := len(holdServer.CompleteCalls) + holdServer.mu.Unlock() + + if abortCount != 1 { + t.Errorf("Expected 1 abort call after a mismatched digest, got %d", abortCount) + } + if completeCount != 0 { + t.Errorf("Expected no complete call after a mismatched digest, got %d", completeCount) + } +} + +// TestCommit_UnsupportedDigestAlgorithm pins that a non-sha256 digest is +// refused rather than silently compared against a sha256 the writer computed. +func TestCommit_UnsupportedDigestAlgorithm(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + store := createTestProxyBlobStore(t, holdServer.URL) + + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + + data := generateTestData(4096) + if _, err := writer.Write(data); err != nil { + t.Fatalf("Write() failed: %v", err) + } + + _, err = writer.Commit(context.Background(), distribution.Descriptor{ + Digest: digest.SHA512.FromBytes(data), + Size: int64(len(data)), + }) + if !errors.Is(err, distribution.ErrBlobDigestUnsupported) { + t.Fatalf("Expected ErrBlobDigestUnsupported, got %T: %v", err, err) + } + + s3Server.mu.Lock() + directPuts := len(s3Server.DirectPuts) + s3Server.mu.Unlock() + if directPuts != 0 { + t.Errorf("Expected no S3 PUT for an unsupported algorithm, got %d", directPuts) + } +} + +// TestCancel_BeforeAnyFlushMakesNoAbortCall pins that cancelling a writer that +// never started a multipart upload does not send an abort for an upload ID +// that was never issued. +func TestCancel_BeforeAnyFlushMakesNoAbortCall(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + store := createTestProxyBlobStore(t, holdServer.URL) + + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + + if _, err := writer.Write(generateTestData(1024)); err != nil { + t.Fatalf("Write() failed: %v", err) + } + + if err := writer.Cancel(context.Background()); err != nil { + t.Fatalf("Cancel() failed: %v", err) + } + + holdServer.mu.Lock() + totalCalls := holdServer.TotalCalls + abortCount := len(holdServer.AbortCalls) + holdServer.mu.Unlock() + + if abortCount != 0 { + t.Errorf("Expected 0 abort calls when no multipart was started, got %d", abortCount) + } + if totalCalls != 0 { + t.Errorf("Expected 0 hold calls for a cancelled sub-threshold upload, got %d", totalCalls) + } +} + +// TestBufferGrowth_StaysWithinThreshold pins the buffer sizing. A config blob +// must not reserve the full threshold, and a buffer past the halfway mark must +// not double past it either. +func TestBufferGrowth_StaysWithinThreshold(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + store := createTestProxyBlobStore(t, holdServer.URL) + + t.Run("small blob keeps a small backing array", func(t *testing.T) { + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + defer writer.Cancel(context.Background()) + + pbw := writer.(*ProxyBlobWriter) + if c := pbw.buffer.Cap(); c != 0 { + t.Errorf("Expected an empty buffer from Create(), got cap %d", c) + } + + // A few KB, written the way distribution streams it + for range 4 { + if _, err := writer.Write(generateTestData(1024)); err != nil { + t.Fatalf("Write() failed: %v", err) + } + } + + if c := pbw.buffer.Cap(); c > 64*1024 { + t.Errorf("Expected a small backing array for a 4KB blob, got cap %d", c) + } + }) + + t.Run("past half the threshold does not overshoot it", func(t *testing.T) { + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + defer writer.Cancel(context.Background()) + + pbw := writer.(*ProxyBlobWriter) + + // Just over half the threshold. The chunk size matters: bytes.Buffer + // grows to max(needed, 2*cap), so a chain that is not a clean power of + // two of the threshold overshoots it. 24KB chunks walk the capacity to + // 12MB and then double to 24MB, which is what this pins against. + written := 0 + target := maxBufferSize - 2*1024*1024 + chunk := generateTestData(24 * 1024) + for written < target { + if _, err := writer.Write(chunk); err != nil { + t.Fatalf("Write() failed: %v", err) + } + written += len(chunk) + } + + if pbw.buffer.Len() < maxBufferSize/2 { + t.Fatalf("Expected the buffer to be past halfway, got len %d", pbw.buffer.Len()) + } + if c := pbw.buffer.Cap(); c > maxBufferSize { + t.Errorf("Expected cap to stop at the threshold %d, got %d", maxBufferSize, c) + } + + // Nothing was flushed, so this stayed off the multipart path entirely + holdServer.mu.Lock() + initiateCount := len(holdServer.InitiateCalls) + holdServer.mu.Unlock() + if initiateCount != 0 { + t.Errorf("Expected no multipart upload for a sub-threshold blob, got %d initiate calls", initiateCount) + } + }) +}