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() }