diff --git a/cmd/bot/webhook.go b/cmd/bot/webhook.go index b232a76..ccc81e7 100644 --- a/cmd/bot/webhook.go +++ b/cmd/bot/webhook.go @@ -5,6 +5,7 @@ import ( "log/slog" "net/http" "os" + "time" tgbot "github.com/go-telegram/bot" ) @@ -23,10 +24,26 @@ func startWebhook(ctx context.Context, b *tgbot.Bot, webhookURL string) { tlsCert := os.Getenv("BOT_TLS_CERT") tlsKey := os.Getenv("BOT_TLS_KEY") + server := &http.Server{ + Addr: listenAddr, + Handler: b.WebhookHandler(), + } + + // Shutdown HTTP server gracefully when context is cancelled + go func() { + <-ctx.Done() + slog.Info("shutting down webhook server") + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := server.Shutdown(shutdownCtx); err != nil { + slog.Error("webhook server shutdown", "error", err) + } + }() + if tlsCert != "" && tlsKey != "" { slog.Info("starting HTTPS webhook", "addr", listenAddr) go func() { - if err := http.ListenAndServeTLS(listenAddr, tlsCert, tlsKey, b.WebhookHandler()); err != nil { + if err := server.ListenAndServeTLS(tlsCert, tlsKey); err != nil && err != http.ErrServerClosed { slog.Error("HTTPS server", "error", err) os.Exit(1) } @@ -34,7 +51,7 @@ func startWebhook(ctx context.Context, b *tgbot.Bot, webhookURL string) { } else { slog.Info("starting HTTP webhook (TLS by reverse proxy)", "addr", listenAddr) go func() { - if err := http.ListenAndServe(listenAddr, b.WebhookHandler()); err != nil { + if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed { slog.Error("HTTP server", "error", err) os.Exit(1) } diff --git a/internal/dl/ytdlp.go b/internal/dl/ytdlp.go index 99353f0..eb6390f 100644 --- a/internal/dl/ytdlp.go +++ b/internal/dl/ytdlp.go @@ -10,6 +10,7 @@ import ( "os" "os/exec" "path/filepath" + "strconv" "strings" "sync" "syscall" @@ -305,6 +306,14 @@ func (c *Client) StartDownload(job *domain.DownloadJob, downloadDir string) (<-c // --newline uses \n instead of \r so bufio.Scanner can read line-by-line args = append(args, "--progress", "--newline") + // Allow resuming partial downloads + args = append(args, "--continue") + + // Enforce max file size limit + if job.MaxFileSize > 0 { + args = append(args, "--max-filesize", strconv.FormatInt(job.MaxFileSize, 10)) + } + // Language selection if specified if job.Language != "" { args = append(args, "--sub-langs", job.Language) diff --git a/internal/domain/types.go b/internal/domain/types.go index 7668d39..9e93aae 100644 --- a/internal/domain/types.go +++ b/internal/domain/types.go @@ -124,8 +124,10 @@ type DownloadJob struct { Speed string ETA string FileSize int64 + MaxFileSize int64 FilePath string Title string + ProgressPrefix string QuotaApplied bool StderrLog string Done chan struct{} diff --git a/internal/handler/download.go b/internal/handler/download.go index 9c94392..ea87655 100644 --- a/internal/handler/download.go +++ b/internal/handler/download.go @@ -31,12 +31,20 @@ type stepState struct { FormatIDs []string PlaylistState *domain.PlaylistState SearchState *domain.SearchState + CreatedAt time.Time } +const stepStateTTL = 30 * time.Minute + func (h *Handler) getUserStepState(chatID int64) *stepState { h.stepStatesMu.Lock() defer h.stepStatesMu.Unlock() - return h.stepStates[chatID] + ss := h.stepStates[chatID] + if ss != nil && time.Since(ss.CreatedAt) > stepStateTTL { + delete(h.stepStates, chatID) + return nil + } + return ss } func (h *Handler) setUserStepState(chatID int64, ss *stepState) { @@ -133,6 +141,7 @@ func (h *Handler) handleURLInput(ctx context.Context, chatID int64, text string, URL: url, MediaInfo: info, Step: 0, + CreatedAt: time.Now(), } h.sendMediaPreview(ctx, chatID, info) @@ -348,6 +357,7 @@ func (h *Handler) startDownload(ctx context.Context, chatID int64, msgID int, us Status: domain.StatusDownloading, Title: ss.MediaInfo.Title, FileSize: fileSize, + MaxFileSize: h.maxFileSize, Done: make(chan struct{}), } @@ -375,9 +385,13 @@ func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) { defer func() { <-h.downloadSem }() kb := progressKeyboard(0, "", "") + progressHeader := "Downloading: " + job.Title + if job.ProgressPrefix != "" { + progressHeader = job.ProgressPrefix + "\n" + progressHeader + } msg, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{ ChatID: job.ChatID, - Text: fmt.Sprintf("Downloading: %s", job.Title), + Text: progressHeader, ReplyMarkup: &kb, }) if err != nil { @@ -390,7 +404,9 @@ func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) { if err != nil { errMsg := err.Error() if job.StderrLog != "" { - errMsg += ": " + job.StderrLog + if lines := strings.SplitN(job.StderrLog, "\n", 2); len(lines) > 0 && lines[0] != "" { + errMsg = strings.TrimSpace(lines[0]) + } } h.editText(ctx, job.ChatID, job.MessageID, fmt.Sprintf("Download failed: %s", errMsg), nil) job.Status = domain.StatusFailed @@ -460,7 +476,10 @@ func (h *Handler) trackDownloadProgress(ctx context.Context, job *domain.Downloa } kb := progressKeyboard(upd.Progress, upd.Speed, upd.ETA) - text := fmt.Sprintf("Downloading: %s", job.Title) + text := "Downloading: " + job.Title + if job.ProgressPrefix != "" { + text = job.ProgressPrefix + "\n" + text + } h.editText(ctx, job.ChatID, job.MessageID, text, &kb) case <-ctx.Done(): @@ -732,7 +751,13 @@ func (h *Handler) sendMediaPreview(ctx context.Context, chatID int64, info *doma // downloadFile downloads a URL to a temp file in the given directory and returns the path. func downloadFile(url, dir string) (string, error) { - client := &http.Client{Timeout: 30 * time.Second} + tr := &http.Transport{ + Proxy: http.ProxyFromEnvironment, + } + client := &http.Client{ + Timeout: 30 * time.Second, + Transport: tr, + } resp, err := client.Get(url) if err != nil { return "", fmt.Errorf("http get: %w", err) diff --git a/internal/handler/handler.go b/internal/handler/handler.go index 13c8057..599dc99 100644 --- a/internal/handler/handler.go +++ b/internal/handler/handler.go @@ -464,21 +464,15 @@ func formatBytes(bytes int64) string { if bytes < unit { return strconv.FormatInt(bytes, 10) + "B" } - div, exp := int64(unit), 0 - for n := bytes / unit; n >= unit; n /= unit { - div *= unit - exp++ - } - switch exp { - case 0: - return fmt.Sprintf("%.1fKB", float64(bytes)/float64(div)) - case 1: - return fmt.Sprintf("%.1fMB", float64(bytes)/float64(div)) - case 2: - return fmt.Sprintf("%.1fGB", float64(bytes)/float64(div)) - default: - return fmt.Sprintf("%.1fTB", float64(bytes)/float64(div)) + units := []string{"KB", "MB", "GB", "TB"} + size := float64(bytes) + for _, u := range units { + size /= unit + if size < unit { + return fmt.Sprintf("%.1f%s", size, u) + } } + return fmt.Sprintf("%.1fTB", size) } func (h *Handler) generateJobID() string { diff --git a/internal/handler/playlist.go b/internal/handler/playlist.go index a352b4d..f74fa52 100644 --- a/internal/handler/playlist.go +++ b/internal/handler/playlist.go @@ -3,7 +3,12 @@ package handler import ( "context" "fmt" + "log/slog" "strconv" + "strings" + "time" + + tgbot "github.com/go-telegram/bot" "uptodownBot/internal/domain" ) @@ -20,6 +25,7 @@ func (h *Handler) handlePlaylist(ctx context.Context, chatID int64, userID int64 URL: url, MediaInfo: info, PlaylistState: state, + CreatedAt: time.Now(), } h.setUserStepState(chatID, ss) @@ -110,6 +116,11 @@ func (h *Handler) handlePlaylistSelectAll(ctx context.Context, chatID int64, msg } func (h *Handler) handlePlaylistConfirm(ctx context.Context, chatID int64, msgID int, userID int64) { + if job := h.getActiveJob(userID); job != nil { + h.sendText(ctx, chatID, "You already have an active download. Cancel it first.") + return + } + state := h.getPlaylistState(chatID) if state == nil { return @@ -127,65 +138,118 @@ func (h *Handler) handlePlaylistConfirm(ctx context.Context, chatID int64, msgID } h.clearStepState(chatID) - h.editText(ctx, chatID, msgID, fmt.Sprintf("Downloading %d videos sequentially...", len(selected)), nil) - for i, entry := range selected { - h.sendText(ctx, chatID, fmt.Sprintf("Downloading (%d/%d): %s", i+1, len(selected), entry.Title)) + // Parent context for the entire playlist — cancelling this stops all entries + playlistCtx, playlistCancel := context.WithCancel(context.Background()) + h.setCancelFunc(userID, playlistCancel) - info, err := h.YtDlp.GetMediaInfo(entry.URL) - if err != nil { - h.sendText(ctx, chatID, fmt.Sprintf("Failed: %s", err.Error())) - continue - } + h.editText(ctx, chatID, msgID, fmt.Sprintf("Downloading %d videos...", len(selected)), nil) - prefs, prefErr := h.Repo.GetOrCreatePreferences(userID) - if prefErr != nil { - prefs = &domain.UserPreferences{ - DefaultMediaType: domain.MediaTypeVideoAudio, - DefaultAudioFormat: "best", - DefaultVideoFormat: "best", + // Run the playlist loop in a goroutine so the handler is free to process + // cancel callbacks and other updates concurrently. + go func() { + for i, entry := range selected { + if playlistCtx.Err() != nil { + break } - } - formatID := prefs.DefaultVideoFormat - if prefs.DefaultMediaType == domain.MediaTypeAudio { - formatID = prefs.DefaultAudioFormat - } - format := &domain.Format{ID: formatID} - for _, f := range info.Formats { - if f.ID == formatID { - format = &f + info, err := h.YtDlp.GetMediaInfo(entry.URL) + if err != nil { + h.sendText(playlistCtx, chatID, fmt.Sprintf("Playlist entry (%d/%d) failed: %s", i+1, len(selected), err.Error())) + continue + } + + prefs, prefErr := h.Repo.GetOrCreatePreferences(userID) + if prefErr != nil { + prefs = &domain.UserPreferences{ + DefaultMediaType: domain.MediaTypeVideoAudio, + DefaultAudioFormat: "best", + DefaultVideoFormat: "best", + } + } + + formatID := prefs.DefaultVideoFormat + if prefs.DefaultMediaType == domain.MediaTypeAudio { + formatID = prefs.DefaultAudioFormat + } + ff := &domain.Format{ID: formatID} + for i := range info.Formats { + if info.Formats[i].ID == formatID { + ff = &info.Formats[i] + break + } + } + + job := &domain.DownloadJob{ + ID: h.generateJobID(), + UserID: userID, + ChatID: chatID, + MessageID: 0, + URL: entry.URL, + MediaType: prefs.DefaultMediaType, + SelectedFormat: ff, + FormatString: "", + Language: prefs.DefaultLanguage, + Status: domain.StatusDownloading, + Title: entry.Title, + FileSize: ff.Filesize, + MaxFileSize: h.maxFileSize, + ProgressPrefix: fmt.Sprintf("Downloading (%d/%d)", i+1, len(selected)), + Done: make(chan struct{}), + } + + ok := h.downloadPlaylistEntry(playlistCtx, userID, chatID, job) + if !ok { break } } - dlCtx, dlCancel := context.WithCancel(context.Background()) + h.clearActiveJob(userID) + h.sendText(playlistCtx, chatID, "Playlist downloads finished.") + }() +} - job := &domain.DownloadJob{ - ID: h.generateJobID(), - UserID: userID, - ChatID: chatID, - URL: entry.URL, - MediaType: prefs.DefaultMediaType, - SelectedFormat: format, - Language: prefs.DefaultLanguage, - Status: domain.StatusDownloading, - Title: entry.Title, - FileSize: format.Filesize, - Done: make(chan struct{}), - } - - h.setActiveJob(userID, job) - h.setCancelFunc(userID, dlCancel) - - h.runDownload(dlCtx, job) - - // runDownload is synchronous — it blocks until the download completes. - // The Done channel is closed inside runDownload when it exits. - // No busy-polling needed. +// downloadPlaylistEntry handles a single entry in a playlist download. +// It returns false if the playlist should stop (cancelled/fatal). +func (h *Handler) downloadPlaylistEntry(ctx context.Context, userID, chatID int64, job *domain.DownloadJob) bool { + if job.Done != nil { + defer close(job.Done) } - h.sendText(ctx, chatID, "All playlist downloads completed.") + h.setActiveJob(userID, job) + + kb := progressKeyboard(0, "", "") + progressHeader := job.ProgressPrefix + "\nDownloading: " + job.Title + msg, sendErr := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{ + ChatID: chatID, + Text: progressHeader, + ReplyMarkup: &kb, + }) + if sendErr != nil { + slog.Error("send playlist progress msg failed", "error", sendErr) + return true // non-fatal, continue to next + } + job.MessageID = msg.ID + + updates, dlErr := h.YtDlp.StartDownload(job, h.downloadDir) + if dlErr != nil { + errMsg := dlErr.Error() + if job.StderrLog != "" { + if lines := strings.SplitN(job.StderrLog, "\n", 2); len(lines) > 0 && lines[0] != "" { + errMsg = strings.TrimSpace(lines[0]) + } + } + h.editText(ctx, chatID, job.MessageID, fmt.Sprintf("Download failed: %s", errMsg), nil) + return true + } + + h.trackDownloadProgress(ctx, job, updates) + + if job.Status == domain.StatusCompleted { + h.handleDownloadCompletion(ctx, job) + } + + return job.Status != domain.StatusCancelled } func countSelected(state *domain.PlaylistState) int { diff --git a/internal/handler/search.go b/internal/handler/search.go index 02c1e2f..51679ea 100644 --- a/internal/handler/search.go +++ b/internal/handler/search.go @@ -5,6 +5,7 @@ import ( "fmt" "strconv" "strings" + "time" "github.com/go-telegram/bot/models" @@ -39,7 +40,7 @@ func (h *Handler) handleSearchQuery(ctx context.Context, chatID int64, text stri PerPage: 5, Source: source, } - ss := &stepState{SearchState: state} + ss := &stepState{SearchState: state, CreatedAt: time.Now()} h.setUserStepState(chatID, ss) kb := searchResultsKeyboard(state)