From 120df54f358898865e26755f0faea798fb9166d2 Mon Sep 17 00:00:00 2001
From: Ben Sarmiento <ben.sarmiento@forto.com>
Date: Sun, 7 Jan 2024 15:43:07 +0100
Subject: [PATCH] Retry when range header is set but not 206

---
 pkg/http/client.go | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/pkg/http/client.go b/pkg/http/client.go
index ffed4ba..873d52f 100644
--- a/pkg/http/client.go
+++ b/pkg/http/client.go
@@ -22,7 +22,7 @@ type HTTPClient struct {
 	client       *http.Client
 	maxRetries   int
 	backoff      func(attempt int) time.Duration
-	getRetryIncr func(resp *http.Response, err error) int
+	getRetryIncr func(resp *http.Response, hasRangeHeader bool, err error) int
 	bearerToken  string
 	cfg          config.ConfigInterface
 	ipv6         cmap.ConcurrentMap[string, string]
@@ -44,11 +44,14 @@ func NewHTTPClient(token string, maxRetries int, timeoutSecs int, cfg config.Con
 			}
 			return time.Duration(backoff) * time.Second
 		},
-		getRetryIncr: func(resp *http.Response, err error) int {
+		getRetryIncr: func(resp *http.Response, hasRangeHeader bool, err error) int {
 			if resp != nil {
 				if resp.StatusCode == 429 || resp.StatusCode == 400 || resp.StatusCode == 403 {
 					return 1 // retry but don't increment attempt
 				}
+				if resp.StatusCode != http.StatusPartialContent && hasRangeHeader {
+					return 1
+				}
 				return 0 // don't retry
 			} else if err != nil {
 				errStr := err.Error()
@@ -112,13 +115,15 @@ func (r *HTTPClient) Do(req *http.Request) (*http.Response, error) {
 	if r.bearerToken != "" {
 		req.Header.Set("Authorization", "Bearer "+r.bearerToken)
 	}
+	// check if Range header is set
+	hasRangeHeader := req.Header.Get("Range") != ""
 
 	var resp *http.Response
 	var err error
 	attempt := 0
 	for {
 		resp, err = r.client.Do(req)
-		if incr := r.getRetryIncr(resp, err); incr > 0 {
+		if incr := r.getRetryIncr(resp, hasRangeHeader, err); incr > 0 {
 			attempt += incr
 			if attempt > r.maxRetries {
 				break