Ensure fetcher always closes body and properly calls release

dmcgowan · k8s-infra-cherrypick-robot · commit add2dcf86880 · 2025-07-19T01:57:47.000Z
Signed-off-by: Derek McGowan &lt;derek@mcg.dev&gt;
diff --git a/core/remotes/docker/fetcher.go b/core/remotes/docker/fetcher.go
@@ -474,7 +474,18 @@ func (r dockerFetcher) open(ctx context.Context, req *request, mediatype string,
 		return nil, err
 	}
 
-	body := resp.Body
+	body := &fnOnClose{
+		BeforeClose: func() {
+			r.Release(1)
+		},
+		ReadCloser: resp.Body,
+	}
+	defer func() {
+		if retErr != nil {
+			body.Close()
+		}
+	}()
+
 	encoding := strings.FieldsFunc(resp.Header.Get("Content-Encoding"), func(r rune) bool {
 		return r == ' ' || r == '\t' || r == ','
 	})
@@ -505,29 +516,33 @@ func (r dockerFetcher) open(ctx context.Context, req *request, mediatype string,
 		for i := range numChunks {
 			readers[i], writers[i] = newPipeWriter(bufPool)
 		}
+		// keep reference of the initial body value to ensure it is closed
+		ibody := body
 		go func() {
 			for i := range numChunks {
 				select {
 				case queue <- i:
 				case <-done:
+					if i == 0 {
+						ibody.Close()
+					}
 					return // avoid leaking a goroutine if we exit early.
 				}
 			}
 			close(queue)
 		}()
-		r.Release(1)
 		for range parallelism {
 			go func() {
 				for i := range queue { // first in first out
 					copy := func() error {
-						if err := r.Acquire(ctx, 1); err != nil {
-							return err
-						}
-						defer r.Release(1)
 						var body io.ReadCloser
 						if i == 0 {
-							body = resp.Body
+							body = ibody
 						} else {
+							if err := r.Acquire(ctx, 1); err != nil {
+								return err
+							}
+							defer r.Release(1)
 							reqClone := req.clone()
 							reqClone.setOffset(offset + i*chunkSize)
 							nresp, err := reqClone.doWithRetries(ctx, lastHost, withErrorCheck)
@@ -564,37 +579,27 @@ func (r dockerFetcher) open(ctx context.Context, req *request, mediatype string,
 			},
 			ReadCloser: io.NopCloser(io.MultiReader(readers...)),
 		}
-	} else {
-		defer func() {
-			if retErr != nil {
-				r.Release(1)
-			}
-		}()
-		body = &fnOnClose{
-			BeforeClose: func() {
-				r.Release(1)
-			},
-			ReadCloser: body,
-		}
 	}
+
 	for i := len(encoding) - 1; i >= 0; i-- {
 		algorithm := strings.ToLower(encoding[i])
 		switch algorithm {
 		case "zstd":
-			r, err := zstd.NewReader(body,
+			r, err := zstd.NewReader(body.ReadCloser,
 				zstd.WithDecoderLowmem(false),
 			)
 			if err != nil {
 				return nil, err
 			}
-			body = r.IOReadCloser()
+			body.ReadCloser = r.IOReadCloser()
 		case "gzip":
-			body, err = gzip.NewReader(body)
+			r, err := gzip.NewReader(body.ReadCloser)
 			if err != nil {
 				return nil, err
 			}
+			body.ReadCloser = r
 		case "deflate":
-			body = flate.NewReader(body)
+			body.ReadCloser = flate.NewReader(body.ReadCloser)
 		case "identity", "":
 			// no content-encoding applied, use raw body
 		default:
