From d61dba5096f7bf05b937685c6dbacce739d00063 Mon Sep 17 00:00:00 2001
From: YoursFunny <admin@yoursfunny.top>
Date: Sat, 8 Aug 2026 20:03:14 +0800
Subject: [PATCH] send: stream media downloads with hard size caps

download_to_temp now uses download_media_limited: non-photos abort the
moment the 10 MiB upload cap is crossed mid-stream (no more full-body
buffering before the size check), photos cap at the 512 MiB decode
budget, and the ugoira frame zip gets a 512 MiB cap. MediaTooLarge
routes to the existing smaller-URL fallback.
---
 crates/x-media/src/site/pixiv/api.rs |  2 +-
 crates/xmedia-bot/src/photo.rs       |  5 +++--
 crates/xmedia-bot/src/send.rs        | 18 ++++++++++++------
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/crates/x-media/src/site/pixiv/api.rs b/crates/x-media/src/site/pixiv/api.rs
index b5a89e1..171938d 100644
--- a/crates/x-media/src/site/pixiv/api.rs
+++ b/crates/x-media/src/site/pixiv/api.rs
@@ -222,7 +222,7 @@ impl PixivAPI {
         let Some(zip_url) = zip_url else {
             return Ok(None);
         };
-        let zip_bytes = crate::site::download_media(&zip_url)
+        let zip_bytes = crate::site::download_media_limited(&zip_url, 512 * 1024 * 1024)
             .await
             .map_err(|e| match e {
                 FetchError::Http(e) => PixivError::Http(e),
diff --git a/crates/xmedia-bot/src/photo.rs b/crates/xmedia-bot/src/photo.rs
index 3920829..25ec765 100644
--- a/crates/xmedia-bot/src/photo.rs
+++ b/crates/xmedia-bot/src/photo.rs
@@ -26,8 +26,9 @@ pub const PHOTO_TARGET_DIMENSION_SUM: u32 = 9900;
 /// to a smaller media URL instead.
 pub const MAX_UPLOAD_BYTES: u64 = 10 * 1024 * 1024;
 /// Decode budget (bytes): a larger intermediate buffer is not worth the peak
-/// memory; the photo degrades to the smaller URL instead.
-const MAX_DECODE_BYTES: u64 = 512 * 1024 * 1024;
+/// memory; the photo degrades to the smaller URL instead. Also the cap for
+/// downloading photos in the send fallback (they must be downloaded whole).
+pub(crate) const MAX_DECODE_BYTES: u64 = 512 * 1024 * 1024;
 /// JPEG output quality (1-100).
 const JPEG_QUALITY: u8 = 90;
 
diff --git a/crates/xmedia-bot/src/send.rs b/crates/xmedia-bot/src/send.rs
index 7adf351..ab2aca6 100644
--- a/crates/xmedia-bot/src/send.rs
+++ b/crates/xmedia-bot/src/send.rs
@@ -472,24 +472,30 @@ async fn download_to_temp(item: &MediaItemPayload) -> Result<NamedTempFile, Fall
         | MediaItemPayload::Video { media, .. }
         | MediaItemPayload::Animation { media, .. } => media,
     };
-    let bytes = match x_media::site::download_media(media_url).await {
+    // Photos are downloaded even over the upload cap so `prepare_photo` can
+    // downscale / transcode them (cap = decode budget); videos/animations
+    // abort as soon as the upload cap is crossed mid-stream.
+    let limit = if matches!(item, MediaItemPayload::Photo { .. }) {
+        photo::MAX_DECODE_BYTES
+    } else {
+        MAX_UPLOAD_BYTES + 1
+    };
+    let bytes = match x_media::site::download_media_limited(media_url, limit).await {
         Ok(bytes) => bytes,
         Err(FetchError::Http(_)) => {
             return Err(FallbackError::Retryable {
                 delay_seconds: retry_delay_seconds(0),
             });
         }
+        Err(FetchError::TooLarge) => {
+            return Err(FallbackError::MediaTooLarge);
+        }
         Err(e) => {
             return Err(FallbackError::Permanent {
                 message: format!("download failed: {e}"),
             });
         }
     };
-    // Photos are downloaded even over the cap so `prepare_photo` can
-    // downscale / transcode them; only videos/animations short-circuit.
-    if !matches!(item, MediaItemPayload::Photo { .. }) && bytes.len() as u64 > MAX_UPLOAD_BYTES {
-        return Err(FallbackError::MediaTooLarge);
-    }
     let ext = sniff_ext(&bytes);
     let mut file = tempfile::Builder::new()
         .suffix(&format!(".{ext}"))