From a42a5eeb80d6203d0a8546c841e19376e4658e15 Mon Sep 17 00:00:00 2001
From: Ben Sarmiento <ben.sarmiento@forto.com>
Date: Tue, 28 Nov 2023 11:39:08 +0100
Subject: [PATCH] Additional configs

---
 cmd/zurg/main.go          |  4 ++--
 config.example.yml        |  9 +++++----
 internal/config/types.go  | 18 ++++++++++++++++++
 internal/universal/get.go |  1 +
 pkg/http/client.go        |  5 ++++-
 5 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/cmd/zurg/main.go b/cmd/zurg/main.go
index d132d50..dd8c169 100644
--- a/cmd/zurg/main.go
+++ b/cmd/zurg/main.go
@@ -39,7 +39,7 @@ func main() {
 		os.Exit(1)
 	}
 
-	apiClient := zurghttp.NewHTTPClient(config.GetToken(), 5, 15, config, log.Named("httpclient"))
+	apiClient := zurghttp.NewHTTPClient(config.GetToken(), config.GetRetriesUntilFailed(), config.GetRealDebridTimeout(), config, log.Named("httpclient"))
 
 	rd := realdebrid.NewRealDebrid(apiClient, log.Named("realdebrid"))
 
@@ -52,7 +52,7 @@ func main() {
 
 	torrentMgr := torrent.NewTorrentManager(config, rd, p, log.Named("manager"))
 
-	downloadClient := zurghttp.NewHTTPClient("", 5, 0, config, log.Named("dlclient"))
+	downloadClient := zurghttp.NewHTTPClient(config.GetToken(), config.GetRetriesUntilFailed(), 0, config, log.Named("dlclient"))
 	getfile := universal.NewGetFile(downloadClient)
 
 	mux := http.NewServeMux()
diff --git a/config.example.yml b/config.example.yml
index b7b504c..d10af00 100644
--- a/config.example.yml
+++ b/config.example.yml
@@ -1,8 +1,7 @@
 # Zurg configuration version
 zurg: v1
-
 token: YOUR_RD_API_TOKEN # https://real-debrid.com/apitoken
-
+# basic functionality
 host: "[::]" # do not change this if you are running it inside a docker container
 port: 9999 # do not change this if you are running it inside a docker container
 concurrent_workers: 200
@@ -11,7 +10,9 @@ check_for_changes_every_secs: 15
 unrestrict_workers: 10 # since unrestricting has a different rate limit, use a different worker pool. decrease this if you are getting 429s
 release_unrestrict_after_ms: 100 # wait time for every unrestrict worker to be released. increase this if you are getting 429s
 rate_limit_sleep_secs: 6 # wait time after getting a 429 from Real-Debrid API
-
+realdebrid_timeout_secs: 60 # api timeout
+retries_until_failed: 5 # api failures until considered failed
+# misc configs
 enable_repair: true # BEWARE! THERE CAN ONLY BE 1 INSTANCE OF ZURG THAT SHOULD REPAIR YOUR TORRENTS
 retain_folder_name_extension: false # if true, zurg won't modify the filenames from real-debrid
 retain_rd_torrent_name: false # if true, it will strictly follow RD API returned torrent name which should make this more compatible with rdt-client
@@ -20,7 +21,7 @@ on_library_update: |
   do
       echo "detected update on: $arg"
   done
-
+# network configs
 network_buffer_size: 1048576 # 1 MiB
 serve_from_rclone: false
 force_ipv6: true
diff --git a/internal/config/types.go b/internal/config/types.go
index be9888f..21e5ac8 100644
--- a/internal/config/types.go
+++ b/internal/config/types.go
@@ -23,6 +23,8 @@ type ConfigInterface interface {
 	GetUnrestrictWorkers() int
 	GetReleaseUnrestrictAfterMs() int
 	GetRateLimitSleepSeconds() int
+	GetRealDebridTimeout() int
+	GetRetriesUntilFailed() int
 }
 
 type ZurgConfig struct {
@@ -44,6 +46,8 @@ type ZurgConfig struct {
 	PreferredHosts            []string `yaml:"preferred_hosts"`
 	ServeFromRclone           bool     `yaml:"serve_from_rclone"`
 	ForceIPv6                 bool     `yaml:"force_ipv6"`
+	RealDebridTimeout         int      `yaml:"realdebrid_timeout_secs"`
+	RetriesUntilFailed        int      `yaml:"retries_until_failed"`
 }
 
 func (z *ZurgConfig) GetToken() string {
@@ -144,3 +148,17 @@ func (z *ZurgConfig) GetRateLimitSleepSeconds() int {
 	}
 	return z.RateLimitSleepSeconds
 }
+
+func (z *ZurgConfig) GetRealDebridTimeout() int {
+	if z.RealDebridTimeout == 0 {
+		return 60
+	}
+	return z.RealDebridTimeout
+}
+
+func (z *ZurgConfig) GetRetriesUntilFailed() int {
+	if z.RetriesUntilFailed == 0 {
+		return 5
+	}
+	return z.RetriesUntilFailed
+}
diff --git a/internal/universal/get.go b/internal/universal/get.go
index 9786258..ab11a78 100644
--- a/internal/universal/get.go
+++ b/internal/universal/get.go
@@ -82,6 +82,7 @@ func (gf *GetFile) HandleGetRequest(w http.ResponseWriter, r *http.Request, t *i
 	if download, exists := t.DownloadCache.Get(link); exists {
 		if c.ShouldServeFromRclone() && t.Api.CanFetchFirstByte(download.Download) {
 			redirect(w, r, download.Download, c)
+			return
 		} else {
 			err := gf.streamCachedLinkToResponse(download.Download, w, r, t, c, log)
 			if err == nil {
diff --git a/pkg/http/client.go b/pkg/http/client.go
index b4ad99b..1c09359 100644
--- a/pkg/http/client.go
+++ b/pkg/http/client.go
@@ -88,6 +88,9 @@ func (r *HTTPClient) Do(req *http.Request) (*http.Response, error) {
 				}
 				time.Sleep(r.Backoff(attempt))
 			}
+			if resp != nil {
+				resp.Body.Close()
+			}
 		}
 	}
 }
@@ -109,7 +112,7 @@ func NewHTTPClient(token string, maxRetries int, timeoutSecs int, cfg config.Con
 		},
 		ShouldRetry: func(resp *http.Response, err error) int {
 			if resp != nil {
-				if resp.StatusCode == 429 || resp.StatusCode == 400 {
+				if resp.StatusCode == 429 || resp.StatusCode == 400 || resp.StatusCode == 403 {
 					return 0 // retry but don't increment attempt
 				}
 				return -1 // don't retry