downloads: cap yt-dlp media size

2 files changed, 27 insertions, 0 deletions

diff --git a/src/commands/utilities.rs b/src/commands/utilities.rs
index 7a2e65d..88829b2 100644
--- a/src/commands/utilities.rs
+++ b/src/commands/utilities.rs
@@ -18,6 +18,9 @@ type Context<'a> = poise::Context<'a, Data, Error>;
 // original and keeps replies snappy.
 const LARGE_FILE_THRESHOLD: usize = 10 * 1024 * 1024;
 
+/// Hard cap for a single `/download` result (before it is read into memory).
+const MAX_DOWNLOAD_BYTES: u64 = 512 * 1024 * 1024;
+
 #[derive(Clone, Copy, PartialEq, Eq)]
 pub(crate) enum FileType {
     Video,
@@ -524,6 +527,8 @@ async fn run_download(
         .arg("bestaudio/bestvideo+bestaudio")
         .arg("--merge-output-format")
         .arg("mp4")
+        .arg("--max-filesize")
+        .arg("512M")
         .arg("-o")
         .arg(&out_template);
     if is_audio {
@@ -574,6 +579,16 @@ async fn run_download(
         None => basename,
     };
 
+    match tokio::fs::metadata(&actual_path).await {
+        Ok(meta) if meta.len() > MAX_DOWNLOAD_BYTES => {
+            let _ = tokio::fs::remove_dir_all(&dl_dir).await;
+            ctx.say(lang::fmt_download_failed("file is larger than 512 MiB"))
+                .await?;
+            return Ok(());
+        }
+        _ => {}
+    }
+
     let bytes = match tokio::fs::read(&actual_path).await {
         Ok(b) => b,
         Err(e) => {
diff --git a/src/voice/mod.rs b/src/voice/mod.rs
index 5043508..d8308eb 100644
--- a/src/voice/mod.rs
+++ b/src/voice/mod.rs
@@ -26,6 +26,10 @@ use crate::recording::{Recorder, RecorderEventHandler};
 /// and the ffmpeg `s16le` encode. Must match the `DecodeConfig` set in `main.rs`.
 const RECORDING_SAMPLE_RATE: u32 = 48000;
 
+/// Upper bound for yt-dlp media buffered in memory. Audio is decoded to f32
+/// PCM afterwards (another 4x), so a multi-hundred-MB download would OOM.
+const MAX_MEDIA_BYTES: usize = 128 * 1024 * 1024;
+
 #[derive(Clone)]
 pub struct VoiceManager {
     pub songbird: Arc<Songbird>,
@@ -247,6 +251,8 @@ pub async fn ytdlp_download(url: &str) -> Result<Vec<u8>> {
                 "-o",
                 "-",
                 "--no-playlist",
+                "--max-filesize",
+                "128M",
                 "--",
                 url,
             ])
@@ -269,6 +275,12 @@ pub async fn ytdlp_download(url: &str) -> Result<Vec<u8>> {
         );
         return Err(Error::Custom(lang::fmt_ytdlp_stderr(&stderr)));
     }
+    if output.stdout.len() > MAX_MEDIA_BYTES {
+        return Err(Error::Custom(format!(
+            "media is too large (limit {} MiB)",
+            MAX_MEDIA_BYTES / (1024 * 1024)
+        )));
+    }
 
     Ok(output.stdout)
 }