From cbd291400f260f94e626c7e3b60bde6efe54009a Mon Sep 17 00:00:00 2001
From: Ben Sarmiento <ben.sarmiento@forto.com>
Date: Fri, 24 Nov 2023 22:03:31 +0100
Subject: [PATCH] Handle preferred hosts in redirects too

---
 internal/universal/get.go | 25 +++++++++++++++++++++----
 pkg/http/client.go        |  4 ++--
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/internal/universal/get.go b/internal/universal/get.go
index c502ce2..c814e6f 100644
--- a/internal/universal/get.go
+++ b/internal/universal/get.go
@@ -3,6 +3,7 @@ package universal
 import (
 	"io"
 	"net/http"
+	"net/url"
 	"path"
 	"path/filepath"
 	"strings"
@@ -75,7 +76,7 @@ func (gf *GetFile) HandleGetRequest(w http.ResponseWriter, r *http.Request, t *i
 
 	if url, exists := cache.Get(requestPath); exists {
 		if c.ShouldServeFromRclone() {
-			http.Redirect(w, r, url, http.StatusFound)
+			redirect(w, r, url, c)
 		} else {
 			gf.streamFileToResponse(file, url, w, r, t, c, log)
 		}
@@ -112,7 +113,7 @@ func (gf *GetFile) HandleGetRequest(w http.ResponseWriter, r *http.Request, t *i
 		}
 		cache.Add(requestPath, resp.Download)
 		if c.ShouldServeFromRclone() {
-			http.Redirect(w, r, resp.Download, http.StatusFound)
+			redirect(w, r, resp.Download, c)
 		} else {
 			gf.streamFileToResponse(file, resp.Download, w, r, t, c, log)
 		}
@@ -174,8 +175,24 @@ func (gf *GetFile) playErrorVideo(link string, w http.ResponseWriter, r *http.Re
 		return
 	}
 	if c.ShouldServeFromRclone() {
-		http.Redirect(w, r, resp.Download, http.StatusFound)
-		return
+		redirect(w, r, resp.Download, c)
 	}
 	gf.streamFileToResponse(nil, resp.Download, w, r, t, c, log)
 }
+
+func redirect(w http.ResponseWriter, r *http.Request, url string, c config.ConfigInterface) {
+	prefHost := c.GetRandomPreferredHost()
+	if prefHost != "" {
+		url = replaceHostInURL(url, prefHost)
+	}
+	http.Redirect(w, r, url, http.StatusFound)
+}
+
+func replaceHostInURL(inputURL string, newHost string) string {
+	u, err := url.Parse(inputURL)
+	if err != nil {
+		return ""
+	}
+	u.Host = newHost
+	return u.String()
+}
diff --git a/pkg/http/client.go b/pkg/http/client.go
index 0be76ee..b06a1d2 100644
--- a/pkg/http/client.go
+++ b/pkg/http/client.go
@@ -23,8 +23,8 @@ type HTTPClient struct {
 func (r *HTTPClient) Do(req *http.Request) (*http.Response, error) {
 	if r.config != nil && strings.Contains(req.Host, "download.real-debrid.") {
 		prefHost := r.config.GetRandomPreferredHost()
-		if host := prefHost; host != "" {
-			req.Host = r.config.GetRandomPreferredHost()
+		if prefHost != "" {
+			req.Host = prefHost
 		}
 	}
 	if r.BearerToken != "" {