Add job for monitoring bw limit status of tokens

pkg/premium/monitor.go

@@ -13,6 +13,7 @@ const (
 	MINIMUM_SLEEP     = 60     // 60 seconds
 )
 
+// MonitorPremiumStatus is a permanent job that monitors the premium status of the user
 func MonitorPremiumStatus(workerPool *ants.Pool, rd *realdebrid.RealDebrid, zurglog *logutil.Logger) {
 	var userInfo *realdebrid.User
 	var err error
@@ -52,8 +53,7 @@ func MonitorPremiumStatus(workerPool *ants.Pool, rd *realdebrid.RealDebrid, zurg
 			userInfo, err = rd.GetUserInformation()
 			if err != nil {
 				zurglog.Errorf("Failed to get user information: %v", err)
-				time.Sleep(5 * time.Minute)
-				continue
+				time.Sleep(time.Duration(MINIMUM_SLEEP) * time.Second)
 			}
 		}
 	})

pkg/realdebrid/api.go

@@ -38,7 +38,7 @@ func NewRealDebrid(apiClient, unrestrictClient, downloadClient *zurghttp.HTTPCli
 
 	rd := &RealDebrid{
 		UnrestrictMap:    cmap.New[cmap.ConcurrentMap[string, *Download]](),
-		TokenManager:     NewDownloadTokenManager(downloadTokens),
+		TokenManager:     NewDownloadTokenManager(downloadTokens, log),
 		torrentsCache:    []Torrent{},
 		verifiedLinks:    cmap.New[int64](),
 		apiClient:        apiClient,
@@ -69,8 +69,8 @@ func (rd *RealDebrid) UnrestrictAndVerify(link string) (*Download, error) {
 		}
 
 		// check if the link is already unrestricted
-		tokenMap, ok := rd.UnrestrictMap.Get(token)
-		if ok && tokenMap.Has(link) {
+		tokenMap, _ := rd.UnrestrictMap.Get(token)
+		if 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() {
@@ -78,8 +78,8 @@ func (rd *RealDebrid) UnrestrictAndVerify(link string) (*Download, error) {
 			}
 
 			err := rd.downloadClient.VerifyLink(download.Download)
-			if utils.IsBWLimitExceeded(err) {
-				rd.TokenManager.SetTokenAsExpired(token)
+			if utils.IsBytesLimitReached(err) {
+				rd.TokenManager.SetTokenAsExpired(token, "bandwidth limit exceeded")
 				continue
 			}
 			if err != nil {
@@ -94,12 +94,13 @@ func (rd *RealDebrid) UnrestrictAndVerify(link string) (*Download, error) {
 		if err != nil {
 			return nil, err
 		}
+		download.Token = token
 
 		tokenMap.Set(link, download)
 
 		err = rd.downloadClient.VerifyLink(download.Download)
-		if utils.IsBWLimitExceeded(err) {
-			rd.TokenManager.SetTokenAsExpired(token)
+		if utils.IsBytesLimitReached(err) {
+			rd.TokenManager.SetTokenAsExpired(token, "bandwidth limit exceeded")
 			continue
 		}
 		if err != nil {
@@ -404,3 +405,35 @@ func (rd *RealDebrid) AvailabilityCheck(hashes []string) (AvailabilityResponse,
 func (rd *RealDebrid) DownloadFile(req *http.Request) (*http.Response, error) {
 	return rd.downloadClient.Do(req)
 }
+
+// MonitorExpiredTokens is a permanent job for monitoring expired tokens if they are still expired
+func (rd *RealDebrid) MonitorExpiredTokens() {
+	sleepPeriod := 5 * time.Minute
+	rd.workerPool.Submit(func() {
+		for {
+			expiredTokens := rd.TokenManager.GetExpiredTokens()
+			for _, token := range expiredTokens {
+				tokenMap, _ := rd.UnrestrictMap.Get(token)
+				stillExpired := true
+				skipAll := false
+				tokenMap.IterCb(func(key string, download *Download) {
+					if skipAll {
+						return
+					}
+					err := rd.downloadClient.VerifyLink(download.Download)
+					if err != nil {
+						skipAll = utils.IsBytesLimitReached(err)
+						return
+					}
+					stillExpired = false
+					skipAll = true
+					rd.verifiedLinks.Set(download.ID, time.Now().Unix()+60*60*24)
+				})
+				if !stillExpired {
+					rd.TokenManager.SetTokenAsUnexpired(token)
+				}
+			}
+			time.Sleep(sleepPeriod)
+		}
+	})
+}

pkg/realdebrid/token_manager.go

@@ -3,6 +3,9 @@ package realdebrid
 import (
 	"fmt"
 	"sync"
+
+	"github.com/debridmediamanager/zurg/pkg/logutil"
+	"github.com/debridmediamanager/zurg/pkg/utils"
 )
 
 type Token struct {
@@ -14,15 +17,16 @@ type DownloadTokenManager struct {
 	tokens  []Token
 	current int
 	mu      sync.RWMutex
+	log     *logutil.Logger
 }
 
 // NewDownloadTokenManager initializes a new DownloadTokenManager with the given tokens.
-func NewDownloadTokenManager(tokenStrings []string) *DownloadTokenManager {
+func NewDownloadTokenManager(tokenStrings []string, log *logutil.Logger) *DownloadTokenManager {
 	tokens := make([]Token, len(tokenStrings))
 	for i, t := range tokenStrings {
 		tokens[i] = Token{value: t, expired: false}
 	}
-	return &DownloadTokenManager{tokens: tokens, current: 0}
+	return &DownloadTokenManager{tokens: tokens, current: 0, log: log}
 }
 
 // GetCurrentToken returns the current non-expired token.
@@ -38,19 +42,36 @@ func (dtm *DownloadTokenManager) GetCurrentToken() (string, error) {
 		dtm.current = (dtm.current + 1) % len(dtm.tokens)
 
 		if dtm.current == 0 {
-			return "", fmt.Errorf("all tokens are bandwidth-limited")
+			return "", fmt.Errorf("all tokens are expired")
 		}
 	}
 }
 
 // SetTokenAsExpired sets the specified token as expired.
-func (dtm *DownloadTokenManager) SetTokenAsExpired(token string) error {
+func (dtm *DownloadTokenManager) SetTokenAsExpired(token, reason string) error {
 	dtm.mu.Lock()
 	defer dtm.mu.Unlock()
 
 	for i, t := range dtm.tokens {
 		if t.value == token {
 			dtm.tokens[i].expired = true
+			dtm.log.Debugf("Token %s set as expired (reason=%s)", utils.MaskToken(token), reason)
+			return nil
+		}
+	}
+
+	return fmt.Errorf("token not found")
+}
+
+// SetTokenAsUnexpired sets the specified token as unexpired.
+func (dtm *DownloadTokenManager) SetTokenAsUnexpired(token string) error {
+	dtm.mu.Lock()
+	defer dtm.mu.Unlock()
+
+	for i, t := range dtm.tokens {
+		if t.value == token {
+			dtm.tokens[i].expired = false
+			dtm.log.Debugf("Token %s set as unexpired", utils.MaskToken(token))
 			return nil
 		}
 	}
@@ -68,3 +89,16 @@ func (dtm *DownloadTokenManager) ResetAllTokens() {
 	}
 	dtm.current = 0
 }
+
+func (dtm *DownloadTokenManager) GetExpiredTokens() []string {
+	dtm.mu.RLock()
+	defer dtm.mu.RUnlock()
+
+	var tokens []string
+	for _, t := range dtm.tokens {
+		if t.expired {
+			tokens = append(tokens, t.value)
+		}
+	}
+	return tokens
+}

pkg/realdebrid/types.go

@@ -24,6 +24,7 @@ type Download struct {
 	Download   string `json:"download"` // Generated link
 	Streamable int    `json:"streamable"`
 	Generated  string `json:"-"` // jsonDate
+	Token      string `json:"-"`
 }
 
 func (d *Download) UnmarshalJSON(data []byte) error {

pkg/utils/bwlimit.go

@@ -2,9 +2,13 @@ package utils
 
 import "github.com/debridmediamanager/zurg/pkg/http"
 
-func IsBWLimitExceeded(err error) bool {
+func IsBytesLimitReached(err error) bool {
 	if dlErr, ok := err.(*http.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
 		return true
 	}
 	return false
 }
+
+func AreAllTokensExpired(err error) bool {
+	return err.Error() == "all tokens are expired"
+}

pkg/utils/mask_token.go

@@ -0,0 +1,7 @@
+package utils
+
+import "strings"
+
+func MaskToken(token string) string {
+	return strings.Repeat("*", len(token)-4) + token[len(token)-4:]
+}