zurg/internal/universal/downloader.go

package universal

import (
	"context"
	"io"
	"net/http"
	"path/filepath"
	"strings"

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

type Downloader struct {
	client *zurghttp.HTTPClient
}

func NewDownloader(client *zurghttp.HTTPClient) *Downloader {
	return &Downloader{
		client: client,
	}
}

// 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.Warnf("Cannot find directory %s", directory)
		http.Error(resp, "File not found", http.StatusNotFound)
		return
	}

	torrent, ok := torrents.Get(torrentName)
	if !ok {
		log.Warnf("Cannot find torrent %sfrom 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("ok_file") {
		// log.Warnf("Cannot find file %s from path %s", fileName, req.URL.Path)
		http.Error(resp, "File not found", http.StatusNotFound)
		return
	}

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

	unrestrict := torMgr.UnrestrictFileUntilOk(file, cfg.ShouldServeFromRclone())
	if unrestrict == nil {
		log.Warnf("File %s cannot be unrestricted (link=%s)", fileName, file.Link)
		if err := file.State.Event(context.Background(), "break_file"); err != nil {
			log.Errorf("File %s is stale: %v", fileName, err)
			http.Error(resp, "File is stale, please try again", http.StatusLocked)
			return
		}
		torMgr.TriggerRepair(torrent)
		http.Error(resp, "File is not available", http.StatusNotFound)
		return
	} else {
		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() {
			redirect(resp, req, unrestrict.Download)
		} else {
			dl.streamFileToResponse(torrent, file, unrestrict, resp, req, torMgr, cfg, log)
		}
		return
	}
}

// DownloadLink handles a GET request for downloads
func (dl *Downloader) DownloadLink(
	fileName,
	link string,
	resp http.ResponseWriter,
	req *http.Request,
	torMgr *intTor.TorrentManager,
	cfg config.ConfigInterface,
	log *logutil.Logger,
) {
	// log.Debugf("Opening file %s (%s)", fileName, link)
	unrestrict := torMgr.UnrestrictLinkUntilOk(link, cfg.ShouldServeFromRclone())
	if unrestrict == nil {
		log.Warnf("File %s cannot be unrestricted (link=%s)", fileName, link)
		http.Error(resp, "File is not available", http.StatusInternalServerError)
		return
	} else {
		if cfg.ShouldServeFromRclone() {
			redirect(resp, req, unrestrict.Download)
		} else {
			dl.streamFileToResponse(nil, nil, unrestrict, resp, req, torMgr, cfg, log)
		}
		return
	}
}

func (dl *Downloader) streamFileToResponse(
	torrent *intTor.Torrent,
	file *intTor.File,
	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", http.StatusNotFound)
		return
	}

	// Add the range header if it exists
	if req.Header.Get("Range") != "" {
		dlReq.Header.Add("Range", req.Header.Get("Range"))
		// log.Debugf("Range request for file %s: %s", unrestrict.Download, req.Header.Get("Range"))
	}

	// Perform the request
	downloadResp, err := dl.client.Do(dlReq)
	if err != nil {
		log.Warnf("Cannot download file %s: %v", unrestrict.Download, err)
		if file != nil {
			if err := file.State.Event(context.Background(), "break_file"); err != nil {
				log.Errorf("File %s is stale: %v", file.Path, err)
				http.Error(resp, "File is stale, please try again", http.StatusLocked)
				return
			}
			torMgr.TriggerRepair(torrent)
		}
		http.Error(resp, "File is not available", http.StatusNotFound)
		return
	}
	defer downloadResp.Body.Close()

	// Check if the download was not successful
	if downloadResp.StatusCode != http.StatusOK && downloadResp.StatusCode != http.StatusPartialContent {
		log.Warnf("Received a %s status code for file %s", downloadResp.Status, unrestrict.Filename)
		if file != nil {
			if err := file.State.Event(context.Background(), "break_file"); err != nil {
				log.Errorf("File %s is stale: %v", file.Path, err)
				http.Error(resp, "File is stale, please try again", http.StatusLocked)
				return
			}
			torMgr.TriggerRepair(torrent)
		}
		http.Error(resp, "File is not available", http.StatusNotFound)
		return
	}

	// Copy the content-range header from the download response to the response
	if cr := downloadResp.Header.Get("Content-Range"); cr != "" {
		resp.Header().Set("Content-Range", cr)
	}

	buf := make([]byte, cfg.GetNetworkBufferSize())
	io.CopyBuffer(resp, downloadResp.Body, buf)
}

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