ytdl-web/app/controllers/download.go

package controllers

import (
	"fmt"
	"io"
	"net/http"
	"net/url"
	"sync"

	"github.com/go-chi/chi/v5"
	"github.com/samber/lo"
	"github.com/spf13/viper"
	"go.fifitido.net/ytdl-web/app"
	"go.fifitido.net/ytdl-web/app/models"
	"go.fifitido.net/ytdl-web/pkg/htmx"
	"go.fifitido.net/ytdl-web/pkg/httpx"
	"go.fifitido.net/ytdl-web/pkg/server"
	"go.fifitido.net/ytdl-web/pkg/view"
	"go.fifitido.net/ytdl-web/pkg/ytdl"
	"go.fifitido.net/ytdl-web/pkg/ytdl/metadata"
	"go.fifitido.net/ytdl-web/version"
	"golang.org/x/exp/slog"
)

type DownloadController struct {
	ytdl ytdl.Ytdl
}

var _ server.Controller = (*DownloadController)(nil)

func NewDownloadController(ytdl ytdl.Ytdl) *DownloadController {
	return &DownloadController{
		ytdl: ytdl,
	}
}

func (c *DownloadController) Router(r chi.Router) {
	r.Get("/", c.ListDownloadLinks)
	r.Head("/proxy", c.ProxyDownload)
	r.Get("/proxy", c.ProxyDownload)
}

func (c *DownloadController) getUrlParam(r *http.Request) (string, bool) {
	urlRaw, err := httpx.Query(r, "url")
	if err != nil {
		return "", false
	}

	urlBytes, err := url.QueryUnescape(urlRaw)
	if err != nil || len(urlBytes) < 1 {
		return "", false
	}

	return urlBytes, true
}

func (c *DownloadController) ListDownloadLinks(w http.ResponseWriter, r *http.Request) {
	hx := htmx.New(w, r)
	var layout []string
	if !hx.IsHtmxRequest() {
		layout = append(layout, "layouts/main")
	}

	isSecure := r.URL.Scheme == "https" || r.Header.Get("X-Forwarded-Proto") == "https"

	videoUrl, ok := c.getUrlParam(r)
	if !ok {
		app.Views.Render(w, "index", view.Data{
			"BasePath":      viper.GetString("base_path"),
			"Version":       version.Version,
			"Build":         version.Build,
			"BinaryVersion": c.ytdl.Version(),
			"IsSecure":      isSecure,
			"Error": view.Data{
				"Message": "Invalid URL",
			},
		}, layout...)
		return
	}

	meta, err := c.ytdl.GetMetadata(videoUrl)
	if err != nil {
		app.Views.Render(w, "index", view.Data{
			"BasePath":      viper.GetString("base_path"),
			"Version":       version.Version,
			"Build":         version.Build,
			"BinaryVersion": c.ytdl.Version(),
			"IsSecure":      isSecure,
			"Error": view.Data{
				"Message":  "Could not find a video at that url",
				"RetryUrl": videoUrl,
			},
		}, layout...)
		return
	}

	if hx.IsHtmxRequest() {
		hx.PushUrl("/download?url=" + url.QueryEscape(videoUrl))
	}

	app.Views.Render(w, "download", view.Data{
		"BasePath":      viper.GetString("base_path"),
		"Version":       version.Version,
		"Build":         version.Build,
		"BinaryVersion": c.ytdl.Version(),
		"IsSecure":      isSecure,
		"Url":           videoUrl,
		"Meta":          meta,
		"Videos":        models.GetVideosFromMetadata(meta),
	}, layout...)
}

func (c *DownloadController) ProxyDownload(w http.ResponseWriter, r *http.Request) {
	videoUrl, ok := c.getUrlParam(r)
	if !ok {
		http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
		return
	}

	formatId, err := httpx.Query(r, "format")
	if err != nil {
		http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
		return
	}

	meta, err := c.ytdl.GetMetadata(videoUrl)
	if err != nil {
		slog.Error("Failed to get metadata", slog.String("error", err.Error()))
		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
		return
	}
	videos := models.GetVideosFromMetadata(meta)

	index, err := httpx.QueryInt(r, "index")
	if err != nil || index < 0 || index >= len(videos) {
		http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
		return
	}

	video := videos[index]
	format, ok := lo.Find(video.Formats, func(format metadata.Format) bool {
		return format.FormatID == formatId
	})
	if !ok {
		http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
		return
	}

	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s-%s.%s\"", meta.ID, format.Resolution, format.Ext))
	if format.Filesize != nil {
		w.Header().Set("Content-Length", fmt.Sprint(*format.Filesize))
	}

	if len(videos) == 1 {
		index = -1
	}

	pread, pwrite := io.Pipe()
	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()

		if err := c.ytdl.Download(pwrite, videoUrl, format.FormatID, index); err != nil {
			slog.Error("Failed to download", slog.String("error", err.Error()))
			http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			pwrite.CloseWithError(err)
			return
		} else {
			pwrite.Close()
		}
	}()

	go func() {
		defer wg.Done()

		_, err = io.Copy(w, pread)
		if err != nil {
			slog.Error("Failed to copy", slog.String("error", err.Error()))
			http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			return
		}
	}()

	wg.Wait()
}