Add token management

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) 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) {
+func (rd *RealDebrid) UnrestrictAndVerify(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" {
-			rd.tokenManager.SetCurrentTokenExpired()
+
+		// check if the link is already unrestricted
+		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) {
-	// 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
-		}
-	}
-
+func (rd *RealDebrid) UnrestrictLink(link string) (*Download, error) {
 	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) {
-	return rd.tokenManager.GetCurrentToken()
+func (rd *RealDebrid) DownloadFile(req *http.Request) (*http.Response, error) {
+	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()
-	defer dtm.mu.Unlock()
+	dtm.mu.RLock()
+	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.
-func (dtm *DownloadTokenManager) SetCurrentTokenExpired() {
+// SetTokenAsExpired sets the specified token as expired.
+func (dtm *DownloadTokenManager) SetTokenAsExpired(token string) error {
 	dtm.mu.Lock()
 	defer dtm.mu.Unlock()
 
-	dtm.tokens[dtm.current].expired = true
-	dtm.current = (dtm.current + 1) % len(dtm.tokens)
+	for i, t := range 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.