diff --git a/internal/handlers/home.go b/internal/handlers/home.go index d145831..467565c 100644 --- a/internal/handlers/home.go +++ b/internal/handlers/home.go @@ -43,7 +43,7 @@ type RootResponse struct { DownloadTokens []string `json:"download_tokens"` IDsToDelete []string `json:"ids_to_delete"` Hosts []string `json:"hosts"` - TrafficServedPerAPI uint64 `json:"traffic_served_per_api"` + TrafficServedPerAPI map[string]uint64 `json:"traffic_served_per_api"` } func (zr *Handlers) generateResponse(resp http.ResponseWriter, req *http.Request) (*RootResponse, error) { @@ -77,15 +77,16 @@ func (zr *Handlers) generateResponse(resp http.ResponseWriter, req *http.Request sortedIDs := zr.torMgr.OnceDoneBin.ToSlice() sort.Strings(sortedIDs) - trafficDetails, err := zr.rd.GetTrafficDetails() - if err != nil { - http.Error(resp, err.Error(), http.StatusInternalServerError) - return nil, err - } - var trafficFromAPI int64 - trafficFromAPI = 0 - if _, ok := trafficDetails["real-debrid.com"]; ok { - trafficFromAPI = trafficDetails["real-debrid.com"] + trafficFromAPI := make(map[string]uint64) + for _, token := range zr.rd.TokenManager.GetAllTokens() { + trafficDetails, err := zr.rd.GetTrafficDetails(token) + if err != nil { + trafficDetails = make(map[string]uint64) + } + trafficFromAPI[token] = 0 + if _, ok := trafficDetails["real-debrid.com"]; ok { + trafficFromAPI[token] = trafficDetails["real-debrid.com"] + } } userInfo.Premium = userInfo.Premium / 86400 @@ -105,7 +106,7 @@ func (zr *Handlers) generateResponse(resp http.ResponseWriter, req *http.Request Infuse: fmt.Sprintf("//%s/infuse/", req.Host), Logs: fmt.Sprintf("//%s/logs/", req.Host), UserInfo: userInfo, - TrafficServedPerAPI: uint64(trafficFromAPI), + TrafficServedPerAPI: trafficFromAPI, LibrarySize: allTorrents.Count(), TorrentsToRepair: repairQueueStr, MemAlloc: bToMb(mem.Alloc), @@ -220,10 +221,6 @@ func (zr *Handlers) handleHome(resp http.ResponseWriter, req *http.Request) { Process ID %d - - Traffic Served (main token) - %d MB (%d MB since startup) - Traffic Served (zurg) %d MB @@ -234,11 +231,22 @@ func (zr *Handlers) handleHome(resp http.ResponseWriter, req *http.Request) { response.Sys, response.NumGC, response.PID, - bToMb(response.TrafficServedPerAPI), // traffic served *api* - bToMb(response.TrafficServedPerAPI-zr.downloader.TrafficOnStartup.Load()), // traffic served *api* since startup - bToMb(zr.downloader.TrafficServed.Load()), // traffic served *zurg* + bToMb(zr.downloader.TrafficServed.Load()), // traffic served *zurg* ) + for token, traffic := range response.TrafficServedPerAPI { + trafficOnStartup, _ := zr.downloader.TrafficOnStartup.Get(token) + out += fmt.Sprintf(` + + Traffic Served (%s) + %d MB (%d MB since startup) + `, + utils.MaskToken(token)[40:], + bToMb(traffic), // traffic served *api* + bToMb(traffic-trafficOnStartup.Load()), // traffic served *api* since startup + ) + } + out += fmt.Sprintf(` Sponsor Zurg diff --git a/internal/universal/downloader.go b/internal/universal/downloader.go index eb11973..9cddf9e 100644 --- a/internal/universal/downloader.go +++ b/internal/universal/downloader.go @@ -15,6 +15,7 @@ import ( "github.com/debridmediamanager/zurg/pkg/logutil" "github.com/debridmediamanager/zurg/pkg/realdebrid" "github.com/debridmediamanager/zurg/pkg/utils" + cmap "github.com/orcaman/concurrent-map/v2" "github.com/panjf2000/ants/v2" ) @@ -22,22 +23,27 @@ type Downloader struct { rd *realdebrid.RealDebrid workerPool *ants.Pool TrafficServed atomic.Uint64 - TrafficOnStartup atomic.Uint64 + TrafficOnStartup cmap.ConcurrentMap[string, *atomic.Uint64] } func NewDownloader(rd *realdebrid.RealDebrid, workerPool *ants.Pool) *Downloader { dl := &Downloader{ - rd: rd, - workerPool: workerPool, + rd: rd, + workerPool: workerPool, + TrafficOnStartup: cmap.New[*atomic.Uint64](), } - trafficDetails, err := dl.rd.GetTrafficDetails() - if err != nil { - trafficDetails = make(map[string]int64) - } - dl.TrafficOnStartup.Store(uint64(0)) - if _, ok := trafficDetails["real-debrid.com"]; ok { - dl.TrafficOnStartup.Store(uint64(trafficDetails["real-debrid.com"])) + for _, token := range rd.TokenManager.GetAllTokens() { + trafficDetails, err := dl.rd.GetTrafficDetails(token) + if err != nil { + trafficDetails = make(map[string]uint64) + } + var trafficOnStartup atomic.Uint64 + trafficOnStartup.Store(0) + if _, ok := trafficDetails["real-debrid.com"]; ok { + trafficOnStartup.Store(uint64(trafficDetails["real-debrid.com"])) + } + dl.TrafficOnStartup.Set(token, &trafficOnStartup) } return dl @@ -62,7 +68,9 @@ func (dl *Downloader) StartResetBandwidthCountersJob() { for { dl.rd.TokenManager.ResetAllTokens() dl.TrafficServed.Store(0) - dl.TrafficOnStartup.Store(0) + dl.TrafficOnStartup.IterCb(func(token string, traffic *atomic.Uint64) { + traffic.Store(0) + }) <-ticker.C } }) diff --git a/pkg/http/client.go b/pkg/http/client.go index 894e6b0..00a4dd4 100644 --- a/pkg/http/client.go +++ b/pkg/http/client.go @@ -129,7 +129,7 @@ func NewHTTPClient( } func (r *HTTPClient) Do(req *http.Request) (*http.Response, error) { - if r.token != "" { + if r.token != "" && req.Header.Get("Authorization") == "" { req.Header.Set("Authorization", "Bearer "+r.token) } diff --git a/pkg/realdebrid/api.go b/pkg/realdebrid/api.go index 65c9b6e..0a9b408 100644 --- a/pkg/realdebrid/api.go +++ b/pkg/realdebrid/api.go @@ -338,12 +338,12 @@ func (rd *RealDebrid) GetUserInformation() (*User, error) { // TrafficDetails represents the structure of the traffic details response type TrafficDetails map[string]struct { - Host map[string]int64 `json:"host"` - Bytes int64 `json:"bytes"` + Host map[string]uint64 `json:"host"` + Bytes int64 `json:"bytes"` } // GetTrafficDetails gets the traffic details from the Real-Debrid API -func (rd *RealDebrid) GetTrafficDetails() (map[string]int64, error) { +func (rd *RealDebrid) GetTrafficDetails(token string) (map[string]uint64, error) { // Construct request URL reqURL := "https://api.real-debrid.com/rest/1.0/traffic/details" req, err := http.NewRequest(http.MethodGet, reqURL, nil) @@ -352,6 +352,8 @@ func (rd *RealDebrid) GetTrafficDetails() (map[string]int64, error) { return nil, err } + req.Header.Set("Authorization", "Bearer "+token) + // Send the request resp, err := rd.apiClient.Do(req) if err != nil { diff --git a/pkg/realdebrid/token_manager.go b/pkg/realdebrid/token_manager.go index d7e08e3..5a8b5f6 100644 --- a/pkg/realdebrid/token_manager.go +++ b/pkg/realdebrid/token_manager.go @@ -102,3 +102,11 @@ func (dtm *DownloadTokenManager) GetExpiredTokens() []string { } return tokens } + +func (dtm *DownloadTokenManager) GetAllTokens() []string { + var tokens []string + for _, t := range dtm.tokens { + tokens = append(tokens, t.value) + } + return tokens +}