Add token management

internal/app.go

@@ -121,7 +121,7 @@ func MainApp(configPath string) {
 	}
 	defer workerPool.Release()
 
-	api := realdebrid.NewRealDebrid(
+	rd := realdebrid.NewRealDebrid(
 		apiClient,
 		unrestrictClient,
 		downloadClient,
@@ -132,7 +132,7 @@ func MainApp(configPath string) {
 
 	premium.MonitorPremiumStatus(
 		workerPool,
-		api,
+		rd,
 		zurglog,
 	)
 
@@ -145,14 +145,14 @@ func MainApp(configPath string) {
 
 	torrentMgr := torrent.NewTorrentManager(
 		config,
-		api,
+		rd,
 		workerPool,
 		hasFFprobe,
 		log.Named("manager"),
 		log.Named("repair"),
 	)
 
-	downloader := universal.NewDownloader(downloadClient, workerPool)
+	downloader := universal.NewDownloader(rd, workerPool)
 
 	router := chi.NewRouter()
 	handlers.AttachHandlers(
@@ -160,7 +160,7 @@ func MainApp(configPath string) {
 		downloader,
 		torrentMgr,
 		config,
-		api,
+		rd,
 		workerPool,
 		hosts,
 		log.Named("router"),

internal/handlers/home.go

@@ -47,7 +47,7 @@ type RootResponse struct {
 }
 
 func (zr *Handlers) generateResponse(resp http.ResponseWriter, req *http.Request) (*RootResponse, error) {
-	userInfo, err := zr.api.GetUserInformation()
+	userInfo, err := zr.rd.GetUserInformation()
 	if err != nil {
 		http.Error(resp, err.Error(), http.StatusInternalServerError)
 		return nil, err
@@ -77,7 +77,7 @@ func (zr *Handlers) generateResponse(resp http.ResponseWriter, req *http.Request
 	sortedIDs := zr.torMgr.OnceDoneBin.ToSlice()
 	sort.Strings(sortedIDs)
 
-	trafficDetails, err := zr.api.GetTrafficDetails()
+	trafficDetails, err := zr.rd.GetTrafficDetails()
 	if err != nil {
 		http.Error(resp, err.Error(), http.StatusInternalServerError)
 		return nil, err

internal/handlers/router.go

@@ -24,7 +24,7 @@ type Handlers struct {
 	downloader       *universal.Downloader
 	torMgr           *torrent.TorrentManager
 	cfg              config.ConfigInterface
-	api              *realdebrid.RealDebrid
+	rd               *realdebrid.RealDebrid
 	workerPool       *ants.Pool
 	hosts            []string
 	trafficOnStartup atomic.Uint64
@@ -37,18 +37,18 @@ func init() {
 	chi.RegisterMethod("MOVE")
 }
 
-func AttachHandlers(router *chi.Mux, downloader *universal.Downloader, torMgr *torrent.TorrentManager, cfg config.ConfigInterface, api *realdebrid.RealDebrid, workerPool *ants.Pool, hosts []string, log *logutil.Logger) {
+func AttachHandlers(router *chi.Mux, downloader *universal.Downloader, torMgr *torrent.TorrentManager, cfg config.ConfigInterface, rd *realdebrid.RealDebrid, workerPool *ants.Pool, hosts []string, log *logutil.Logger) {
 	hs := &Handlers{
 		downloader: downloader,
 		torMgr:     torMgr,
 		cfg:        cfg,
-		api:        api,
+		rd:         rd,
 		workerPool: workerPool,
 		hosts:      hosts,
 		log:        log,
 	}
 
-	trafficDetails, err := api.GetTrafficDetails()
+	trafficDetails, err := rd.GetTrafficDetails()
 	if err != nil {
 		log.Errorf("Failed to get traffic details: %v", err)
 		trafficDetails = make(map[string]int64)

internal/torrent/manager.go

@@ -12,7 +12,6 @@ import (
 
 	"github.com/debridmediamanager/zurg/internal/config"
 	"github.com/debridmediamanager/zurg/internal/fs"
-	"github.com/debridmediamanager/zurg/pkg/http"
 	"github.com/debridmediamanager/zurg/pkg/logutil"
 	"github.com/debridmediamanager/zurg/pkg/realdebrid"
 	"github.com/debridmediamanager/zurg/pkg/utils"
@@ -144,7 +143,7 @@ func (t *TorrentManager) UnrestrictFile(file *File) (*realdebrid.Download, error
 	} else if file.State.Is("broken_file") {
 		return nil, fmt.Errorf("file %s is broken", file.Path)
 	}
-	return t.rd.UnrestrictLink(file.Link)
+	return t.rd.UnrestrictAndVerify(file.Link)
 }
 
 func (t *TorrentManager) GetKey(torrent *Torrent) string {
@@ -220,7 +219,7 @@ func (t *TorrentManager) applyMediaInfoDetails(torrent *Torrent) error {
 			return
 		}
 		unrestrict, err := t.UnrestrictFile(file)
-		if dlErr, ok := err.(*http.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
+		if utils.IsBWLimitExceeded(err) {
 			bwLimitReached = true
 			return
 		}
@@ -332,11 +331,8 @@ func (t *TorrentManager) deleteInfoFile(torrentID string) {
 /// end info functions
 
 func (t *TorrentManager) mountNewDownloads() {
-	token, _ := t.rd.GetToken()
+	token := t.Config.GetToken()
-	var tokenMap cmap.ConcurrentMap[string, *realdebrid.Download]
+	tokenMap, _ := t.rd.UnrestrictMap.Get(token)
-	if token != "" {
-		tokenMap, _ = t.rd.UnrestrictMap.Get(token)
-	}
 
 	downloads := t.rd.GetDownloads()
 	mountedCount := 0

internal/torrent/repair.go

@@ -9,7 +9,6 @@ import (
 	"time"
 
 	"github.com/debridmediamanager/zurg/internal/config"
-	"github.com/debridmediamanager/zurg/pkg/http"
 	"github.com/debridmediamanager/zurg/pkg/realdebrid"
 	"github.com/debridmediamanager/zurg/pkg/utils"
 	mapset "github.com/deckarep/golang-set/v2"
@@ -197,7 +196,7 @@ func (t *TorrentManager) repair(torrent *Torrent, wg *sync.WaitGroup) {
 			return
 		}
 		_, err := t.UnrestrictFile(file)
-		if dlErr, ok := err.(*http.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
+		if utils.IsBWLimitExceeded(err) {
 			bwLimitReached = true
 			return
 		}
@@ -253,7 +252,7 @@ func (t *TorrentManager) repair(torrent *Torrent, wg *sync.WaitGroup) {
 	info, err := t.redownloadTorrent(torrent, []string{}) // reinsert the whole torrent, passing empty selection
 	if info != nil && info.Progress == 100 {
 		err = t.checkIfBroken(info, brokenFiles)
-		if dlErr, ok := err.(*http.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
+		if utils.IsBWLimitExceeded(err) {
 			t.repairLog.Warnf("Your account has reached the bandwidth limit, cannot continue repairing torrent %s", t.GetKey(torrent))
 			return
 		}
@@ -339,8 +338,8 @@ 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.rd.UnrestrictLink(link)
+		unrestrict, err := t.rd.UnrestrictAndVerify(link)
-		if dlErr, ok := err.(*http.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
+		if utils.IsBWLimitExceeded(err) {
 			bwLimitReached = true
 			return true
 		}

internal/universal/downloader.go

@@ -13,22 +13,22 @@ import (
 
 	"github.com/debridmediamanager/zurg/internal/config"
 	intTor "github.com/debridmediamanager/zurg/internal/torrent"
-	zurghttp "github.com/debridmediamanager/zurg/pkg/http"
 	"github.com/debridmediamanager/zurg/pkg/logutil"
 	"github.com/debridmediamanager/zurg/pkg/realdebrid"
+	"github.com/debridmediamanager/zurg/pkg/utils"
 	"github.com/panjf2000/ants/v2"
 )
 
 type Downloader struct {
-	client         *zurghttp.HTTPClient
+	rd             *realdebrid.RealDebrid
 	workerPool     *ants.Pool
 	RequestedBytes atomic.Uint64
 	TotalBytes     atomic.Uint64
 }
 
-func NewDownloader(client *zurghttp.HTTPClient, workerPool *ants.Pool) *Downloader {
+func NewDownloader(rd *realdebrid.RealDebrid, workerPool *ants.Pool) *Downloader {
 	dl := &Downloader{
-		client:     client,
+		rd:         rd,
 		workerPool: workerPool,
 	}
 
@@ -42,10 +42,12 @@ func NewDownloader(client *zurghttp.HTTPClient, workerPool *ants.Pool) *Download
 	nextMidnightInCET := time.Date(tomorrow.Year(), tomorrow.Month(), tomorrow.Day(), 0, 0, 0, 0, cetTZ)
 	duration := nextMidnightInCET.Sub(now)
 	timer := time.NewTimer(duration)
+	// permanent job for bandwidth reset
 	workerPool.Submit(func() {
 		<-timer.C
 		ticker := time.NewTicker(24 * time.Hour)
 		for {
+			rd.TokenManager.ResetAllTokens()
 			dl.RequestedBytes.Store(0)
 			dl.TotalBytes.Store(0)
 			<-ticker.C
@@ -88,7 +90,7 @@ func (dl *Downloader) DownloadFile(
 	}
 
 	unrestrict, err := torMgr.UnrestrictFile(file)
-	if dlErr, ok := err.(*zurghttp.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
+	if utils.IsBWLimitExceeded(err) {
 		// log.Errorf("Your account has reached the bandwidth limit, please try again after 12AM CET")
 		http.Error(resp, "File is not available (bandwidth limit reached)", http.StatusBadRequest)
 		return
@@ -161,8 +163,8 @@ func (dl *Downloader) streamFileToResponse(
 		dlReq.Header.Add("Range", req.Header.Get("Range"))
 	}
 
-	downloadResp, err := dl.client.Do(dlReq)
+	downloadResp, err := dl.rd.DownloadFile(dlReq)
-	if dlErr, ok := err.(*zurghttp.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
+	if utils.IsBWLimitExceeded(err) {
 		// log.Errorf("Your account has reached the bandwidth limit, please try again after 12AM CET")
 		http.Error(resp, "File is not available (bandwidth limit reached)", http.StatusBadRequest)
 		return

pkg/realdebrid/api.go

@@ -11,18 +11,19 @@ import (
 	"github.com/debridmediamanager/zurg/internal/config"
 	zurghttp "github.com/debridmediamanager/zurg/pkg/http"
 	"github.com/debridmediamanager/zurg/pkg/logutil"
+	"github.com/debridmediamanager/zurg/pkg/utils"
 	cmap "github.com/orcaman/concurrent-map/v2"
 	"github.com/panjf2000/ants/v2"
 )
 
 type RealDebrid struct {
-	torrentsCache    []Torrent
 	UnrestrictMap    cmap.ConcurrentMap[string, cmap.ConcurrentMap[string, *Download]]
+	TokenManager     *DownloadTokenManager
+	torrentsCache    []Torrent
 	verifiedLinks    cmap.ConcurrentMap[string, int64]
 	apiClient        *zurghttp.HTTPClient
 	unrestrictClient *zurghttp.HTTPClient
 	downloadClient   *zurghttp.HTTPClient
-	tokenManager     *DownloadTokenManager
 	workerPool       *ants.Pool
 	cfg              config.ConfigInterface
 	log              *logutil.Logger
@@ -36,13 +37,13 @@ func NewRealDebrid(apiClient, unrestrictClient, downloadClient *zurghttp.HTTPCli
 	}
 
 	rd := &RealDebrid{
-		torrentsCache:    []Torrent{},
 		UnrestrictMap:    cmap.New[cmap.ConcurrentMap[string, *Download]](),
+		TokenManager:     NewDownloadTokenManager(downloadTokens),
+		torrentsCache:    []Torrent{},
 		verifiedLinks:    cmap.New[int64](),
 		apiClient:        apiClient,
 		unrestrictClient: unrestrictClient,
 		downloadClient:   downloadClient,
-		tokenManager:     NewDownloadTokenManager(downloadTokens),
 		workerPool:       workerPool,
 		cfg:              cfg,
 		log:              log,
@@ -59,76 +60,59 @@ func NewRealDebrid(apiClient, unrestrictClient, downloadClient *zurghttp.HTTPCli
 	return rd
 }
 
-// currently unused
+func (rd *RealDebrid) UnrestrictAndVerify(link string) (*Download, error) {
-func (rd *RealDebrid) UnrestrictCheck(link string) (*Download, error) {
-	data := url.Values{}
-	data.Set("link", link)
-	requestBody := strings.NewReader(data.Encode())
-
-	req, err := http.NewRequest(http.MethodPost, "https://api.real-debrid.com/rest/1.0/unrestrict/check", requestBody)
-	if err != nil {
-		rd.log.Errorf("Error when creating a unrestrict check request: %v", err)
-		return nil, err
-	}
-	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
-
-	resp, err := rd.unrestrictClient.Do(req)
-	if err != nil {
-		rd.log.Errorf("Error when executing the unrestrict check request: %v", err)
-		return nil, err
-	}
-	defer resp.Body.Close()
-
-	body, err := io.ReadAll(resp.Body)
-	if err != nil {
-		rd.log.Errorf("Error when reading the body of unrestrict check response: %v", err)
-		return nil, err
-	}
-
-	var response Download
-	err = json.Unmarshal(body, &response)
-	if err != nil {
-		rd.log.Errorf("Error when decoding unrestrict check JSON: %v", err)
-		return nil, err
-	}
-
-	rd.log.Debugf("Link %s is streamable? %v", response.Streamable)
-	return &response, nil
-}
-
-func (rd *RealDebrid) UnrestrictLink(link string) (*Download, error) {
 	for {
-		token, err := rd.tokenManager.GetCurrentToken()
+		token, err := rd.TokenManager.GetCurrentToken()
 		if err != nil {
 			// when all tokens are expired
 			return nil, err
 		}
-		download, err := rd.UnrestrictLinkWithToken(token, link)
+
-		if dlErr, ok := err.(*zurghttp.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
+		// check if the link is already unrestricted
-			rd.tokenManager.SetCurrentTokenExpired()
+		tokenMap, ok := rd.UnrestrictMap.Get(token)
+		if ok && tokenMap.Has(link) {
+			download, _ := tokenMap.Get(link)
+			// check if the link is in the verified links cache
+			if expiry, ok := rd.verifiedLinks.Get(download.ID); ok && expiry > time.Now().Unix() {
+				return download, nil
+			}
+
+			err := rd.downloadClient.VerifyLink(download.Download)
+			if utils.IsBWLimitExceeded(err) {
+				rd.TokenManager.SetTokenAsExpired(token)
 				continue
 			}
+			if err != nil {
+				return nil, err
+			}
+			rd.verifiedLinks.Set(download.ID, time.Now().Unix()+60*60*24)
+
+			return download, nil
+		}
+
+		download, err := rd.UnrestrictLink(link)
+		if err != nil {
+			return nil, err
+		}
+
+		tokenMap.Set(link, download)
+
+		err = rd.downloadClient.VerifyLink(download.Download)
+		if utils.IsBWLimitExceeded(err) {
+			rd.TokenManager.SetTokenAsExpired(token)
+			continue
+		}
+		if err != nil {
+			return nil, err
+		}
+
+		rd.verifiedLinks.Set(download.ID, time.Now().Unix()+60*60*24)
+
 		return download, err
 	}
 }
 
-func (rd *RealDebrid) UnrestrictLinkWithToken(token, link string) (*Download, error) {
+func (rd *RealDebrid) UnrestrictLink(link string) (*Download, error) {
-	// check if the link is already unrestricted
-	if tokenMap, ok := rd.UnrestrictMap.Get(token); ok {
-		if d, ok := tokenMap.Get(link); ok {
-			// check if the link is in the verified links cache
-			if expiry, ok := rd.verifiedLinks.Get(d.Download); ok && expiry > time.Now().Unix() {
-				return d, nil
-			}
-			err := rd.downloadClient.VerifyLink(d.Download)
-			if err != nil {
-				return nil, err
-			}
-			rd.verifiedLinks.Set(d.Download, time.Now().Unix()+60*60*24)
-			return d, nil
-		}
-	}
-
 	data := url.Values{}
 	if strings.HasPrefix(link, "https://real-debrid.com/d/") {
 		// set link to max 39 chars (26 + 13)
@@ -145,7 +129,6 @@ func (rd *RealDebrid) UnrestrictLinkWithToken(token, link string) (*Download, er
 
 	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
 
-	rd.unrestrictClient.SetToken(token)
 	// at this point, any errors mean that the link has expired and we need to repair it
 	resp, err := rd.unrestrictClient.Do(req)
 	if err != nil {
@@ -167,16 +150,6 @@ func (rd *RealDebrid) UnrestrictLinkWithToken(token, link string) (*Download, er
 		return nil, fmt.Errorf("undecodable response: %v", err)
 	}
 
-	tokenMap, _ := rd.UnrestrictMap.Get(token)
-	tokenMap.Set(link, &response)
-
-	err = rd.downloadClient.VerifyLink(response.Download)
-	if err != nil {
-		return nil, err
-	}
-
-	rd.verifiedLinks.Set(response.Download, time.Now().Unix()+60*60*24)
-
 	// rd.log.Debugf("Unrestricted link %s into %s", link, response.Download)
 	return &response, nil
 }
@@ -428,6 +401,6 @@ func (rd *RealDebrid) AvailabilityCheck(hashes []string) (AvailabilityResponse,
 	return response, nil
 }
 
-func (rd *RealDebrid) GetToken() (string, error) {
+func (rd *RealDebrid) DownloadFile(req *http.Request) (*http.Response, error) {
-	return rd.tokenManager.GetCurrentToken()
+	return rd.downloadClient.Do(req)
 }

pkg/realdebrid/token_manager.go

@@ -13,7 +13,7 @@ type Token struct {
 type DownloadTokenManager struct {
 	tokens  []Token
 	current int
-	mu      sync.Mutex
+	mu      sync.RWMutex
 }
 
 // NewDownloadTokenManager initializes a new DownloadTokenManager with the given tokens.
@@ -27,8 +27,8 @@ func NewDownloadTokenManager(tokenStrings []string) *DownloadTokenManager {
 
 // GetCurrentToken returns the current non-expired token.
 func (dtm *DownloadTokenManager) GetCurrentToken() (string, error) {
-	dtm.mu.Lock()
+	dtm.mu.RLock()
-	defer dtm.mu.Unlock()
+	defer dtm.mu.RUnlock()
 
 	for {
 		if !dtm.tokens[dtm.current].expired {
@@ -43,13 +43,19 @@ func (dtm *DownloadTokenManager) GetCurrentToken() (string, error) {
 	}
 }
 
-// SetCurrentTokenExpired sets the current token as expired.
+// SetTokenAsExpired sets the specified token as expired.
-func (dtm *DownloadTokenManager) SetCurrentTokenExpired() {
+func (dtm *DownloadTokenManager) SetTokenAsExpired(token string) error {
 	dtm.mu.Lock()
 	defer dtm.mu.Unlock()
 
-	dtm.tokens[dtm.current].expired = true
+	for i, t := range dtm.tokens {
-	dtm.current = (dtm.current + 1) % len(dtm.tokens)
+		if t.value == token {
+			dtm.tokens[i].expired = true
+			return nil
+		}
+	}
+
+	return fmt.Errorf("token not found")
 }
 
 // ResetAllTokens resets all tokens to expired=false.

pkg/utils/bwlimit.go

@@ -0,0 +1,10 @@
+package utils
+
+import "github.com/debridmediamanager/zurg/pkg/http"
+
+func IsBWLimitExceeded(err error) bool {
+	if dlErr, ok := err.(*http.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
+		return true
+	}
+	return false
+}