From 1ff8cf2dfc09dda16ced121e0a35c1e2d306d071 Mon Sep 17 00:00:00 2001
From: Ben Sarmiento <ben.sarmiento@forto.com>
Date: Fri, 26 Jan 2024 18:33:15 +0100
Subject: [PATCH] Replace host on every retry

---
 pkg/http/client.go | 44 ++++++++++++++++++++++++++++----------------
 1 file changed, 28 insertions(+), 16 deletions(-)

diff --git a/pkg/http/client.go b/pkg/http/client.go
index d41e958..5ee4b12 100644
--- a/pkg/http/client.go
+++ b/pkg/http/client.go
@@ -10,6 +10,7 @@ import (
 	"net"
 	"net/http"
 	"net/url"
+	"strconv"
 	"strings"
 	"time"
 
@@ -25,6 +26,7 @@ import (
 type HTTPClient struct {
 	client         *http.Client
 	maxRetries     int
+	timeoutSecs    int
 	backoff        func(attempt int) time.Duration
 	getRetryIncr   func(resp *http.Response, reqHasRangeHeader bool, err error) int
 	bearerToken    string
@@ -54,6 +56,7 @@ func NewHTTPClient(token string, maxRetries int, timeoutSecs int, ensureIPv6Host
 		bearerToken: token,
 		client:      &http.Client{},
 		maxRetries:  maxRetries,
+		timeoutSecs: timeoutSecs,
 		backoff: func(attempt int) time.Duration {
 			maxDuration := 60
 			backoff := int(math.Pow(2, float64(attempt)))
@@ -182,7 +185,6 @@ func (r *HTTPClient) Do(req *http.Request) (*http.Response, error) {
 	if r.bearerToken != "" {
 		req.Header.Set("Authorization", "Bearer "+r.bearerToken)
 	}
-	r.replaceHostIfNeeded(req)
 	// check if Range header is set
 	reqHasRangeHeader := req.Header.Get("Range") != "" && req.Header.Get("Range") != "bytes=0-"
 
@@ -190,7 +192,21 @@ func (r *HTTPClient) Do(req *http.Request) (*http.Response, error) {
 	var err error
 	attempt := 0
 	for {
+		r.replaceHostIfNeeded(req)
+
+		ctx, cancel := context.WithTimeout(req.Context(), time.Duration(r.timeoutSecs)*time.Second)
+		defer cancel()
+		req = req.WithContext(ctx)
+
 		resp, err = r.client.Do(req)
+		select {
+		case <-ctx.Done():
+			if ctx.Err() == context.DeadlineExceeded {
+				err = fmt.Errorf("request timed out after %d seconds", r.timeoutSecs)
+			}
+		default:
+		}
+
 		if resp != nil && (resp.StatusCode < http.StatusOK || resp.StatusCode > http.StatusPartialContent) {
 			body, _ := io.ReadAll(resp.Body)
 			if body != nil {
@@ -221,19 +237,18 @@ func (r *HTTPClient) Do(req *http.Request) (*http.Response, error) {
 }
 
 func (r *HTTPClient) replaceHostIfNeeded(req *http.Request) {
-	if !r.ensureIPv6Host || !r.cfg.ShouldForceIPv6() {
+	if !r.ensureIPv6Host && !r.cfg.ShouldForceIPv6() || !strings.HasSuffix(req.URL.Host, "real-debrid.com") {
 		return
 	}
-	// if no hosts are found, just replace .com with .cloud
-	if len(r.ipv6Hosts) == 0 {
-		host := req.URL.Host
-		if strings.HasSuffix(host, ".com") {
-			newHost := strings.Replace(host, ".com", ".cloud", 1)
-			req.Host = newHost
-			req.URL.Host = newHost
-		}
+	// get subdomain of req.URL.Host
+	subdomain := strings.Split(req.URL.Host, ".")[0]
+	// check if subdomain is numeric
+	_, err := strconv.Atoi(subdomain)
+	if err == nil {
+		// subdomain is numeric, replace it with .cloud
+		req.URL.Host = strings.Replace(req.URL.Host, ".com", ".cloud", 1)
 	}
-	// if hosts are found, ensure the host is an IPv6 host
+	// check if host is in the list of IPv6 hosts
 	found := false
 	for _, h := range r.ipv6Hosts {
 		if h == req.URL.Host {
@@ -241,12 +256,9 @@ func (r *HTTPClient) replaceHostIfNeeded(req *http.Request) {
 			break
 		}
 	}
-	// if host is not an IPv6 host, replace it with a random IPv6 host
 	if !found {
-		r.log.Warnf("Host %s is not an IPv6 host, replacing with a random IPv6 host (ensure you have preferred_hosts properly set in your config.yml, if unset, run `zurg network-test -t ipv6`)", req.URL.Host)
-		newHost := r.ipv6Hosts[rand.Intn(len(r.ipv6Hosts))]
-		req.Host = newHost
-		req.URL.Host = newHost
+		// random IPv6 host
+		req.URL.Host = r.ipv6Hosts[rand.Intn(len(r.ipv6Hosts))]
 	}
 }