appview: upload small blobs with one presigned PUT, and verify every digest

Every blob went through the multipart machinery: an S3 multipart started
on Docker's initial POST, a hold round trip per part, and a complete on
the hold that finished the multipart, HEADed the temp object, copied it
to its final key, and deleted the temp. For a 2KB config blob that was
three hold calls and six S3 operations. On production data 86% of
distinct layers and every config blob fit in a 16MB buffer, and 49% of
image manifests have no layer larger than that.

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

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

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

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

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018Yf1ZVA7sXYhQNb9tCo1m5
This commit is contained in:
Evan Jarrett
2026-09-09 09:50:08 -05:00
co-authored by Claude Fable 5.1
parent 034ea5988b
commit f4343d7956
7 changed files with 862 additions and 103 deletions
+1 -1
View File
@@ -88,7 +88,7 @@ ATCR uses **distribution/distribution** as a library, extending it via middlewar
### Request Flow Summary
**Push:** Client pushes to `atcr.io/<identity>/<image>:<tag>`. 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/<identity>/<image>:<tag>`. 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.
+1 -1
View File
@@ -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
+19 -5
View File
@@ -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)
```
+4 -4
View File
@@ -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/<name>/blobs/uploads/<id>`, 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. |
+2
View File
@@ -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
+189 -35
View File
@@ -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
+646 -57
View File
@@ -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)
}
})
}