Merge pull request #486 from jmorganca/mxyng/fix-push

fix: retry push on expired token

server/auth.go

@@ -103,7 +103,7 @@ func getAuthToken(ctx context.Context, redirData AuthRedirect, regOpts *Registry
 
 	headers := make(http.Header)
 	headers.Set("Authorization", sig)
-	resp, err := makeRequest(ctx, "GET", redirectURL, headers, nil, regOpts)
+	resp, err := makeRequest(ctx, "GET", redirectURL, headers, nil, nil)
 	if err != nil {
 		log.Printf("couldn't get token: %q", err)
 	}

server/images.go

@@ -1313,11 +1313,13 @@ func makeRequest(ctx context.Context, method string, requestURL *url.URL, header
 		req.Header = headers
 	}
 
+	if regOpts != nil {
 		if regOpts.Token != "" {
 			req.Header.Set("Authorization", "Bearer "+regOpts.Token)
 		} else if regOpts.Username != "" && regOpts.Password != "" {
 			req.SetBasicAuth(regOpts.Username, regOpts.Password)
 		}
+	}
 
 	req.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))
 

server/upload.go

@@ -66,12 +66,19 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r
 
 		sectionReader := io.NewSectionReader(f, int64(offset), chunk)
 		for try := 0; try < MaxRetries; try++ {
+			ch := make(chan error, 1)
+
 			r, w := io.Pipe()
 			defer r.Close()
 			go func() {
 				defer w.Close()
 
 				for chunked := int64(0); chunked < chunk; {
+					select {
+					case err := <-ch:
+						log.Printf("chunk interrupted: %v", err)
+						return
+					default:
 						n, err := io.CopyN(w, sectionReader, 1024*1024)
 						if err != nil && !errors.Is(err, io.EOF) {
 							fn(api.ProgressResponse{
@@ -92,6 +99,7 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r
 							Completed: int(offset) + int(chunked),
 						})
 					}
+				}
 			}()
 
 			headers := make(http.Header)
@@ -113,6 +121,8 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r
 
 			switch {
 			case resp.StatusCode == http.StatusUnauthorized:
+				ch <- errors.New("unauthorized")
+
 				auth := resp.Header.Get("www-authenticate")
 				authRedir := ParseAuthRedirectString(auth)
 				token, err := getAuthToken(ctx, authRedir, regOpts)
@@ -121,10 +131,7 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r
 				}
 
 				regOpts.Token = token
-				if _, err := sectionReader.Seek(0, io.SeekStart); err != nil {
-					return err
-				}
-
+				sectionReader = io.NewSectionReader(f, int64(offset), chunk)
 				continue
 			case resp.StatusCode >= http.StatusBadRequest:
 				body, _ := io.ReadAll(resp.Body)