Multi-token support

2024-06-28 04:47:43 +02:00
parent 5e06f04f33
commit 962845fb81
15 changed files with 214 additions and 108 deletions
--- a/internal/torrent/delete.go
+++ b/internal/torrent/delete.go
@@ -38,6 +38,6 @@ func (t *TorrentManager) Delete(accessKey string, deleteInRD bool) {
 }

 func (t *TorrentManager) DeleteByID(torrentID string) {
-	t.api.DeleteTorrent(torrentID)
+	t.rd.DeleteTorrent(torrentID)
 	t.deleteInfoFile(torrentID)
 }
--- a/internal/torrent/latestState.go
+++ b/internal/torrent/latestState.go
@@ -35,7 +35,7 @@ func (t *TorrentManager) setNewLatestState(checksum LibraryState) {
 func (t *TorrentManager) getCurrentState() LibraryState {
 	var state LibraryState

-	torrents, totalCount, err := t.api.GetTorrents(true)
+	torrents, totalCount, err := t.rd.GetTorrents(true)
 	if err != nil {
 		t.log.Errorf("Checksum API Error (GetTorrents): %v", err)
 		return LibraryState{}
@@ -45,7 +45,7 @@ func (t *TorrentManager) getCurrentState() LibraryState {
 		state.FirstTorrentId = torrents[0].ID
 	}

-	count, err := t.api.GetActiveTorrentCount()
+	count, err := t.rd.GetActiveTorrentCount()
 	if err != nil {
 		t.log.Errorf("Checksum API Error (GetActiveTorrentCount): %v", err)
 		return LibraryState{}
--- a/internal/torrent/manager.go
+++ b/internal/torrent/manager.go
@@ -30,14 +30,13 @@ type TorrentManager struct {
 	requiredVersion string

 	Config     config.ConfigInterface
-	api        *realdebrid.RealDebrid
+	rd         *realdebrid.RealDebrid
 	workerPool *ants.Pool
 	log        *logutil.Logger
 	repairLog  *logutil.Logger

-	DirectoryMap  cmap.ConcurrentMap[string, cmap.ConcurrentMap[string, *Torrent]] // directory -> accessKey -> Torrent
-	DownloadMap   cmap.ConcurrentMap[string, *realdebrid.Download]
-	UnrestrictMap cmap.ConcurrentMap[string, *realdebrid.Download]
+	DirectoryMap cmap.ConcurrentMap[string, cmap.ConcurrentMap[string, *Torrent]] // directory -> accessKey -> Torrent
+	DownloadMap  cmap.ConcurrentMap[string, *realdebrid.Download]

 	RootNode *fs.FileNode

@@ -67,14 +66,13 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, w
 		requiredVersion: "0.10.0",

 		Config:     cfg,
-		api:        api,
+		rd:         api,
 		workerPool: workerPool,
 		log:        log,
 		repairLog:  repairLog,

-		DirectoryMap:  cmap.New[cmap.ConcurrentMap[string, *Torrent]](),
-		DownloadMap:   cmap.New[*realdebrid.Download](),
-		UnrestrictMap: cmap.New[*realdebrid.Download](),
+		DirectoryMap: cmap.New[cmap.ConcurrentMap[string, *Torrent]](),
+		DownloadMap:  cmap.New[*realdebrid.Download](),

 		RootNode: fs.NewFileNode("root", true),

@@ -140,34 +138,13 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, w
 }

 // proxy function
-func (t *TorrentManager) UnrestrictFile(file *File, checkFirstByte bool) (*realdebrid.Download, error) {
+func (t *TorrentManager) UnrestrictFile(file *File) (*realdebrid.Download, error) {
 	if file.State.Is("deleted_file") {
 		return nil, fmt.Errorf("file %s has been deleted", file.Path)
 	} else if file.State.Is("broken_file") {
 		return nil, fmt.Errorf("file %s is broken", file.Path)
 	}
-	return t.UnrestrictLink(file.Link, checkFirstByte)
-}
-
-func (t *TorrentManager) UnrestrictLink(link string, verifyURL bool) (*realdebrid.Download, error) {
-	isRealDebrid := strings.HasPrefix(link, "https://real-debrid.com/d/")
-	if isRealDebrid && t.UnrestrictMap.Has(link[0:39]) {
-		ret, _ := t.UnrestrictMap.Get(link[0:39])
-		return ret, nil
-	} else if !isRealDebrid && t.UnrestrictMap.Has(link) {
-		ret, _ := t.UnrestrictMap.Get(link)
-		return ret, nil
-	}
-	ret, err := t.api.UnrestrictLink(link, verifyURL)
-	if err != nil {
-		return nil, err
-	}
-	if isRealDebrid {
-		t.UnrestrictMap.Set(ret.Link[0:39], ret)
-	} else {
-		t.UnrestrictMap.Set(ret.Link, ret)
-	}
-	return ret, nil
+	return t.rd.UnrestrictLink(file.Link)
 }

 func (t *TorrentManager) GetKey(torrent *Torrent) string {
@@ -242,7 +219,7 @@ func (t *TorrentManager) applyMediaInfoDetails(torrent *Torrent) error {
 		if file.MediaInfo != nil || !file.State.Is("ok_file") || !isPlayable {
 			return
 		}
-		unrestrict, err := t.UnrestrictFile(file, true)
+		unrestrict, err := t.UnrestrictFile(file)
 		if dlErr, ok := err.(*http.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
 			bwLimitReached = true
 			return
@@ -355,17 +332,29 @@ func (t *TorrentManager) deleteInfoFile(torrentID string) {
 /// end info functions

 func (t *TorrentManager) mountNewDownloads() {
-	downloads := t.api.GetDownloads()
+	token, _ := t.rd.GetToken()
+	var tokenMap cmap.ConcurrentMap[string, *realdebrid.Download]
+	if token != "" {
+		tokenMap, _ = t.rd.UnrestrictMap.Get(token)
+	}
+
+	downloads := t.rd.GetDownloads()
+	mountedCount := 0
 	for i := range downloads {
 		isRealDebrid := strings.HasPrefix(downloads[i].Link, "https://real-debrid.com/d/")
-		if isRealDebrid {
-			t.UnrestrictMap.SetIfAbsent(downloads[i].Link[0:39], &downloads[i])
-		} else {
-			t.UnrestrictMap.SetIfAbsent(downloads[i].Link, &downloads[i])
+		if !isRealDebrid {
 			filename := filepath.Base(downloads[i].Filename)
 			t.DownloadMap.Set(filename, &downloads[i])
+			mountedCount++
+		} else if token != "" {
+			tokenMap.Set(downloads[i].Link, &downloads[i])
 		}
 	}
+	if mountedCount > 0 {
+		t.log.Infof("Mounted %d new downloads", mountedCount)
+	} else {
+		t.log.Debugf("No new downloads to mount")
+	}
 }

 // StartDownloadsJob: permanent job for remounting downloads
--- a/internal/torrent/refresh.go
+++ b/internal/torrent/refresh.go
@@ -17,7 +17,7 @@ import (
 )

 func (t *TorrentManager) refreshTorrents(initialRun bool) {
-	instances, _, err := t.api.GetTorrents(false)
+	instances, _, err := t.rd.GetTorrents(false)
 	if err != nil {
 		t.log.Warnf("Cannot get torrents: %v", err)
 		return
@@ -173,7 +173,7 @@ func (t *TorrentManager) getMoreInfo(rdTorrent realdebrid.Torrent) *realdebrid.T
 	info := t.readInfoFromFile(rdTorrent.ID)
 	if info == nil {
 		var err error
-		info, err = t.api.GetTorrentInfo(rdTorrent.ID)
+		info, err = t.rd.GetTorrentInfo(rdTorrent.ID)
 		if err != nil {
 			t.log.Warnf("Cannot get info for id=%s: %v", rdTorrent.ID, err)
 			return nil
--- a/internal/torrent/repair.go
+++ b/internal/torrent/repair.go
@@ -196,7 +196,7 @@ func (t *TorrentManager) repair(torrent *Torrent, wg *sync.WaitGroup) {
 		if bwLimitReached || !file.State.Is("ok_file") {
 			return
 		}
-		_, err := t.UnrestrictFile(file, true)
+		_, err := t.UnrestrictFile(file)
 		if dlErr, ok := err.(*http.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
 			bwLimitReached = true
 			return
@@ -339,7 +339,7 @@ func (t *TorrentManager) assignLinks(torrent *Torrent) bool {
 	bwLimitReached := false
 	torrent.UnassignedLinks.Clone().Each(func(link string) bool {
 		// unrestrict each unassigned link that was filled out during torrent init
-		unrestrict, err := t.UnrestrictLink(link, true)
+		unrestrict, err := t.rd.UnrestrictLink(link)
 		if dlErr, ok := err.(*http.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
 			bwLimitReached = true
 			return true
@@ -481,7 +481,7 @@ func (t *TorrentManager) redownloadTorrent(torrent *Torrent, selection []string)
 	// redownload torrent
 	var newTorrentID string
 	prevState := t.latestState
-	resp, err := t.api.AddMagnetHash(torrent.Hash)
+	resp, err := t.rd.AddMagnetHash(torrent.Hash)
 	if err != nil {
 		if strings.Contains(err.Error(), "timeout") {
 			newState := t.getCurrentState()
@@ -523,14 +523,14 @@ func (t *TorrentManager) redownloadTorrent(torrent *Torrent, selection []string)
 			return nil, fmt.Errorf("cannot start redownloading: too many retries")
 		}

-		err = t.api.SelectTorrentFiles(newTorrentID, finalSelection)
+		err = t.rd.SelectTorrentFiles(newTorrentID, finalSelection)
 		if err != nil {
 			t.DeleteByID(newTorrentID)
 			return nil, fmt.Errorf("cannot start redownloading: %v", err)
 		}
 		time.Sleep(2 * time.Second)

-		info, err = t.api.GetTorrentInfo(newTorrentID)
+		info, err = t.rd.GetTorrentInfo(newTorrentID)
 		if err != nil {
 			t.DeleteByID(newTorrentID)
 			return nil, fmt.Errorf("cannot get info on redownloaded : %v", err)
@@ -568,7 +568,7 @@ func (t *TorrentManager) canCapacityHandle() bool {
 	const maxDelay = 60 * time.Second
 	retryCount := 0
 	for {
-		count, err := t.api.GetActiveTorrentCount()
+		count, err := t.rd.GetActiveTorrentCount()
 		if err != nil {
 			t.repairLog.Warnf("Cannot get active downloads count: %v", err)
 			if retryCount >= maxRetries {
@@ -679,7 +679,7 @@ func (t *TorrentManager) checkIfBroken(info *realdebrid.TorrentInfo, brokenFiles
 			if oldFile.ID != newFile.ID {
 				continue
 			}
-			if _, err := t.UnrestrictFile(selectedFiles[idx], true); err != nil {
+			if _, err := t.UnrestrictFile(selectedFiles[idx]); err != nil {
 				return err
 			}
 		}
--- a/internal/torrent/uncached.go
+++ b/internal/torrent/uncached.go
@@ -34,7 +34,7 @@ func (t *TorrentManager) GetUncachedTorrents() ([]*Torrent, error) {
 			break
 		}

-		resp, err := t.api.AvailabilityCheck(hashGroups[i].ToSlice())
+		resp, err := t.rd.AvailabilityCheck(hashGroups[i].ToSlice())
 		if err != nil {
 			return nil, fmt.Errorf("availability check is incomplete, skipping uncached check: %v", err)
 		}