hpz937/zurg
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add token management

Browse Source
This commit is contained in:
Ben Adrian Sarmiento
2024-06-28 13:19:09 +02:00
parent 962845fb81
commit c3aea427d0
9 changed files with 100 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
internal/app.go
View File

@@ -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"),

4
internal/handlers/home.go
View File

@@ -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

8
internal/handlers/router.go
View File

@@ -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)

12
internal/torrent/manager.go
View File

@@ -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()
var tokenMap cmap.ConcurrentMap[string, *realdebrid.Download]
if token != "" {
tokenMap, _ = t.rd.UnrestrictMap.Get(token)
}
token := t.Config.GetToken()
tokenMap, _ := t.rd.UnrestrictMap.Get(token)
downloads := t.rd.GetDownloads()
mountedCount := 0

9
internal/torrent/repair.go
View File

@@ -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)
if dlErr, ok := err.(*http.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
unrestrict, err := t.rd.UnrestrictAndVerify(link)
if utils.IsBWLimitExceeded(err) {
bwLimitReached = true
return true
}

16
internal/universal/downloader.go
View File

@@ -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)
if dlErr, ok := err.(*zurghttp.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
downloadResp, err := dl.rd.DownloadFile(dlReq)
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

125
pkg/realdebrid/api.go
View File

@@ -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)
}

20
pkg/realdebrid/token_manager.go
View File

@@ -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.

10
pkg/utils/bwlimit.go Normal file
View File

@@ -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
}