optimize(transport): optimize resumable transport implementaion

derekhjray · derekhjray · commit 6e297fbc1e04 · 2025-11-14T10:40:35.000+08:00
diff --git a/pkg/v1/remote/transport/resumable.go b/pkg/v1/remote/transport/resumable.go
@@ -22,51 +22,70 @@ func NewResumable(inner http.RoundTripper) http.RoundTripper {
 
 var (
 	contentRangeRe = regexp.MustCompile(`^bytes (\d+)-(\d+)/(\d+|\*)$`)
+	rangeRe        = regexp.MustCompile(`bytes=(\d+)-(\d+)?`)
 )
 
 type resumableTransport struct {
 	inner http.RoundTripper
 }
 
-func (rt *resumableTransport) RoundTrip(in *http.Request) (*http.Response, error) {
-	if in.Method != http.MethodGet {
-		return rt.inner.RoundTrip(in)
+func (rt *resumableTransport) RoundTrip(in *http.Request) (resp *http.Response, err error) {
+	var total, start, end int64
+	// check initial request, maybe resumable transport is already enabled
+	if contentRange := in.Header.Get("Range"); contentRange != "" {
+		if matches := rangeRe.FindStringSubmatch(contentRange); len(matches) == 3 {
+			if start, err = strconv.ParseInt(matches[1], 10, 64); err != nil {
+				return nil, fmt.Errorf("invalid content range %q: %w", contentRange, err)
+			}
+
+			if len(matches[2]) == 0 {
+				// request whole file
+				end = -1
+			} else if end, err = strconv.ParseInt(matches[2], 10, 64); err == nil {
+				if start > end {
+					return nil, fmt.Errorf("invalid content range %q", contentRange)
+				}
+			} else {
+				return nil, fmt.Errorf("invalid content range %q: %w", contentRange, err)
+			}
+		}
 	}
 
-	req := in.Clone(in.Context())
-	req.Header.Set("Range", "bytes=0-")
-	resp, err := rt.inner.RoundTrip(req)
-	if err != nil {
+	if resp, err = rt.inner.RoundTrip(in); err != nil {
 		return resp, err
 	}
 
+	if in.Method != http.MethodGet {
+		return resp, nil
+	}
+
 	switch resp.StatusCode {
+	case http.StatusOK:
+		if end != 0 {
+			// request range content, but unexpected status code, cant not resume for this request
+			return resp, nil
+		}
+
+		total = resp.ContentLength
 	case http.StatusPartialContent:
-	case http.StatusRequestedRangeNotSatisfiable:
-		// fallback to previous behavior
-		resp.Body.Close()
-		return rt.inner.RoundTrip(in)
+		// keep original response status code, which should be processed by original transport or operation
+		if start, _, total, err = parseContentRange(resp.Header.Get("Content-Range")); err != nil || total <= 0 {
+			return resp, nil
+		} else if end > 0 {
+			total = end + 1
+		}
 	default:
 		return resp, nil
 	}
 
-	var contentLength int64
-	if _, _, contentLength, err = parseContentRange(resp.Header.Get("Content-Range")); err != nil || contentLength <= 0 {
-		// fallback to previous behavior
-		resp.Body.Close()
-		return rt.inner.RoundTrip(in)
-	}
-
-	// modify response status to 200, ensure caller error checking works
-	resp.StatusCode = http.StatusOK
-	resp.Status = "200 OK"
-	resp.ContentLength = contentLength
-	resp.Body = &resumableBody{
-		rc:          resp.Body,
-		inner:       rt.inner,
-		req:         req,
-		total:       contentLength,
-		transferred: 0,
+	if total > 0 {
+		resp.Body = &resumableBody{
+			rc:          resp.Body,
+			inner:       rt.inner,
+			req:         in,
+			total:       total,
+			transferred: start,
+		}
 	}
 
 	return resp, nil
@@ -94,14 +113,18 @@ func (rb *resumableBody) Read(p []byte) (n int, err error) {
 
 	for {
 		if n, err = rb.rc.Read(p); n > 0 {
+			if rb.transferred+int64(n) >= rb.total {
+				n = int(rb.total - rb.transferred)
+				err = io.EOF
+			}
 			rb.transferred += int64(n)
 		}
 
 		if err == nil {
 			return
 		}
 
-		if errors.Is(err, io.EOF) && rb.total >= 0 && rb.transferred == rb.total {
+		if errors.Is(err, io.EOF) && rb.total >= 0 && rb.transferred >= rb.total {
 			return
 		}
 
@@ -148,7 +171,8 @@ func (rb *resumableBody) resume(reason error) error {
 
 	if err = rb.validate(resp); err != nil {
 		resp.Body.Close()
-		return err
+		// wraps original error
+		return fmt.Errorf("%w, %v", reason, err)
 	}
 
 	if atomic.LoadUint32(&rb.closed) == 1 {
@@ -162,21 +186,21 @@ func (rb *resumableBody) resume(reason error) error {
 	return nil
 }
 
+const size100m = 100 << 20
+
 func (rb *resumableBody) validate(resp *http.Response) (err error) {
 	var start, total int64
 	switch resp.StatusCode {
 	case http.StatusPartialContent:
-		if start, _, total, err = parseContentRange(resp.Header.Get("Content-Range")); err != nil {
+		// donot using total size from Content-Range header, keep rb.total unchanged
+		if start, _, _, err = parseContentRange(resp.Header.Get("Content-Range")); err != nil {
 			return err
 		}
 
-		if total > rb.total {
-			rb.total = total
-		}
-
 		if start == rb.transferred {
 			break
 		} else if start < rb.transferred {
+			// incoming data is overlapped for somehow, just discard it
 			if _, err := io.CopyN(io.Discard, resp.Body, rb.transferred-start); err != nil {
 				return fmt.Errorf("discard overlapped data failed, %v", err)
 			}
@@ -185,6 +209,12 @@ func (rb *resumableBody) validate(resp *http.Response) (err error) {
 		}
 	case http.StatusOK:
 		if rb.transferred > 0 {
+			// range is not supported, and transferred data is too large, stop resuming
+			if rb.transferred > size100m {
+				return fmt.Errorf("too large data transferred: %d", rb.transferred)
+			}
+
+			// try resume from unsupported range request
 			if _, err = io.CopyN(io.Discard, resp.Body, rb.transferred); err != nil {
 				return err
 			}
diff --git a/pkg/v1/remote/transport/resumable_test.go b/pkg/v1/remote/transport/resumable_test.go
@@ -1,6 +1,7 @@
 package transport
 
 import (
+	"bytes"
 	"context"
 	"crypto/sha256"
 	"encoding/hex"
@@ -11,7 +12,6 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"os"
-	"regexp"
 	"strconv"
 	"testing"
 	"time"
@@ -23,50 +23,52 @@ import (
 	"github.com/google/go-containerregistry/pkg/v1/types"
 )
 
-var rangeRe = regexp.MustCompile(`bytes=(\d+)-(\d+)?`)
-
 func handleResumableLayer(data []byte, w http.ResponseWriter, r *http.Request, t *testing.T) {
 	if r.Method != http.MethodGet {
 		w.WriteHeader(http.StatusMethodNotAllowed)
 		return
 	}
 
-	contentRange := r.Header.Get("Range")
-	if contentRange == "" {
-		w.WriteHeader(http.StatusBadRequest)
-		return
-	}
-
-	matches := rangeRe.FindStringSubmatch(contentRange)
-	if len(matches) != 3 {
-		w.WriteHeader(http.StatusBadRequest)
-		return
-	}
-
-	contentLength := int64(len(data))
-	start, err := strconv.ParseInt(matches[1], 10, 64)
-	if err != nil || start < 0 {
-		w.WriteHeader(http.StatusBadRequest)
-		return
-	}
+	var (
+		contentLength, start, end int64
+		statusCode                = http.StatusOK
+		err                       error
+	)
 
-	if start >= int64(contentLength) {
-		w.WriteHeader(http.StatusRequestedRangeNotSatisfiable)
-		return
-	}
+	contentLength = int64(len(data))
+	end = contentLength - 1
+	contentRange := r.Header.Get("Range")
+	if contentRange != "" {
+		matches := rangeRe.FindStringSubmatch(contentRange)
+		if len(matches) != 3 {
+			w.WriteHeader(http.StatusBadRequest)
+			return
+		}
 
-	var end = int64(contentLength) - 1
-	if matches[2] != "" {
-		end, err = strconv.ParseInt(matches[2], 10, 64)
-		if err != nil || end < 0 {
+		if start, err = strconv.ParseInt(matches[1], 10, 64); err != nil || start < 0 {
 			w.WriteHeader(http.StatusBadRequest)
 			return
 		}
 
-		if end >= int64(contentLength) {
+		if start >= int64(contentLength) {
 			w.WriteHeader(http.StatusRequestedRangeNotSatisfiable)
 			return
 		}
+
+		if matches[2] != "" {
+			end, err = strconv.ParseInt(matches[2], 10, 64)
+			if err != nil || end < 0 {
+				w.WriteHeader(http.StatusBadRequest)
+				return
+			}
+
+			if end >= int64(contentLength) {
+				w.WriteHeader(http.StatusRequestedRangeNotSatisfiable)
+				return
+			}
+		}
+
+		statusCode = http.StatusPartialContent
 	}
 
 	var currentContentLength = end - start + 1
@@ -91,14 +93,19 @@ func handleResumableLayer(data []byte, w http.ResponseWriter, r *http.Request, t
 
 	end = start + currentContentLength - 1
 
-	w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, contentLength))
-	w.Header().Set("Content-Length", strconv.FormatInt(currentContentLength, 10))
-	w.WriteHeader(http.StatusPartialContent)
+	if statusCode == http.StatusPartialContent {
+		w.Header().Set("Content-Length", strconv.FormatInt(currentContentLength, 10))
+		w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, contentLength))
+	} else {
+		w.Header().Set("Content-Length", strconv.FormatInt(contentLength, 10))
+	}
+
+	w.WriteHeader(statusCode)
 	w.Write(data[start : end+1])
 	time.Sleep(time.Second)
 }
 
-func resumableRequest(client *http.Client, url string, size int64, digest string, overlap bool, t *testing.T) {
+func resumableRequest(client *http.Client, url string, leading, trailing []byte, size int64, digest string, overlap bool, t *testing.T) {
 	req, err := http.NewRequest(http.MethodGet, url, http.NoBody)
 	if err != nil {
 		t.Fatalf("http.NewRequest(): %v", err)
@@ -108,6 +115,16 @@ func resumableRequest(client *http.Client, url string, size int64, digest string
 		req.Header.Set("X-Overlap", "true")
 	}
 
+	if len(leading) > 0 || len(trailing) > 0 {
+		var buf bytes.Buffer
+		buf.WriteString("bytes=")
+		buf.WriteString(fmt.Sprintf("%d-", len(leading)))
+		if len(trailing) > 0 {
+			buf.WriteString(fmt.Sprintf("%d", size-int64(len(trailing))-1))
+		}
+		req.Header.Set("Range", buf.String())
+	}
+
 	resp, err := client.Do(req.WithContext(t.Context()))
 	if err != nil {
 		t.Fatalf("client.Do(): %v", err)
@@ -120,12 +137,19 @@ func resumableRequest(client *http.Client, url string, size int64, digest string
 	}
 
 	hash := sha256.New()
+	if len(leading) > 0 {
+		io.Copy(hash, bytes.NewReader(leading))
+	}
 
 	if _, err = io.Copy(hash, resp.Body); err != nil {
 		t.Errorf("unexpected error: %v", err)
 		return
 	}
 
+	if len(trailing) > 0 {
+		io.Copy(hash, bytes.NewReader(trailing))
+	}
+
 	actualDigest := "sha256:" + hex.EncodeToString(hash.Sum(nil))
 
 	if actualDigest != digest {
@@ -248,23 +272,40 @@ func TestResumableTransport(t *testing.T) {
 	}
 
 	tests := []struct {
-		name         string
-		digest       string
-		size         int64
-		timeout      bool
-		cancel       bool
-		nonResumable bool
-		overlap      bool
+		name              string
+		digest            string
+		leading, trailing int64
+		timeout           bool
+		cancel            bool
+		nonResumable      bool
+		overlap           bool
+		ranged            bool
 	}{
 		{
-			name:   "resumable",
-			digest: digest.String(),
-			size:   size,
+			name:    "resumable",
+			digest:  digest.String(),
+			leading: 0,
+		},
+		{
+			name:    "resumable-range-leading",
+			digest:  digest.String(),
+			leading: 3,
+		},
+		{
+			name:    "resumable-range-trailing",
+			digest:  digest.String(),
+			leading: 0,
+		},
+		{
+			name:     "resumable-range-leading-trailing",
+			digest:   digest.String(),
+			leading:  3,
+			trailing: 6,
 		},
 		{
 			name:    "resumable-overlap",
 			digest:  digest.String(),
-			size:    size,
+			leading: 0,
 			overlap: true,
 		},
 		{
@@ -290,8 +331,8 @@ func TestResumableTransport(t *testing.T) {
 				resumableStopByCancelRequest(client, url, t)
 			} else if tt.timeout {
 				resumableStopByTimeoutRequest(client, url, t)
-			} else if tt.digest != "" && tt.size > 0 {
-				resumableRequest(client, url, tt.size, tt.digest, tt.overlap, t)
+			} else if tt.digest != "" {
+				resumableRequest(client, url, data[:tt.leading], data[size-tt.trailing:], size, tt.digest, tt.overlap, t)
 			}
 		})
 	}
