From aa45d7c1df895c444aacefdec7e5f13d09fff1a4 Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Thu, 14 Sep 2023 15:42:50 -0700
Subject: [PATCH] draft: explicitly follow upload redirects

---
 server/upload.go | 39 ++++++++++++++++++++++++++++++---------
 1 file changed, 30 insertions(+), 9 deletions(-)

diff --git a/server/upload.go b/server/upload.go
index 1f02cef3..161709b5 100644
--- a/server/upload.go
+++ b/server/upload.go
@@ -55,7 +55,7 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r
 	}
 	defer f.Close()
 
-	// 95MB chunk size
+	// 95MiB chunk size
 	chunkSize := 95 * 1024 * 1024
 	pw := ProgressWriter{
 		status: fmt.Sprintf("uploading %s", layer.Digest),
@@ -70,10 +70,10 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r
 			chunk = int64(chunkSize)
 		}
 
-		resp, err := uploadBlobChunk(ctx, requestURL, f, offset, chunk, regOpts, &pw)
+		resp, err := uploadBlobChunk(ctx, http.MethodPatch, requestURL, f, offset, chunk, regOpts, &pw)
 		if err != nil {
 			fn(api.ProgressResponse{
-				Status:    fmt.Sprintf("error uploading limit: %v", err),
+				Status:    fmt.Sprintf("error uploading chunk: %v", err),
 				Digest:    layer.Digest,
 				Total:     layer.Size,
 				Completed: int(offset),
@@ -83,12 +83,15 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r
 		}
 
 		offset += chunk
-		location, err := resp.Location()
+		location := resp.Header.Get("Docker-Upload-Location")
+		if location == "" {
+			location = resp.Header.Get("Location")
+		}
+
+		requestURL, err = url.Parse(location)
 		if err != nil {
 			return err
 		}
-
-		requestURL = location
 	}
 
 	values := requestURL.Query()
@@ -114,22 +117,40 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r
 	return nil
 }
 
-func uploadBlobChunk(ctx context.Context, requestURL *url.URL, r io.ReaderAt, offset, limit int64, opts *RegistryOptions, pw *ProgressWriter) (*http.Response, error) {
+func uploadBlobChunk(ctx context.Context, method string, requestURL *url.URL, r io.ReaderAt, offset, limit int64, opts *RegistryOptions, pw *ProgressWriter) (*http.Response, error) {
 	sectionReader := io.NewSectionReader(r, int64(offset), limit)
 
 	headers := make(http.Header)
 	headers.Set("Content-Type", "application/octet-stream")
 	headers.Set("Content-Length", strconv.Itoa(int(limit)))
-	headers.Set("Content-Range", fmt.Sprintf("%d-%d", offset, offset+sectionReader.Size()-1))
+	headers.Set("X-Redirect-Uploads", "1")
+
+	if method == http.MethodPatch {
+		headers.Set("Content-Range", fmt.Sprintf("%d-%d", offset, offset+sectionReader.Size()-1))
+	}
 
 	for try := 0; try < MaxRetries; try++ {
-		resp, err := makeRequest(ctx, "PATCH", requestURL, headers, io.TeeReader(sectionReader, pw), opts)
+		resp, err := makeRequest(ctx, method, requestURL, headers, io.TeeReader(sectionReader, pw), opts)
 		if err != nil && !errors.Is(err, io.EOF) {
 			return nil, err
 		}
 		defer resp.Body.Close()
 
 		switch {
+		case resp.StatusCode == http.StatusTemporaryRedirect:
+			location, err := resp.Location()
+			if err != nil {
+				return nil, err
+			}
+
+			pw.completed = int(offset)
+			if _, err := uploadBlobChunk(ctx, http.MethodPut, location, r, offset, limit, nil, pw); err != nil {
+				// retry
+				log.Printf("retrying redirected upload: %v", err)
+				continue
+			}
+
+			return resp, nil
 		case resp.StatusCode == http.StatusUnauthorized:
 			auth := resp.Header.Get("www-authenticate")
 			authRedir := ParseAuthRedirectString(auth)