From 9f944c00f1a9c28ae4a2e49675ec9f1f3fc6a48d Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Wed, 16 Aug 2023 11:12:24 -0700
Subject: [PATCH] push: retry on unauthorized

---
 server/images.go | 25 +++++++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/server/images.go b/server/images.go
index 974c3fc30..31e51b681 100644
--- a/server/images.go
+++ b/server/images.go
@@ -1106,7 +1106,11 @@ func GetSHA256Digest(r io.Reader) (string, int) {
 	return fmt.Sprintf("sha256:%x", h.Sum(nil)), int(n)
 }
 
+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)
@@ -1119,8 +1123,25 @@ func startUpload(ctx context.Context, mp ModelPath, layer *Layer, regOpts *Regis
 	}
 	defer resp.Body.Close()
 
-	// Check for success
-	if resp.StatusCode != http.StatusAccepted && resp.StatusCode != http.StatusCreated {
+	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)
 	}