From 3cbc6a5c01120b59cf38378d5f18836bcd90f6c9 Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Thu, 17 Aug 2023 12:35:29 -0700
Subject: [PATCH] fix push manifest

---
 server/images.go | 74 +++++++++++++++++++++++++++---------------------
 1 file changed, 41 insertions(+), 33 deletions(-)

diff --git a/server/images.go b/server/images.go
index df1b26fd..c187490f 100644
--- a/server/images.go
+++ b/server/images.go
@@ -963,18 +963,12 @@ func PushModel(ctx context.Context, name string, regOpts *RegistryOptions, fn fu
 		return err
 	}
 
-	resp, err := makeRequest(ctx, "PUT", url, headers, bytes.NewReader(manifestJSON), regOpts)
+	resp, err := makeRequestWithRetry(ctx, "PUT", url, headers, bytes.NewReader(manifestJSON), regOpts)
 	if err != nil {
 		return err
 	}
 	defer resp.Body.Close()
 
-	// Check for success: For a successful upload, the Docker registry will respond with a 201 Created
-	if resp.StatusCode != http.StatusCreated {
-		body, _ := io.ReadAll(resp.Body)
-		return fmt.Errorf("on push registry responded with code %d: %v", resp.StatusCode, string(body))
-	}
-
 	fn(api.ProgressResponse{Status: "success"})
 
 	return nil
@@ -1116,43 +1110,18 @@ func GetSHA256Digest(r io.Reader) (string, int) {
 type requestContextKey string
 
 func startUpload(ctx context.Context, mp ModelPath, layer *Layer, regOpts *RegistryOptions) (string, error) {
-	retry, _ := ctx.Value(requestContextKey("retry")).(int)
-
 	url := fmt.Sprintf("%s/v2/%s/blobs/uploads/", mp.Registry, mp.GetNamespaceRepository())
 	if layer.From != "" {
 		url = fmt.Sprintf("%s/v2/%s/blobs/uploads/?mount=%s&from=%s", mp.Registry, mp.GetNamespaceRepository(), layer.Digest, layer.From)
 	}
 
-	resp, err := makeRequest(ctx, "POST", url, nil, nil, regOpts)
+	resp, err := makeRequestWithRetry(ctx, "POST", url, nil, nil, regOpts)
 	if err != nil {
 		log.Printf("couldn't start upload: %v", err)
 		return "", err
 	}
 	defer resp.Body.Close()
 
-	switch resp.StatusCode {
-	case http.StatusAccepted, http.StatusCreated:
-		// noop
-	case http.StatusUnauthorized:
-		if retry > MaxRetries {
-			return "", fmt.Errorf("max retries exceeded: %s", resp.Status)
-		}
-
-		auth := resp.Header.Get("www-authenticate")
-		authRedir := ParseAuthRedirectString(auth)
-		token, err := getAuthToken(ctx, authRedir, regOpts)
-		if err != nil {
-			return "", err
-		}
-
-		regOpts.Token = token
-		ctx = context.WithValue(ctx, requestContextKey("retry"), retry+1)
-		return startUpload(ctx, mp, layer, regOpts)
-	default:
-		body, _ := io.ReadAll(resp.Body)
-		return "", fmt.Errorf("on upload registry responded with code %d: %s", resp.StatusCode, body)
-	}
-
 	// Extract UUID location from header
 	location := resp.Header.Get("Location")
 	if location == "" {
@@ -1277,6 +1246,45 @@ func uploadBlobChunked(ctx context.Context, mp ModelPath, url string, layer *Lay
 	return nil
 }
 
+func makeRequestWithRetry(ctx context.Context, method, url string, headers map[string]string, body io.ReadSeeker, regOpts *RegistryOptions) (*http.Response, error) {
+	var status string
+	for try := 0; try < MaxRetries; try++ {
+		resp, err := makeRequest(ctx, method, url, headers, body, regOpts)
+		if err != nil {
+			log.Printf("couldn't start upload: %v", err)
+			return nil, err
+		}
+
+		status = resp.Status
+
+		switch resp.StatusCode {
+		case http.StatusAccepted, http.StatusCreated:
+			return resp, nil
+		case http.StatusUnauthorized:
+			auth := resp.Header.Get("www-authenticate")
+			authRedir := ParseAuthRedirectString(auth)
+			token, err := getAuthToken(ctx, authRedir, regOpts)
+			if err != nil {
+				return nil, err
+			}
+
+			regOpts.Token = token
+			if body != nil {
+				if _, err := body.Seek(0, io.SeekStart); err != nil {
+					return nil, err
+				}
+			}
+
+			continue
+		default:
+			body, _ := io.ReadAll(resp.Body)
+			return nil, fmt.Errorf("on upload registry responded with code %d: %s", resp.StatusCode, body)
+		}
+	}
+
+	return nil, fmt.Errorf("max retry exceeded: %v", status)
+}
+
 func makeRequest(ctx context.Context, method, url string, headers map[string]string, body io.Reader, regOpts *RegistryOptions) (*http.Response, error) {
 	if !strings.HasPrefix(url, "http") {
 		if regOpts.Insecure {