From cc4069396f3ad2c370c53eed2e4a42ac13adab84 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Mon, 28 Sep 2026 13:31:38 -0700 Subject: [PATCH] mlx: bound pull stall retries and let the watchdog interrupt them (#18625) * transfer: bound stall retries and let the watchdog interrupt them A download that stalls could never recover. copy derived its own cancellable context, but the request was built from the parent, so cancelling on a stall never reached the in-flight body and the read stayed parked. The stall watchdog could detect a stall it had no way to stop. Stall retries were also unbounded: they never incremented attempt, so the retry loop had no limit, and because backoff was gated on attempt > 0 they retried with no delay at all. The request now owns the cancellable context and passes its cancel down to copy, so a stall aborts the read. Stalls get their own retry budget before counting against the limit, and backoff is keyed to every retry rather than only the counted ones. * move out of x/ --- server/download.go | 19 +++- server/download_test.go | 113 ++++++++++++++++++++- server/images.go | 2 +- {x/transfer => transfer}/download.go | 59 +++++++---- {x/transfer => transfer}/redirect.go | 26 ++++- {x/transfer => transfer}/redirect_test.go | 6 +- {x/transfer => transfer}/sparse_other.go | 0 {x/transfer => transfer}/sparse_windows.go | 0 {x/transfer => transfer}/transfer.go | 23 +++++ {x/transfer => transfer}/transfer_test.go | 87 ++++++++++++++++ {x/transfer => transfer}/upload.go | 8 +- 11 files changed, 302 insertions(+), 41 deletions(-) rename {x/transfer => transfer}/download.go (87%) rename {x/transfer => transfer}/redirect.go (84%) rename {x/transfer => transfer}/redirect_test.go (98%) rename {x/transfer => transfer}/sparse_other.go (100%) rename {x/transfer => transfer}/sparse_windows.go (100%) rename {x/transfer => transfer}/transfer.go (93%) rename {x/transfer => transfer}/transfer_test.go (96%) rename {x/transfer => transfer}/upload.go (99%) diff --git a/server/download.go b/server/download.go index bb2de7d24..b77342e91 100644 --- a/server/download.go +++ b/server/download.go @@ -25,6 +25,7 @@ import ( "github.com/ollama/ollama/api" "github.com/ollama/ollama/format" "github.com/ollama/ollama/manifest" + "github.com/ollama/ollama/transfer" "github.com/ollama/ollama/types/model" ) @@ -267,13 +268,23 @@ func (b *blobDownload) run(ctx context.Context, requestURL *url.URL, opts *regis if resp.StatusCode != http.StatusTemporaryRedirect && resp.StatusCode != http.StatusOK { return nil, fmt.Errorf("unexpected status code %d", resp.StatusCode) } - return resp.Location() + loc, err := resp.Location() + if err != nil { + return nil, err + } + if err := transfer.ValidateRedirectTarget(ctx, loc, requestURL.String(), opts.Insecure); err != nil { + return nil, err + } + return loc, nil } }() if err != nil { return err } + client := transfer.NewRedirectClient(requestURL.String(), opts.Insecure) + defer client.CloseIdleConnections() + g, inner := errgroup.WithContext(ctx) g.SetLimit(numDownloadParts) for i := range b.Parts { @@ -286,7 +297,7 @@ func (b *blobDownload) run(ctx context.Context, requestURL *url.URL, opts *regis var err error for try := 0; try < maxRetries; try++ { w := io.NewOffsetWriter(file, part.StartsAt()) - err = b.downloadChunk(inner, directURL, w, part) + err = b.downloadChunk(inner, client, directURL, w, part) switch { case errors.Is(err, context.Canceled), errors.Is(err, syscall.ENOSPC): // return immediately if the context is canceled or the device is out of space @@ -330,7 +341,7 @@ func (b *blobDownload) run(ctx context.Context, requestURL *url.URL, opts *regis return nil } -func (b *blobDownload) downloadChunk(ctx context.Context, requestURL *url.URL, w io.Writer, part *blobDownloadPart) error { +func (b *blobDownload) downloadChunk(ctx context.Context, client *http.Client, requestURL *url.URL, w io.Writer, part *blobDownloadPart) error { g, ctx := errgroup.WithContext(ctx) attemptStarted := time.Now() transferDone := make(chan struct{}) @@ -342,7 +353,7 @@ func (b *blobDownload) downloadChunk(ctx context.Context, requestURL *url.URL, w return err } req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", part.StartsAt(), part.StopsAt()-1)) - resp, err := http.DefaultClient.Do(req) + resp, err := client.Do(req) if err != nil { return err } diff --git a/server/download_test.go b/server/download_test.go index 1fd8e3b80..8bae9f2a3 100644 --- a/server/download_test.go +++ b/server/download_test.go @@ -1,19 +1,126 @@ package server import ( + "bytes" "context" "crypto/sha256" "errors" "fmt" "io" + "net" "net/http" "net/http/httptest" "net/url" "path/filepath" + "strings" + "sync/atomic" "testing" "time" + + "github.com/ollama/ollama/transfer" ) +func TestDownloadChunkRedirects(t *testing.T) { + for _, tc := range []struct { + name string + location string + insecure bool + wantErr string + wantHits int32 + }{ + {name: "public", location: "https://203.0.113.2/data", wantHits: 2}, + {name: "private", location: "https://127.0.0.1/data", wantErr: "not allowed", wantHits: 1}, + {name: "downgrade", location: "http://203.0.113.2/data", wantErr: "not allowed", wantHits: 1}, + {name: "lookup", wantErr: "not allowed"}, + {name: "insecure", insecure: true, wantHits: 2}, + {name: "loop", location: "https://203.0.113.1/start", wantErr: "stopped after 10 redirects", wantHits: 10}, + } { + t.Run(tc.name, func(t *testing.T) { + var hits atomic.Int32 + var location string + server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + if r.URL.Path == "/start" { + http.Redirect(w, r, location, http.StatusTemporaryRedirect) + return + } + w.Write([]byte("x")) + })) + t.Cleanup(server.Close) + location = tc.location + if location == "" { + location = server.URL + "/data" + } + client := transfer.NewRedirectClient("https://registry.example", tc.insecure) + t.Cleanup(client.CloseIdleConnections) + tr := client.Transport.(*http.Transport) + tr.TLSClientConfig = server.Client().Transport.(*http.Transport).TLSClientConfig.Clone() + tr.TLSClientConfig.ServerName = "127.0.0.1" + dial := tr.DialContext + tr.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) { + // Route only the public fixture addresses to the test server. + if addr == "203.0.113.1:443" || addr == "203.0.113.2:443" || addr == "203.0.113.2:80" { + return new(net.Dialer).DialContext(ctx, network, server.Listener.Addr().String()) + } + return dial(ctx, network, addr) + } + requestURL, _ := url.Parse("https://203.0.113.1/start") + if tc.name == "lookup" { + requestURL, _ = url.Parse(strings.Replace(server.URL, "127.0.0.1", "localhost", 1) + "/data") + } + download := &blobDownload{Name: filepath.Join(t.TempDir(), "blob")} + part := &blobDownloadPart{Size: 1, blobDownload: download} + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + var body bytes.Buffer + err := download.downloadChunk(ctx, client, requestURL, &body, part) + if tc.wantErr != "" { + if err == nil || !strings.Contains(err.Error(), tc.wantErr) { + t.Fatalf("downloadChunk() error = %v, want %q", err, tc.wantErr) + } + } else if err != nil || body.String() != "x" { + t.Fatalf("downloadChunk() = %q, %v, want x, nil", body.String(), err) + } + if got := hits.Load(); got != tc.wantHits { + t.Errorf("requests = %d, want %d", got, tc.wantHits) + } + }) + } +} + +func TestDownloadRedirect(t *testing.T) { + var hits atomic.Int32 + cdn := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + w.Write([]byte("x")) + })) + t.Cleanup(cdn.Close) + registry := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, cdn.URL, http.StatusTemporaryRedirect) + })) + t.Cleanup(registry.Close) + requestURL, _ := url.Parse(strings.Replace(registry.URL, "127.0.0.1", "localhost", 1)) + for _, insecure := range []bool{false, true} { + download := &blobDownload{ + Name: filepath.Join(t.TempDir(), "blob"), Total: 1, + Digest: "sha256:0000000000000000000000000000000000000000000000000000000000000000", + } + if err := download.newPart(0, 1); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + err := download.run(ctx, requestURL, ®istryOptions{Insecure: insecure}) + cancel() + if insecure { + if err != nil || hits.Load() != 1 { + t.Fatalf("insecure download: error = %v, requests = %d, want nil, 1", err, hits.Load()) + } + } else if err == nil || !strings.Contains(err.Error(), "not allowed") || hits.Load() != 0 { + t.Fatalf("download: error = %v, requests = %d, want rejection before request", err, hits.Load()) + } + } +} + func BenchmarkDownloadChunkCompletion(b *testing.B) { data := make([]byte, 1024*1024) digest := fmt.Sprintf("sha256:%x", sha256.Sum256(data)) @@ -36,7 +143,7 @@ func BenchmarkDownloadChunkCompletion(b *testing.B) { for range b.N { download := &blobDownload{Name: downloadPath, Digest: digest} part := &blobDownloadPart{Size: int64(len(data)), blobDownload: download} - if err := download.downloadChunk(b.Context(), requestURL, io.Discard, part); err != nil { + if err := download.downloadChunk(b.Context(), server.Client(), requestURL, io.Discard, part); err != nil { b.Fatal(err) } } @@ -63,7 +170,7 @@ func TestDownloadChunkReturnsWhenTransferCompletes(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), 250*time.Millisecond) defer cancel() - if err := download.downloadChunk(ctx, requestURL, io.Discard, part); err != nil { + if err := download.downloadChunk(ctx, server.Client(), requestURL, io.Discard, part); err != nil { t.Fatalf("downloadChunk() error = %v, want nil", err) } } @@ -96,7 +203,7 @@ func TestDownloadChunkDetectsStallBeforeFirstByte(t *testing.T) { }) started := time.Now() - err = download.downloadChunk(ctx, requestURL, io.Discard, part) + err = download.downloadChunk(ctx, server.Client(), requestURL, io.Discard, part) elapsed := time.Since(started) select { diff --git a/server/images.go b/server/images.go index 1c19d8154..ac4f77e8b 100644 --- a/server/images.go +++ b/server/images.go @@ -29,9 +29,9 @@ import ( "github.com/ollama/ollama/parser" "github.com/ollama/ollama/template" "github.com/ollama/ollama/thinking" + "github.com/ollama/ollama/transfer" "github.com/ollama/ollama/types/model" "github.com/ollama/ollama/version" - "github.com/ollama/ollama/x/transfer" ) // Blobs newer than this may belong to another process that has not written its diff --git a/x/transfer/download.go b/transfer/download.go similarity index 87% rename from x/transfer/download.go rename to transfer/download.go index 8a2aa4dc6..e80a32bfc 100644 --- a/x/transfer/download.go +++ b/transfer/download.go @@ -167,12 +167,14 @@ func download(ctx context.Context, opts DownloadOptions) error { func (d *downloader) download(ctx context.Context, blob Blob) error { var lastErr error - var slowRetries int + var slowRetries, stallRetries, retries int attempt := 0 for attempt < maxRetries { - if attempt > 0 { - if err := backoff(ctx, attempt, time.Second< 0 { + if err := backoff(ctx, retries, time.Second<= maxTransientRetries { + attempt++ + } case errors.Is(err, errSlow): - if slowRetries++; slowRetries >= 3 { - attempt++ // Only count after 3 slow retries + if slowRetries++; slowRetries >= maxTransientRetries { + attempt++ } default: attempt++ @@ -251,6 +257,11 @@ func (d *downloader) downloadOnce(ctx context.Context, blob Blob) (int64, error) } } + // The body read is cancelled through this context, so the stall watchdog in + // copy can actually interrupt a request that has stopped delivering bytes. + ctx, cancel := context.WithCancelCause(ctx) + defer cancel(nil) + req, _ := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil) req.Header.Set("User-Agent", d.userAgent) // Add auth only for same-host (not CDN) @@ -279,10 +290,10 @@ func (d *downloader) downloadOnce(ctx context.Context, blob Blob) (int64, error) return 0, fmt.Errorf("status %d", resp.StatusCode) } - return d.save(ctx, blob, resp.Body, existingSize) + return d.save(ctx, cancel, blob, resp.Body, existingSize) } -func (d *downloader) save(ctx context.Context, blob Blob, r io.Reader, existingSize int64) (int64, error) { +func (d *downloader) save(ctx context.Context, cancel context.CancelCauseFunc, blob Blob, r io.Reader, existingSize int64) (int64, error) { dest := filepath.Join(d.destDir, digestToPath(blob.Digest)) tmp := dest + ".tmp" os.MkdirAll(filepath.Dir(dest), 0o755) @@ -320,7 +331,7 @@ func (d *downloader) save(ctx context.Context, blob Blob, r io.Reader, existingS } defer f.Close() - n, err := d.copy(ctx, f, r, h) + n, err := d.copy(ctx, cancel, f, r, h) if err != nil { // Don't remove .tmp here — download() handles cleanup based on blob size return existingSize + n, err @@ -339,20 +350,26 @@ func (d *downloader) save(ctx context.Context, blob Blob, r io.Reader, existingS return totalWritten, os.Rename(tmp, dest) } -func (d *downloader) copy(ctx context.Context, dst io.Writer, src io.Reader, h io.Writer) (int64, error) { - var n int64 +// copy streams src to dst, cancelling the caller's request context when the +// transfer stalls or falls far behind the median speed. +func (d *downloader) copy(ctx context.Context, cancel context.CancelCauseFunc, dst io.Writer, src io.Reader, h io.Writer) (int64, error) { + var n atomic.Int64 var lastRead atomic.Int64 lastRead.Store(time.Now().UnixNano()) start := time.Now() - ctx, cancel := context.WithCancelCause(ctx) - defer cancel(nil) + // The watchdog cancels the caller's context, so it has to stop when this + // copy ends rather than when that context does. + done := make(chan struct{}) + defer close(done) go func() { tick := time.NewTicker(time.Second) defer tick.Stop() for { select { + case <-done: + return case <-ctx.Done(): return case <-tick.C: @@ -361,7 +378,7 @@ func (d *downloader) copy(ctx context.Context, dst io.Writer, src io.Reader, h i return } if e := time.Since(start); e > 5*time.Second { - if m := d.speeds.median(); m > 0 && float64(n)/e.Seconds() < m*0.1 { + if m := d.speeds.median(); m > 0 && float64(n.Load())/e.Seconds() < m*0.1 { cancel(errSlow) return } @@ -374,9 +391,9 @@ func (d *downloader) copy(ctx context.Context, dst io.Writer, src io.Reader, h i for { if err := ctx.Err(); err != nil { if c := context.Cause(ctx); c != nil { - return n, c + return n.Load(), c } - return n, err + return n.Load(), err } nr, err := src.Read(buf) @@ -385,13 +402,13 @@ func (d *downloader) copy(ctx context.Context, dst io.Writer, src io.Reader, h i dst.Write(buf[:nr]) h.Write(buf[:nr]) d.progress.add(int64(nr)) - n += int64(nr) + n.Add(int64(nr)) } if err == io.EOF { - return n, nil + return n.Load(), nil } if err != nil { - return n, err + return n.Load(), err } } } @@ -430,11 +447,11 @@ func (d *downloader) resolve(ctx context.Context, rawURL string) (*url.URL, erro } case http.StatusTemporaryRedirect, http.StatusFound, http.StatusMovedPermanently: loc, _ := resp.Location() - if err := validateRedirectScheme(loc, d.baseURL); err != nil { + if err := ValidateRedirectScheme(loc, d.baseURL); err != nil { return nil, err } if loc.Host != u.Host { - if err := validateRedirectTarget(ctx, loc, d.baseURL, d.allowPrivate); err != nil { + if err := ValidateRedirectTarget(ctx, loc, d.baseURL, d.allowPrivate); err != nil { return nil, err } return loc, nil diff --git a/x/transfer/redirect.go b/transfer/redirect.go similarity index 84% rename from x/transfer/redirect.go rename to transfer/redirect.go index cbe68915f..4570c7f03 100644 --- a/x/transfer/redirect.go +++ b/transfer/redirect.go @@ -48,12 +48,12 @@ func isPublicIP(ip net.IP) bool { return true } -// validateRedirectScheme rejects redirects that downgrade an https session +// ValidateRedirectScheme rejects redirects that downgrade an https session // to cleartext http, even when the target host is unchanged — a hostile or // compromised registry must not be able to strip TLS off follow-up requests. // It applies even under allowPrivate (the --insecure opt-in), which relaxes // address checks but never scheme checks. -func validateRedirectScheme(loc *url.URL, baseURL string) error { +func ValidateRedirectScheme(loc *url.URL, baseURL string) error { if loc == nil { return fmt.Errorf("%w: missing Location", errRedirectNotAllowed) } @@ -69,16 +69,16 @@ func validateRedirectScheme(loc *url.URL, baseURL string) error { return nil } -// validateRedirectTarget rejects redirect targets that aren't public HTTPS +// ValidateRedirectTarget rejects redirect targets that aren't public HTTPS // endpoints, unless allowPrivate is set. -func validateRedirectTarget(ctx context.Context, loc *url.URL, baseURL string, allowPrivate bool) error { +func ValidateRedirectTarget(ctx context.Context, loc *url.URL, baseURL string, allowPrivate bool) error { if loc == nil { return fmt.Errorf("%w: missing Location", errRedirectNotAllowed) } if allowPrivate { return nil } - if err := validateRedirectScheme(loc, baseURL); err != nil { + if err := ValidateRedirectScheme(loc, baseURL); err != nil { return err } host := loc.Hostname() @@ -144,6 +144,22 @@ func checkedDialer(d *net.Dialer, allowPrivate bool) func(ctx context.Context, n } } +// NewRedirectClient returns a client for download URLs supplied by a registry. +func NewRedirectClient(baseURL string, allowPrivate bool) *http.Client { + // Redirect destinations do not get the registry host's exemption. + client := checkedClient("", allowPrivate) + client.CheckRedirect = func(req *http.Request, via []*http.Request) error { + if len(via) >= 10 { + return errors.New("stopped after 10 redirects") + } + if err := ValidateRedirectScheme(req.URL, baseURL); err != nil { + return err + } + return ValidateRedirectTarget(req.Context(), req.URL, baseURL, allowPrivate) + } + return client +} + // checkedClient returns an HTTP client using checkedDialer; the registry // base host is exempt since the caller explicitly directed traffic at it. func checkedClient(baseURL string, allowPrivate bool) *http.Client { diff --git a/x/transfer/redirect_test.go b/transfer/redirect_test.go similarity index 98% rename from x/transfer/redirect_test.go rename to transfer/redirect_test.go index 68c1895f9..5a1760e80 100644 --- a/x/transfer/redirect_test.go +++ b/transfer/redirect_test.go @@ -259,7 +259,7 @@ func TestValidateRedirectTarget(t *testing.T) { for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { u, _ := url.Parse(tc.raw) - err := validateRedirectTarget(ctx, u, tc.base, tc.allowPrivate) + err := ValidateRedirectTarget(ctx, u, tc.base, tc.allowPrivate) if tc.wantErr && err == nil { t.Error("expected error, got nil") } @@ -270,7 +270,7 @@ func TestValidateRedirectTarget(t *testing.T) { } // The public-IP literal case must pass too (no DNS involved). - if err := validateRedirectTarget(ctx, okPublic, httpsBase, false); err != nil { + if err := ValidateRedirectTarget(ctx, okPublic, httpsBase, false); err != nil { t.Errorf("public https target rejected: %v", err) } } @@ -358,7 +358,7 @@ func TestValidateRedirectScheme(t *testing.T) { for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { u, _ := url.Parse(tc.raw) - err := validateRedirectScheme(u, tc.base) + err := ValidateRedirectScheme(u, tc.base) if tc.wantErr && err == nil { t.Error("expected error, got nil") } diff --git a/x/transfer/sparse_other.go b/transfer/sparse_other.go similarity index 100% rename from x/transfer/sparse_other.go rename to transfer/sparse_other.go diff --git a/x/transfer/sparse_windows.go b/transfer/sparse_windows.go similarity index 100% rename from x/transfer/sparse_windows.go rename to transfer/sparse_windows.go diff --git a/x/transfer/transfer.go b/transfer/transfer.go similarity index 93% rename from x/transfer/transfer.go rename to transfer/transfer.go index 49b14d6b4..887c831db 100644 --- a/x/transfer/transfer.go +++ b/transfer/transfer.go @@ -30,12 +30,15 @@ package transfer import ( "context" "errors" + "fmt" "log/slog" "math/rand/v2" "net/http" "strings" "sync/atomic" "time" + + "github.com/ollama/ollama/manifest" ) // Blob represents a content-addressed blob to transfer. @@ -109,6 +112,11 @@ const ( maxRetries = 6 defaultUserAgent = "ollama-transfer/1.0" + // maxTransientRetries is how many stalled or slow transfers a blob may + // absorb before they start counting against maxRetries. Both are usually + // recoverable, but a connection that only ever stalls still has to give up. + maxTransientRetries = 3 + // resumeThreshold is the minimum blob size for resume support. // Only blobs above this size keep partial .tmp files on failure. resumeThreshold = 64 << 20 // 64 MB @@ -161,14 +169,29 @@ func (p *progressTracker) add(n int64) { // Download downloads blobs in parallel with streaming hash verification. func Download(ctx context.Context, opts DownloadOptions) error { + if err := validateBlobs(opts.Blobs); err != nil { + return err + } return download(ctx, opts) } // Upload uploads blobs in parallel. func Upload(ctx context.Context, opts UploadOptions) error { + if err := validateBlobs(opts.Blobs); err != nil { + return err + } return upload(ctx, opts) } +func validateBlobs(blobs []Blob) error { + for _, blob := range blobs { + if err := manifest.ValidateDigest(blob.Digest); err != nil { + return fmt.Errorf("blob %q: %w", blob.Digest, err) + } + } + return nil +} + // digestToPath converts sha256:abc123 to sha256-abc123 func digestToPath(digest string) string { if len(digest) > 7 && digest[6] == ':' { diff --git a/x/transfer/transfer_test.go b/transfer/transfer_test.go similarity index 96% rename from x/transfer/transfer_test.go rename to transfer/transfer_test.go index c183a929f..4dd019e5a 100644 --- a/x/transfer/transfer_test.go +++ b/transfer/transfer_test.go @@ -12,11 +12,14 @@ import ( "net/http/httptest" "os" "path/filepath" + "strconv" "strings" "sync" "sync/atomic" "testing" "time" + + "github.com/ollama/ollama/manifest" ) // chunkedSession tracks accumulated PATCH body bytes for an upload session. @@ -83,6 +86,35 @@ func createTestBlob(t *testing.T, dir string, size int) (Blob, []byte) { return Blob{Digest: digest, Size: int64(size)}, data } +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +func TestTransferDigests(t *testing.T) { + root := t.TempDir() + dir := filepath.Join(root, "blobs") + blob, data := createTestBlob(t, dir, 8) + // A traversal must be rejected even when it resolves to a same-size cached file. + if err := os.WriteFile(filepath.Join(root, "outside"), data, 0o644); err != nil { + t.Fatal(err) + } + for _, digest := range []string{"", "../outside", `..\outside`, "sha256:../outside", "sha256:" + strings.Repeat("a", 63), "sha256:" + strings.Repeat("z", 64)} { + t.Run(digest, func(t *testing.T) { + client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Error("invalid digest reached HTTP transport") + return nil, errors.New("unexpected request") + })} + blobs := []Blob{blob, {Digest: digest, Size: blob.Size}} + if err := Download(t.Context(), DownloadOptions{Blobs: blobs, DestDir: dir, Client: client}); !errors.Is(err, manifest.ErrInvalidDigestFormat) { + t.Errorf("Download: got %v, want invalid digest", err) + } + if err := Upload(t.Context(), UploadOptions{Blobs: blobs, SrcDir: dir, Client: client}); !errors.Is(err, manifest.ErrInvalidDigestFormat) { + t.Errorf("Upload: got %v, want invalid digest", err) + } + }) + } +} + func TestDownload(t *testing.T) { // Create test blobs on "server" serverDir := t.TempDir() @@ -1801,6 +1833,61 @@ func TestResumeFromPartialFile(t *testing.T) { } } +// A connection that always stalls must eventually fail the download. Stall +// retries are budgeted rather than unlimited, so the loop cannot spin forever. +func TestDownloadStallGivesUp(t *testing.T) { + const blobSize = 4096 + blob, data := createTestBlob(t, t.TempDir(), blobSize) + + release := make(chan struct{}) + + // Each attempt issues two GETs: resolve, then the body. Answer the resolve + // normally and stall only the body, so the stall lands in copy. + var gets atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Length", strconv.Itoa(blobSize)) + if r.Method == http.MethodHead { + return + } + if gets.Add(1)%2 == 1 { + w.Write(data) + return + } + // Headers only: the body never arrives, so the transfer stalls. + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + <-release + })) + defer server.Close() + // Released before Close so the parked handlers cannot block shutdown. + defer close(release) + + done := make(chan error, 1) + go func() { + done <- Download(context.Background(), DownloadOptions{ + Blobs: []Blob{blob}, + BaseURL: server.URL, + DestDir: t.TempDir(), + StallTimeout: 20 * time.Millisecond, + }) + }() + + select { + case err := <-done: + if err == nil { + t.Fatal("Download() = nil, want an error after repeated stalls") + } + if !errors.Is(err, errMaxRetriesExceeded) { + t.Fatalf("Download() = %v, want %v", err, errMaxRetriesExceeded) + } + // maxTransientRetries + maxRetries attempts, each detected on the watchdog's + // one-second tick, plus backoff: ~13s here. The budget is loose because only + // the unbounded-retry case has to fail, and that case never finishes at all. + case <-time.After(90 * time.Second): + t.Fatal("Download retried forever after repeated stalls") + } +} + func TestResumeCorruptPartialFile(t *testing.T) { blobSize := resumeThreshold + 1024 data := make([]byte, blobSize) diff --git a/x/transfer/upload.go b/transfer/upload.go similarity index 99% rename from x/transfer/upload.go rename to transfer/upload.go index 0cffdb29e..e22aa9c35 100644 --- a/x/transfer/upload.go +++ b/transfer/upload.go @@ -376,11 +376,11 @@ func (u *uploader) initUpload(ctx context.Context, blob Blob) (uploadEndpoint, e if !sessionURL.IsAbs() { sessionURL = base.ResolveReference(sessionURL) } - if err := validateRedirectScheme(sessionURL, u.baseURL); err != nil { + if err := ValidateRedirectScheme(sessionURL, u.baseURL); err != nil { return uploadEndpoint{}, err } if base != nil && sessionURL.Host != base.Host { - if err := validateRedirectTarget(ctx, sessionURL, u.baseURL, u.allowPrivate); err != nil { + if err := ValidateRedirectTarget(ctx, sessionURL, u.baseURL, u.allowPrivate); err != nil { return uploadEndpoint{}, err } } @@ -397,7 +397,7 @@ func (u *uploader) initUpload(ctx context.Context, blob Blob) (uploadEndpoint, e // (percent-encoding case, query ordering) which can change the // canonical form a signed URL was computed over. if d, err := url.Parse(directURL); err == nil && d.IsAbs() { - if err := validateRedirectTarget(ctx, d, u.baseURL, u.allowPrivate); err != nil { + if err := ValidateRedirectTarget(ctx, d, u.baseURL, u.allowPrivate); err != nil { return uploadEndpoint{}, err } ep.directUploadURL = directURL @@ -670,7 +670,7 @@ func (u *uploader) uploadOnePart(ctx context.Context, sessionURL *url.URL, part if redirectURL == nil { return nil, nil, pr.bytes(), fmt.Errorf("patch part %d: 307 without Location", part.n) } - if err := validateRedirectTarget(ctx, redirectURL, u.baseURL, u.allowPrivate); err != nil { + if err := ValidateRedirectTarget(ctx, redirectURL, u.baseURL, u.allowPrivate); err != nil { return nil, nil, pr.bytes(), err } // The PATCH attempt's progress is wasted — we re-upload to CDN.