Multi-token support

pkg/realdebrid/api.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"time"
 
 	"github.com/debridmediamanager/zurg/internal/config"
 	zurghttp "github.com/debridmediamanager/zurg/pkg/http"
@@ -17,25 +18,44 @@ import (
 type RealDebrid struct {
 	torrentsCache    []Torrent
 	UnrestrictMap    cmap.ConcurrentMap[string, cmap.ConcurrentMap[string, *Download]]
+	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
 }
 
 func NewRealDebrid(apiClient, unrestrictClient, downloadClient *zurghttp.HTTPClient, workerPool *ants.Pool, cfg config.ConfigInterface, log *logutil.Logger) *RealDebrid {
+	mainToken := cfg.GetToken()
+	downloadTokens := cfg.GetDownloadTokens()
+	if !strings.Contains(strings.Join(downloadTokens, ","), mainToken) {
+		downloadTokens = append([]string{mainToken}, downloadTokens...)
+	}
+
 	rd := &RealDebrid{
 		torrentsCache:    []Torrent{},
+		UnrestrictMap:    cmap.New[cmap.ConcurrentMap[string, *Download]](),
+		verifiedLinks:    cmap.New[int64](),
 		apiClient:        apiClient,
 		unrestrictClient: unrestrictClient,
 		downloadClient:   downloadClient,
+		tokenManager:     NewDownloadTokenManager(downloadTokens),
 		workerPool:       workerPool,
 		cfg:              cfg,
 		log:              log,
 	}
+
+	apiClient.SetToken(mainToken)
+	unrestrictClient.SetToken(mainToken)
+	for _, token := range downloadTokens {
+		rd.UnrestrictMap.Set(token, cmap.New[*Download]())
+	}
+
 	rd.loadCachedTorrents()
+
 	return rd
 }
 
@@ -76,10 +96,42 @@ func (rd *RealDebrid) UnrestrictCheck(link string) (*Download, error) {
 	return &response, nil
 }
 
-func (rd *RealDebrid) UnrestrictLink(link string, verifyDownloadURL bool) (*Download, error) {
+func (rd *RealDebrid) UnrestrictLink(link string) (*Download, error) {
+	for {
+		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" {
+			rd.tokenManager.SetCurrentTokenExpired()
+			continue
+		}
+		return download, err
+	}
+}
+
+func (rd *RealDebrid) UnrestrictLinkWithToken(token, 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
+		// set link to max 39 chars (26 + 13)
 		link = link[0:39]
 	}
 	data.Set("link", link)
@@ -87,12 +139,13 @@ func (rd *RealDebrid) UnrestrictLink(link string, verifyDownloadURL bool) (*Down
 
 	req, err := http.NewRequest(http.MethodPost, "https://api.real-debrid.com/rest/1.0/unrestrict/link", requestBody)
 	if err != nil {
-		rd.log.Errorf("Error when creating a unrestrict link request: %v", err)
+		// rd.log.Errorf("Error when creating a unrestrict link request: %v", err)
 		return nil, err
 	}
 
 	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 {
@@ -114,14 +167,16 @@ func (rd *RealDebrid) UnrestrictLink(link string, verifyDownloadURL bool) (*Down
 		return nil, fmt.Errorf("undecodable response: %v", err)
 	}
 
-	// will only check for first byte if serving from rclone
-	if verifyDownloadURL {
-		err := rd.downloadClient.VerifyURL(response.Download)
-		if err != nil {
-			return nil, 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
 }
@@ -372,3 +427,7 @@ func (rd *RealDebrid) AvailabilityCheck(hashes []string) (AvailabilityResponse,
 
 	return response, nil
 }
+
+func (rd *RealDebrid) GetToken() (string, error) {
+	return rd.tokenManager.GetCurrentToken()
+}

pkg/realdebrid/token_manager.go

@@ -0,0 +1,64 @@
+package realdebrid
+
+import (
+	"fmt"
+	"sync"
+)
+
+type Token struct {
+	value   string
+	expired bool
+}
+
+type DownloadTokenManager struct {
+	tokens  []Token
+	current int
+	mu      sync.Mutex
+}
+
+// NewDownloadTokenManager initializes a new DownloadTokenManager with the given tokens.
+func NewDownloadTokenManager(tokenStrings []string) *DownloadTokenManager {
+	tokens := make([]Token, len(tokenStrings))
+	for i, t := range tokenStrings {
+		tokens[i] = Token{value: t, expired: false}
+	}
+	return &DownloadTokenManager{tokens: tokens, current: 0}
+}
+
+// GetCurrentToken returns the current non-expired token.
+func (dtm *DownloadTokenManager) GetCurrentToken() (string, error) {
+	dtm.mu.Lock()
+	defer dtm.mu.Unlock()
+
+	for {
+		if !dtm.tokens[dtm.current].expired {
+			return dtm.tokens[dtm.current].value, nil
+		}
+
+		dtm.current = (dtm.current + 1) % len(dtm.tokens)
+
+		if dtm.current == 0 {
+			return "", fmt.Errorf("all tokens are bandwidth-limited")
+		}
+	}
+}
+
+// SetCurrentTokenExpired sets the current token as expired.
+func (dtm *DownloadTokenManager) SetCurrentTokenExpired() {
+	dtm.mu.Lock()
+	defer dtm.mu.Unlock()
+
+	dtm.tokens[dtm.current].expired = true
+	dtm.current = (dtm.current + 1) % len(dtm.tokens)
+}
+
+// ResetAllTokens resets all tokens to expired=false.
+func (dtm *DownloadTokenManager) ResetAllTokens() {
+	dtm.mu.Lock()
+	defer dtm.mu.Unlock()
+
+	for i := range dtm.tokens {
+		dtm.tokens[i].expired = false
+	}
+	dtm.current = 0
+}