zurg/internal/universal/downloader.go

package universal

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"path/filepath"
	"strings"
	"sync/atomic"
	"time"

	"github.com/debridmediamanager/zurg/internal/config"
	intTor "github.com/debridmediamanager/zurg/internal/torrent"
	"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"
)

type Downloader struct {
	rd               *realdebrid.RealDebrid
	workerPool       *ants.Pool
	TrafficServed    atomic.Uint64
	TrafficOnStartup cmap.ConcurrentMap[string, *atomic.Uint64]
}

func NewDownloader(rd *realdebrid.RealDebrid, workerPool *ants.Pool) *Downloader {
	dl := &Downloader{
		rd:               rd,
		workerPool:       workerPool,
		TrafficOnStartup: cmap.New[*atomic.Uint64](),
	}

	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
}

// StartResetBandwidthCountersJob is a permanent job that resets the bandwidth counters at 12AM CET
func (dl *Downloader) StartResetBandwidthCountersJob() {
	// track bandwidth usage and reset at 12AM CET
	now := time.Now()
	tomorrow := now.AddDate(0, 0, 1)
	cetTZ, err := time.LoadLocation("CET")
	if err != nil {
		cetTZ = time.FixedZone("CET", 1*60*60)
	}
	nextMidnightInCET := time.Date(tomorrow.Year(), tomorrow.Month(), tomorrow.Day(), 0, 5, 0, 0, cetTZ)
	duration := nextMidnightInCET.Sub(now)
	timer := time.NewTimer(duration)

	dl.workerPool.Submit(func() {
		<-timer.C
		ticker := time.NewTicker(24 * time.Hour)
		for {
			dl.rd.TokenManager.ResetAllTokens()
			dl.TrafficServed.Store(0)
			dl.TrafficOnStartup.IterCb(func(token string, traffic *atomic.Uint64) {
				traffic.Store(0)
			})
			<-ticker.C
		}
	})
}

// DownloadFile handles a GET request for files in torrents
func (dl *Downloader) DownloadFile(
	directory,
	torrentName,
	fileName string,
	resp http.ResponseWriter,
	req *http.Request,
	torMgr *intTor.TorrentManager,
	cfg config.ConfigInterface,
	log *logutil.Logger,
) {
	torrents, ok := torMgr.DirectoryMap.Get(directory)
	if !ok {
		log.Errorf("Cannot find directory %s", directory)
		http.Error(resp, "File not found", http.StatusNotFound)
		return
	}

	torrent, ok := torrents.Get(torrentName)
	if !ok {
		log.Errorf("Cannot find torrent %s from path %s", torrentName, req.URL.Path)
		http.Error(resp, "File not found", http.StatusNotFound)
		return
	}

	file, ok := torrent.SelectedFiles.Get(fileName)
	if !ok || file.State.Is("deleted_file") {
		log.Errorf("Cannot find file %s from path %s", fileName, req.URL.Path)
		http.Error(resp, "File not found", http.StatusNotFound)
		return
	}

	if file.State.Is("broken_file") {
		http.Error(resp, "File is not available (being repaired)", http.StatusNotFound)
		return
	}

	unrestrict, err := torMgr.UnrestrictFile(file)
	if utils.AreAllTokensExpired(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
	}
	if err != nil {
		if file.State.Event(context.Background(), "break_file") == nil {
			torMgr.EnqueueForRepair(torrent)
		}
		log.Errorf("Error unrestricting file %s: %v", file.Path, err)
		http.Error(resp, "File is not available (can't unrestrict)", http.StatusBadRequest)
		return
	}

	if unrestrict.Filesize != file.Bytes {
		// this is possible if there's only 1 streamable file in the torrent
		// and then suddenly it's a rar file
		actualExt := strings.ToLower(filepath.Ext(unrestrict.Filename))
		expectedExt := strings.ToLower(filepath.Ext(fileName))
		if actualExt != expectedExt && unrestrict.Streamable != 1 {
			log.Warnf("File was changed and is not streamable: %s and %s (link=%s)", fileName, unrestrict.Filename, unrestrict.Link)
		} else {
			log.Warnf("Filename mismatch: %s and %s", fileName, unrestrict.Filename)
		}
	}
	if cfg.ShouldServeFromRclone() {
		// log.Debugf("Redirecting to %s", unrestrict.Download)
		redirect(resp, req, unrestrict.Download)
	} else {
		// log.Debugf("Streaming %s", unrestrict.Download)
		dl.streamFileToResponse(torrent, file, unrestrict, resp, req, torMgr, cfg, log)
	}
}

// DownloadLink handles a GET request for downloads
func (dl *Downloader) DownloadLink(
	unrestrict *realdebrid.Download,
	resp http.ResponseWriter,
	req *http.Request,
	torMgr *intTor.TorrentManager,
	cfg config.ConfigInterface,
	log *logutil.Logger,
) {
	if cfg.ShouldServeFromRclone() {
		// log.Debugf("Redirecting to %s", unrestrict.Download)
		redirect(resp, req, unrestrict.Download)
	} else {
		log.Debugf("Streaming %s", unrestrict.Download)
		dl.streamFileToResponse(nil, nil, unrestrict, resp, req, torMgr, cfg, log)
	}
}

func (dl *Downloader) streamFileToResponse(
	torrent *intTor.Torrent, // can be nil if downloading a link
	file *intTor.File, // can be nil if downloading a link
	unrestrict *realdebrid.Download,
	resp http.ResponseWriter,
	req *http.Request,
	torMgr *intTor.TorrentManager,
	cfg config.ConfigInterface,
	log *logutil.Logger,
) {
	// Create a new request for the file download
	dlReq, err := http.NewRequest(http.MethodGet, unrestrict.Download, nil)
	if err != nil {
		if file != nil {
			log.Errorf("Error creating new request for file %s: %v", file.Path, err)
		}
		http.Error(resp, "File is not available (can't create request)", http.StatusBadRequest)
		return
	}

	// Add the range header if it exists
	if req.Header.Get("Range") != "" {
		rangeVal := req.Header.Get("Range")
		// check if open-ended range request (e.g. "bytes=999-")
		if strings.HasSuffix(rangeVal, "-") {
			rangeVal += fmt.Sprintf("%d", unrestrict.Filesize-1)
		}
		dlReq.Header.Add("Range", rangeVal)
	}

	downloadResp, err := dl.rd.DownloadFile(dlReq)
	if utils.IsBytesLimitReached(err) {
		dl.rd.TokenManager.SetTokenAsExpired(unrestrict.Token, "bandwidth limit exceeded")
		// 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
	} else if utils.IsInvalidDownloadCode(err) {
		http.Error(resp, "File is not available (invalid download code)", http.StatusInternalServerError)
		return
	} else if err != nil {
		log.Errorf("Cannot download file %s: %v", unrestrict.Download, err)
		if file != nil && file.State.Event(context.Background(), "break_file") == nil {
			torMgr.EnqueueForRepair(torrent)
		}
		http.Error(resp, "File is not available (can't download)", http.StatusBadRequest)
		return
	}
	defer downloadResp.Body.Close()

	// Check if the download was not successful
	if downloadResp.StatusCode != http.StatusOK && downloadResp.StatusCode != http.StatusPartialContent {
		log.Errorf("Received a %s status code for file %s", downloadResp.Status, unrestrict.Filename)
		if file != nil && file.State.Event(context.Background(), "break_file") == nil {
			torMgr.EnqueueForRepair(torrent)
		}
		http.Error(resp, "File is not available (download error)", http.StatusBadRequest)
		return
	}

	if cr := downloadResp.Header.Get("Content-Range"); cr != "" {
		resp.Header().Set("Content-Range", cr)
	}

	var n int64

	if cfg.ShouldLogRequests() {
		var totalBytes int64
		ticker := time.NewTicker(60 * time.Second)
		done := make(chan bool)
		go func() {
			for {
				select {
				case <-ticker.C:
					mbps := float64(totalBytes*8) / 60_000_000
					if mbps > 0 {
						log.Debugf("%s | %.2f MB/s", dlReq.URL, mbps)
					}
					totalBytes = 0
				case <-done:
					ticker.Stop()
					return
				}
			}
		}()
		defer func() {
			done <- true
		}()
		n, _ = io.Copy(io.MultiWriter(resp, &byteCounter{&totalBytes}), downloadResp.Body)
	} else {
		n, _ = io.Copy(resp, downloadResp.Body)
	}

	if !strings.HasPrefix(unrestrict.Link, "https://real-debrid.com/d/") {
		return
	}

	go func() {
		dl.TrafficServed.Add(uint64(n))

		if cfg.ShouldLogRequests() && bToMb(uint64(n)) > 0 {
			log.Debugf("Served %d MB of file %s (range=%s)", bToMb(uint64(n)), unrestrict.Filename, req.Header.Get("Range"))
		}
	}()
}

func redirect(resp http.ResponseWriter, req *http.Request, url string) {
	http.Redirect(resp, req, url, http.StatusFound)
}

type byteCounter struct {
	totalBytes *int64
}

func (bc *byteCounter) Write(p []byte) (int, error) {
	n := len(p)
	*bc.totalBytes += int64(n)
	return n, nil
}

func bToMb(b uint64) uint64 {
	return b / 1024 / 1024
}