From 4dcd2e64bbcf79cca315fa4336e5fc093c339768 Mon Sep 17 00:00:00 2001
From: Ben Sarmiento <ben.sarmiento@forto.com>
Date: Fri, 8 Dec 2023 14:38:49 +0100
Subject: [PATCH] Verify downloads

---
 internal/config/types.go  |  6 ++++++
 internal/universal/get.go | 20 +++++++++++++++++---
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/internal/config/types.go b/internal/config/types.go
index bf08dff..7a3e923 100644
--- a/internal/config/types.go
+++ b/internal/config/types.go
@@ -19,6 +19,7 @@ type ConfigInterface interface {
 	EnableRetainRDTorrentName() bool
 	GetRandomPreferredHost() string
 	ShouldServeFromRclone() bool
+	ShouldVerifyDownloadLink() bool
 	ShouldForceIPv6() bool
 	GetRealDebridTimeout() int
 	GetRetriesUntilFailed() int
@@ -49,6 +50,7 @@ type ZurgConfig struct {
 	PreferredHosts        []string `yaml:"preferred_hosts" json:"preferred_hosts"`
 	NetworkBufferSize     int      `yaml:"network_buffer_size" json:"network_buffer_size"`
 	ServeFromRclone       bool     `yaml:"serve_from_rclone" json:"serve_from_rclone"`
+	VerifyDownloadLink    bool     `yaml:"verify_download_link" json:"verify_download_link"`
 	ForceIPv6             bool     `yaml:"force_ipv6" json:"force_ipv6"`
 
 	OnLibraryUpdate string `yaml:"on_library_update" json:"on_library_update"`
@@ -125,6 +127,10 @@ func (z *ZurgConfig) ShouldServeFromRclone() bool {
 	return z.ServeFromRclone
 }
 
+func (z *ZurgConfig) ShouldVerifyDownloadLink() bool {
+	return z.VerifyDownloadLink
+}
+
 func (z *ZurgConfig) ShouldForceIPv6() bool {
 	return z.ForceIPv6
 }
diff --git a/internal/universal/get.go b/internal/universal/get.go
index d8b5e53..4441e30 100644
--- a/internal/universal/get.go
+++ b/internal/universal/get.go
@@ -54,9 +54,16 @@ func (gf *GetFile) HandleGetRequest(directory, torrentName, fileName string, res
 	link := file.Link
 
 	if download, exists := torMgr.DownloadCache.Get(link); exists {
-		if cfg.ShouldServeFromRclone() && torMgr.Api.CanFetchFirstByte(download.Download) {
-			redirect(resp, req, download.Download, cfg)
-			return
+		if cfg.ShouldServeFromRclone() {
+			if cfg.ShouldVerifyDownloadLink() {
+				if torMgr.Api.CanFetchFirstByte(download.Download) {
+					redirect(resp, req, download.Download, cfg)
+					return
+				}
+			} else {
+				redirect(resp, req, download.Download, cfg)
+				return
+			}
 		} else {
 			err := gf.streamCachedLinkToResponse(download.Download, resp, req, torMgr, cfg, log)
 			if err == nil {
@@ -93,6 +100,13 @@ func (gf *GetFile) HandleGetRequest(directory, torrentName, fileName string, res
 		}
 		torMgr.DownloadCache.Set(link, unrestrict)
 		if cfg.ShouldServeFromRclone() {
+			if cfg.ShouldVerifyDownloadLink() {
+				if !torMgr.Api.CanFetchFirstByte(unrestrict.Download) {
+					log.Warnf("File %s is not available", fileName)
+					http.Error(resp, "File is not available", http.StatusNotFound)
+					return
+				}
+			}
 			redirect(resp, req, unrestrict.Download, cfg)
 		} else {
 			gf.streamFileToResponse(torrent, file, unrestrict, resp, req, torMgr, cfg, log)