Fix unrestrict issue

2023-11-28 00:41:15 +01:00
parent c8334ecb3b
commit 3d380e468f
9 changed files with 131 additions and 68 deletions
--- a/cmd/zurg/main.go
+++ b/cmd/zurg/main.go
@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"net/http"
 	"os"
-	"time"

 	"github.com/debridmediamanager.com/zurg/internal/config"
 	"github.com/debridmediamanager.com/zurg/internal/net"
@@ -16,7 +15,6 @@ import (
 	zurghttp "github.com/debridmediamanager.com/zurg/pkg/http"
 	"github.com/debridmediamanager.com/zurg/pkg/logutil"
 	"github.com/debridmediamanager.com/zurg/pkg/realdebrid"
-	"github.com/hashicorp/golang-lru/v2/expirable"

 	_ "net/http/pprof"
 )
@@ -43,8 +41,6 @@ func main() {
 		os.Exit(1)
 	}

-	cache := expirable.NewLRU[string, string](1e4, nil, time.Hour)
-
 	apiClient := zurghttp.NewHTTPClient(config.GetToken(), 5, 15, config, log.Named("httpclient"))

 	rd := realdebrid.NewRealDebrid(apiClient, log.Named("realdebrid"))
@@ -59,7 +55,6 @@ func main() {
 	torrentMgr := torrent.NewTorrentManager(config, rd, p, log.Named("manager"))

 	downloadClient := zurghttp.NewHTTPClient("", 5, 0, config, log.Named("dlclient"))
-
 	getfile := universal.NewGetFile(downloadClient)

 	go func() {
@@ -67,7 +62,7 @@ func main() {
 	}()

 	mux := http.NewServeMux()
-	net.Router(mux, getfile, config, torrentMgr, cache, log.Named("net"))
+	net.Router(mux, getfile, config, torrentMgr, log.Named("net"))

 	addr := fmt.Sprintf("%s:%s", config.GetHost(), config.GetPort())
 	server := &http.Server{Addr: addr, Handler: mux}
--- a/internal/net/router.go
+++ b/internal/net/router.go
@@ -10,24 +10,23 @@ import (
 	intHttp "github.com/debridmediamanager.com/zurg/internal/http"
 	"github.com/debridmediamanager.com/zurg/internal/torrent"
 	"github.com/debridmediamanager.com/zurg/internal/universal"
-	"github.com/hashicorp/golang-lru/v2/expirable"
 	"go.uber.org/zap"
 )

 // Router creates a WebDAV router
-func Router(mux *http.ServeMux, getfile *universal.GetFile, c config.ConfigInterface, t *torrent.TorrentManager, cache *expirable.LRU[string, string], log *zap.SugaredLogger) {
+func Router(mux *http.ServeMux, getfile *universal.GetFile, c config.ConfigInterface, t *torrent.TorrentManager, log *zap.SugaredLogger) {
 	mux.HandleFunc("/http/", func(w http.ResponseWriter, r *http.Request) {
 		switch r.Method {
 		case http.MethodGet:
 			requestPath := path.Clean(r.URL.Path)
 			if countNonEmptySegments(strings.Split(requestPath, "/")) > 3 {
-				getfile.HandleGetRequest(w, r, t, c, cache, log)
+				getfile.HandleGetRequest(w, r, t, c, log)
 			} else {
 				intHttp.HandleDirectoryListing(w, r, t, log)
 			}

 		case http.MethodHead:
-			universal.HandleHeadRequest(w, r, t, cache, log)
+			universal.HandleHeadRequest(w, r, t, log)

 		default:
 			log.Errorf("Request %s %s not supported yet", r.Method, r.URL.Path)
@@ -44,7 +43,7 @@ func Router(mux *http.ServeMux, getfile *universal.GetFile, c config.ConfigInter
 			dav.HandleDeleteRequest(w, r, t, log)

 		case http.MethodGet:
-			getfile.HandleGetRequest(w, r, t, c, cache, log)
+			getfile.HandleGetRequest(w, r, t, c, log)

 		case http.MethodOptions:
 			w.WriteHeader(http.StatusOK)
--- a/internal/torrent/manager.go
+++ b/internal/torrent/manager.go
@@ -27,6 +27,7 @@ const (

 type TorrentManager struct {
 	DirectoryMap    cmap.ConcurrentMap[string, cmap.ConcurrentMap[string, *Torrent]] // directory -> accessKey -> Torrent
+	DownloadCache   cmap.ConcurrentMap[string, *realdebrid.Download]
 	checksum        string
 	latestAdded     string
 	requiredVersion string
@@ -70,10 +71,35 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p
 		t.DirectoryMap.Set(directory, cmap.New[*Torrent]())
 	}

-	newTorrents, _, err := t.api.GetTorrents(0)
-	if err != nil {
-		t.log.Fatalf("Cannot get torrents: %v\n", err)
-	}
+	var initWait sync.WaitGroup
+	initWait.Add(2)
+
+	// Fetch downloads
+	go func() {
+		defer initWait.Done()
+		downloads, _, err := t.api.GetDownloads()
+		if err != nil {
+			t.log.Fatalf("Cannot get downloads: %v\n", err)
+		}
+		t.DownloadCache = cmap.New[*realdebrid.Download]()
+		for _, download := range downloads {
+			if !t.DownloadCache.Has(download.Link) {
+				t.DownloadCache.Set(download.Link, &download)
+			}
+		}
+	}()
+
+	// Fetch torrents
+	var newTorrents []realdebrid.Torrent
+	go func() {
+		defer initWait.Done()
+		newTorrents, _, err = t.api.GetTorrents(0)
+		if err != nil {
+			t.log.Fatalf("Cannot get torrents: %v\n", err)
+		}
+	}()
+
+	initWait.Wait()

 	torrentsChan := make(chan *Torrent, len(newTorrents))
 	var wg sync.WaitGroup
@@ -178,8 +204,8 @@ func (t *TorrentManager) mergeToMain(mainTorrent, torrentToMerge *Torrent) *Torr
 }

 // proxy
-func (t *TorrentManager) UnrestrictUntilOk(link string) *realdebrid.UnrestrictResponse {
-	retChan := make(chan *realdebrid.UnrestrictResponse, 1)
+func (t *TorrentManager) UnrestrictUntilOk(link string) *realdebrid.Download {
+	retChan := make(chan *realdebrid.Download, 1)
 	t.unrestrictPool.Submit(func() {
 		retChan <- t.api.UnrestrictUntilOk(link, t.cfg.ShouldServeFromRclone())
 		time.Sleep(1 * time.Second)
@@ -388,7 +414,7 @@ func (t *TorrentManager) getMoreInfo(rdTorrent realdebrid.Torrent) *Torrent {
 	var info *realdebrid.TorrentInfo
 	var err error
 	// file cache
-	torrentFromFile := t.readFromFile(rdTorrent.ID)
+	torrentFromFile := t.readTorrentFromFile(rdTorrent.ID)
 	if torrentFromFile != nil && len(torrentFromFile.ID) > 0 && len(torrentFromFile.Links) > 0 && len(torrentFromFile.Links) == len(rdTorrent.Links) && torrentFromFile.Links[0] == rdTorrent.Links[0] {
 		info = torrentFromFile
 		info.Progress = rdTorrent.Progress
@@ -440,7 +466,7 @@ func (t *TorrentManager) getMoreInfo(rdTorrent realdebrid.Torrent) *Torrent {
 	}

 	if len(selectedFiles) > 0 && torrentFromFile == nil {
-		t.writeToFile(info) // only when there are selected files, else it's useless
+		t.writeTorrentToFile(info) // only when there are selected files, else it's useless
 	}

 	infoCache.Set(rdTorrent.ID, &torrent)
@@ -462,7 +488,7 @@ func (t *TorrentManager) getName(name, originalName string) string {
 	}
 }

-func (t *TorrentManager) writeToFile(torrent *realdebrid.TorrentInfo) error {
+func (t *TorrentManager) writeTorrentToFile(torrent *realdebrid.TorrentInfo) error {
 	filePath := DATA_DIR + "/" + torrent.ID + ".bin"
 	file, err := os.Create(filePath)
 	if err != nil {
@@ -481,7 +507,7 @@ func (t *TorrentManager) writeToFile(torrent *realdebrid.TorrentInfo) error {
 	return nil
 }

-func (t *TorrentManager) readFromFile(torrentID string) *realdebrid.TorrentInfo {
+func (t *TorrentManager) readTorrentFromFile(torrentID string) *realdebrid.TorrentInfo {
 	filePath := DATA_DIR + "/" + torrentID + ".bin"
 	file, err := os.Open(filePath)
 	if err != nil {
@@ -506,7 +532,7 @@ func (t *TorrentManager) readFromFile(torrentID string) *realdebrid.TorrentInfo

 func (t *TorrentManager) organizeChaos(links []string, selectedFiles []*File) ([]*File, bool) {
 	type Result struct {
-		Response *realdebrid.UnrestrictResponse
+		Response *realdebrid.Download
 	}

 	resultsChan := make(chan Result, len(links))
--- a/internal/universal/get.go
+++ b/internal/universal/get.go
@@ -13,7 +13,6 @@ import (
 	intHttp "github.com/debridmediamanager.com/zurg/internal/http"
 	intTor "github.com/debridmediamanager.com/zurg/internal/torrent"
 	zurghttp "github.com/debridmediamanager.com/zurg/pkg/http"
-	"github.com/hashicorp/golang-lru/v2/expirable"
 	"go.uber.org/zap"
 )

@@ -26,7 +25,7 @@ func NewGetFile(client *zurghttp.HTTPClient) *GetFile {
 }

 // HandleGetRequest handles a GET request universally for both WebDAV and HTTP
-func (gf *GetFile) HandleGetRequest(w http.ResponseWriter, r *http.Request, t *intTor.TorrentManager, c config.ConfigInterface, cache *expirable.LRU[string, string], log *zap.SugaredLogger) {
+func (gf *GetFile) HandleGetRequest(w http.ResponseWriter, r *http.Request, t *intTor.TorrentManager, c config.ConfigInterface, log *zap.SugaredLogger) {
 	requestPath := path.Clean(r.URL.Path)
 	isDav := true
 	if strings.Contains(requestPath, "/http") {
@@ -71,15 +70,6 @@ func (gf *GetFile) HandleGetRequest(w http.ResponseWriter, r *http.Request, t *i
 		return
 	}

-	if url, exists := cache.Get(requestPath); exists {
-		if c.ShouldServeFromRclone() {
-			redirect(w, r, url, c)
-		} else {
-			gf.streamFileToResponse(file, url, w, r, t, c, log)
-		}
-		return
-	}
-
 	if !strings.HasPrefix(file.Link, "http") {
 		// This is a dead file, serve an alternate file
 		log.Warnf("File %s is not available", filename)
@@ -88,6 +78,15 @@ func (gf *GetFile) HandleGetRequest(w http.ResponseWriter, r *http.Request, t *i
 	}
 	link := file.Link

+	if download, exists := t.DownloadCache.Get(link); exists {
+		if c.ShouldServeFromRclone() {
+			redirect(w, r, download.Download, c)
+		} else {
+			gf.streamFileToResponse(file, download.Download, w, r, t, c, log)
+		}
+		return
+	}
+
 	resp := t.UnrestrictUntilOk(link)
 	if resp == nil {
 		// log.Warnf("File %s is no longer available, link %s", filepath.Base(file.Path), link)
@@ -112,7 +111,7 @@ func (gf *GetFile) HandleGetRequest(w http.ResponseWriter, r *http.Request, t *i
 				log.Warnf("Filename mismatch: %s and %s", filename, resp.Filename)
 			}
 		}
-		cache.Add(requestPath, resp.Download)
+		t.DownloadCache.Set(link, resp)
 		if c.ShouldServeFromRclone() {
 			redirect(w, r, resp.Download, c)
 		} else {
--- a/internal/universal/head.go
+++ b/internal/universal/head.go
@@ -8,7 +8,6 @@ import (
 	"strings"

 	"github.com/debridmediamanager.com/zurg/internal/torrent"
-	"github.com/hashicorp/golang-lru/v2/expirable"
 	"go.uber.org/zap"
 )

@@ -16,7 +15,7 @@ const (
 	SPLIT_TOKEN = "$"
 )

-func HandleHeadRequest(w http.ResponseWriter, r *http.Request, t *torrent.TorrentManager, cache *expirable.LRU[string, string], log *zap.SugaredLogger) {
+func HandleHeadRequest(w http.ResponseWriter, r *http.Request, t *torrent.TorrentManager, log *zap.SugaredLogger) {
 	requestPath := path.Clean(r.URL.Path)
 	requestPath = strings.Replace(requestPath, "/http", "", 1)
 	if requestPath == "/favicon.ico" {
@@ -31,18 +30,6 @@ func HandleHeadRequest(w http.ResponseWriter, r *http.Request, t *torrent.Torren
 		return
 	}

-	if data, exists := cache.Get("head:" + requestPath); exists {
-		splits := strings.Split(data, SPLIT_TOKEN)
-		contentType := splits[0]
-		contentLength := splits[1]
-		lastModified := splits[2]
-		w.Header().Set("Content-Type", contentType)
-		w.Header().Set("Content-Length", contentLength)
-		w.Header().Set("Last-Modified", lastModified)
-		w.WriteHeader(http.StatusOK)
-		return
-	}
-
 	baseDirectory := segments[len(segments)-3]
 	accessKey := segments[len(segments)-2]
 	filename := segments[len(segments)-1]
@@ -78,8 +65,6 @@ func HandleHeadRequest(w http.ResponseWriter, r *http.Request, t *torrent.Torren
 	w.Header().Set("Content-Type", contentType)
 	w.Header().Set("Content-Length", contentLength)
 	w.Header().Set("Last-Modified", lastModified)
-	cacheVal := strings.Join([]string{contentType, contentLength, lastModified}, SPLIT_TOKEN)
-	cache.Add("head:"+requestPath, cacheVal)
 	w.WriteHeader(http.StatusOK)
 }

--- a/pkg/realdebrid/api.go
+++ b/pkg/realdebrid/api.go
@@ -26,7 +26,7 @@ func NewRealDebrid(client *zurghttp.HTTPClient, log *zap.SugaredLogger) *RealDeb
 	}
 }

-func (rd *RealDebrid) UnrestrictCheck(link string) (*UnrestrictResponse, error) {
+func (rd *RealDebrid) UnrestrictCheck(link string) (*Download, error) {
 	data := url.Values{}
 	data.Set("link", link)

@@ -50,7 +50,7 @@ func (rd *RealDebrid) UnrestrictCheck(link string) (*UnrestrictResponse, error)
 		return nil, err
 	}

-	var response UnrestrictResponse
+	var response Download
 	err = json.Unmarshal(body, &response)
 	if err != nil {
 		rd.log.Errorf("Error when decoding unrestrict check JSON: %v", err)
@@ -62,7 +62,7 @@ func (rd *RealDebrid) UnrestrictCheck(link string) (*UnrestrictResponse, error)
 }

 // GetTorrents returns all torrents, paginated
-// if customLimit is 0, the default limit of 2500 is used
+// if customLimit is 0, the default limit of 1000 is used
 func (rd *RealDebrid) GetTorrents(customLimit int) ([]Torrent, int, error) {
 	baseURL := "https://api.real-debrid.com/rest/1.0/torrents"
 	var allTorrents []Torrent
@@ -116,6 +116,8 @@ func (rd *RealDebrid) GetTorrents(customLimit int) ([]Torrent, int, error) {
 			break
 		}

+		rd.log.Debugf("Got %d torrents (page %d), total count is %d", len(allTorrents), page, totalCount)
+
 		page++
 	}
 	return allTorrents, totalCount, nil
@@ -263,7 +265,8 @@ func (rd *RealDebrid) GetActiveTorrentCount() (*ActiveTorrentCountResponse, erro
 	return &response, nil
 }

-func (rd *RealDebrid) UnrestrictLink(link string, checkFirstByte bool) (*UnrestrictResponse, error) {
+func (rd *RealDebrid) UnrestrictLink(link string, checkFirstByte bool) (*Download, error) {
+	fmt.Println("Unrestricting link via api", link)
 	data := url.Values{}
 	data.Set("link", link)

@@ -294,7 +297,7 @@ func (rd *RealDebrid) UnrestrictLink(link string, checkFirstByte bool) (*Unrestr
 		return nil, fmt.Errorf("unreadable body so likely it has expired")
 	}

-	var response UnrestrictResponse
+	var response Download
 	err = json.Unmarshal(body, &response)
 	if err != nil {
 		// rd.log.Errorf("Error when decoding unrestrict link JSON: %v", err)
@@ -308,3 +311,61 @@ func (rd *RealDebrid) UnrestrictLink(link string, checkFirstByte bool) (*Unrestr
 	// rd.log.Debugf("Unrestricted link %s into %s", link, response.Download)
 	return &response, nil
 }
+
+// GetDownloads returns all torrents, paginated
+func (rd *RealDebrid) GetDownloads() ([]Download, int, error) {
+	baseURL := "https://api.real-debrid.com/rest/1.0/downloads"
+	var allDownloads []Download
+	page := 1
+	limit := 1000
+	totalCount := 0
+
+	for {
+		params := url.Values{}
+		params.Set("page", fmt.Sprintf("%d", page))
+		params.Set("limit", fmt.Sprintf("%d", limit))
+		// params.Set("filter", "active")
+
+		reqURL := baseURL + "?" + params.Encode()
+
+		req, err := http.NewRequest("GET", reqURL, nil)
+		if err != nil {
+			rd.log.Errorf("Error when creating a get downloads request: %v", err)
+			return nil, 0, err
+		}
+
+		resp, err := rd.client.Do(req)
+		if err != nil {
+			rd.log.Errorf("Error when executing the get downloads request: %v", err)
+			return nil, 0, err
+		}
+		defer resp.Body.Close()
+
+		// if status code is not 2xx, return erro
+
+		var downloads []Download
+		decoder := json.NewDecoder(resp.Body)
+		err = decoder.Decode(&downloads)
+		if err != nil {
+			rd.log.Errorf("Error when decoding get downloads JSON: %v", err)
+			return nil, 0, err
+		}
+
+		allDownloads = append(allDownloads, downloads...)
+
+		totalCountHeader := resp.Header.Get("x-total-count")
+		totalCount, err = strconv.Atoi(totalCountHeader)
+		if err != nil {
+			break
+		}
+
+		if len(allDownloads) >= totalCount {
+			break
+		}
+
+		rd.log.Debugf("Got %d downloads (page %d), total count is %d", len(allDownloads), page, totalCount)
+
+		page++
+	}
+	return allDownloads, totalCount, nil
+}
--- a/pkg/realdebrid/types.go
+++ b/pkg/realdebrid/types.go
@@ -11,12 +11,12 @@ type FileJSON struct {
 	Link     string `json:"link"`
 }

-type UnrestrictResponse struct {
+type Download struct {
 	Filename   string `json:"filename"`
-	Filesize   int64  `json:"filesize"`
-	Link       string `json:"link"`
-	Host       string `json:"host"`
-	Download   string `json:"download,omitempty"`
+	Filesize   int64  `json:"filesize"` // bytes, 0 if unknown
+	Link       string `json:"link"`     // Original link
+	Host       string `json:"host"`     // Host main domain
+	Download   string `json:"download"` // Generated link
 	Streamable int    `json:"streamable"`
 }

--- a/pkg/realdebrid/unrestrict.go
+++ b/pkg/realdebrid/unrestrict.go
@@ -5,7 +5,7 @@ import (
 	"strings"
 )

-func (rd *RealDebrid) UnrestrictUntilOk(link string, serveFromRclone bool) *UnrestrictResponse {
+func (rd *RealDebrid) UnrestrictUntilOk(link string, serveFromRclone bool) *Download {
 	if !strings.HasPrefix(link, "http") {
 		return nil
 	}
--- a/scan_test.sh
+++ b/scan_test.sh
@@ -1,14 +1,12 @@
 #!/bin/bash

-process_file() {
-    echo "Processing $1"
-    # echo -n "First 100 bytes of $file: "
-    #  | od -An -t x1
-    # dd bs=1 count=100 if="$file" 2>/dev/null
+mock_scanner() {
+    echo "Processing $1..."
    dd bs=1 count=100 if="$1" >/dev/null 2>&1
+    echo "Processed $1"
 }

-export -f process_file
+export -f mock_scanner

 if [ "$#" -ne 1 ]; then
    echo "Usage: $0 directory"
@@ -20,4 +18,4 @@ if [ ! -d "$1" ]; then
    exit 1
 fi

-find "$1" -type f -print0 | xargs -0 -n1 -P20 -I{} bash -c 'process_file "$@"' _ {}
+find "$1" -type f -print0 | xargs -0 -n1 -P20 -I{} bash -c 'mock_scanner "$@"' _ {}