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

@@ -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
 		}

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, &registryOptions{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 {

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

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<<uint(attempt-1)); err != nil {
+		// Keyed to every retry, not just the ones that count against the limit,
+		// so repeated stalls cannot spin without pausing.
+		if retries > 0 {
+			if err := backoff(ctx, retries, time.Second<<uint(min(retries, maxRetries)-1)); err != nil {
 				return err
 			}
 		}
@@ -199,14 +201,18 @@ func (d *downloader) download(ctx context.Context, blob Blob) error {
 			os.Remove(dest + ".tmp")
 		}
 
+		retries++
+
 		switch {
 		case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded):
 			return err
 		case errors.Is(err, errStalled):
-			// Don't count stall retries against limit
+			if stallRetries++; stallRetries >= 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

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 {

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

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] == ':' {

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)

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.