Files
zurg/pkg/realdebrid/token_manager.go
Ben Adrian Sarmiento c3aea427d0 Add token management
2024-06-28 13:19:09 +02:00

71 lines
1.5 KiB
Go

package realdebrid
import (
"fmt"
"sync"
)
type Token struct {
value string
expired bool
}
type DownloadTokenManager struct {
tokens []Token
current int
mu sync.RWMutex
}
// 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.RLock()
defer dtm.mu.RUnlock()
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")
}
}
}
// SetTokenAsExpired sets the specified token as expired.
func (dtm *DownloadTokenManager) SetTokenAsExpired(token string) error {
dtm.mu.Lock()
defer dtm.mu.Unlock()
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.
func (dtm *DownloadTokenManager) ResetAllTokens() {
dtm.mu.Lock()
defer dtm.mu.Unlock()
for i := range dtm.tokens {
dtm.tokens[i].expired = false
}
dtm.current = 0
}