From 70d4f2ec88cedec841d27a18b36a62630954bcb0 Mon Sep 17 00:00:00 2001 From: YoursFunny Date: Mon, 21 Sep 2026 18:35:58 +0800 Subject: [PATCH] refactor(x-media): move the download stack into site/download.rs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit site/mod.rs held the registry, the Site trait, Fetched, the caption helpers, FetchError and, at the end, the whole media-download stack: two HTTP clients with different timeouts, the guard that refuses a URL inside the host's own network (start URL and every redirect hop), and the two streaming entry points. That stack is its own reason to change and moves whole, with its tests; CLIENT and the download fns stay re-exported at site::… so no adapter or bot call site moved. SITES becomes pub(crate) for the header rule. --- AGENTS.md | 4 +- crates/x-media/src/site/download.rs | 426 ++++++++++++++++++++++++++++ crates/x-media/src/site/mod.rs | 414 +-------------------------- 3 files changed, 433 insertions(+), 411 deletions(-) create mode 100644 crates/x-media/src/site/download.rs diff --git a/AGENTS.md b/AGENTS.md index b01b031..f956e1d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -84,7 +84,7 @@ Docker: `docker build -t tgxmb .` then `docker run --rm -d --name tgxmb --env-fi | `crates/xmedia-bot/src/handlers/` | `statics.rs` = `CHAT_STORE`/`TASK_QUEUE`/`CONFIG` singletons (open `$DATA_DIR/task_queue.db`, default `data/` **relative to CWD**, dir auto-created); `mod.rs` also holds `apply_caption_edit`, the one place a caption edit is applied and its failure classified: a short retryable delay is retried once, anything else is reported to the user instead of being swallowed (`callback.rs`'s template button answers its toast with the failure and leaves the record alone); `commands.rs` = command dispatch (incl. `/test ` send-only, `/debug ` parse-only, the read-only `/settings` every chat member can read — unlike the admin-only `/bot_dict` raw dump — and template removal; `/start`/`/help` carry the guidance teloxide's `descriptions()` cannot render, and `/set_format` rejects unknown `{…}` placeholders, resetting with `-`); `urls.rs` = URL extraction + the per-URL pipeline (`url_media` takes a `PostSend` mode: chat settings vs `/test`'s suppressed actions); `inline.rs` = debounced inline queries (hotlink-protected and local media skipped); `callback.rs` = edit-before-forward buttons (dptree entry + testable `handle_callback` core, incl. `skip`) | | `crates/xmedia-bot/src/send/` | `mod.rs`: constants `MAX_MEDIA_GROUP = 10` and the senders; `error.rs`: `classify_request_error` (5xx/non-JSON bodies retry, see the Retries bullet) and the media-fetch markers that route a URL send into the reupload fallback — including `failed to get HTTP url content`, the description single-media URL sends answer with. `upload.rs`: download-and-reupload fallback triggered only by Telegram API errors (`is_media_fetch_failure` / `is_size_error`), with a download's class from `classify_download_error` (transport/429/5xx retry; 4xx is permanent — the media itself is gone or refused — and a temp-file *write* failure retries, being resource exhaustion far more often than a broken temp dir). Item preparation is bounded **process-wide** (`PREP_SLOTS` in `upload.rs`: URL workers and queue workers can each be inside a batch, so a per-batch bound is not a memory bound), and the check that routes an oversized item to `fallback_url` is the download's own declared-Content-Length abort (`FetchError::TooLarge` → `MediaTooLarge`) — there is no separate size probe, which used to cost a second request per item. `post_send.rs`: settlement (`settle_task`), cache write, post-send actions (dead-letter text via `failure_text`: post key + cause, since the raw error alone does not say which link died), queue handlers. `input_media.rs`: payload → `InputMedia` | | `crates/xmedia-bot/src/photo.rs` | Pure-Rust photo processing (no ffmpeg): `png` (image-png) decode/encode + `zune-jpeg` decode + `fast_image_resize` Lanczos3 downscale + `jpeg-encoder`. Photos over Telegram's limits (width + height > 10000 px → `PHOTO_INVALID_DIMENSIONS`; bytes > 10 MiB) are decoded, downscaled keeping the format, PNG bit depth > 24 (RGBA 32-bit / 16-bit per channel) reduced to 24-bit RGB with alpha flattened white (≤24-bit untouched, never upconverted), and transcoded to JPEG only if still over the cap; memory budget guarded, otherwise the item's smaller fallback URL. Two budgets, not one: `MAX_PHOTO_DOWNLOAD_BYTES` (32 MiB) caps the *download* in the send fallback — the whole body is buffered, once per prep slot — while `MAX_DECODE_BYTES` (512 MiB) stays the pre-allocation guard that decides whether a decoded photo can be processed at all; over either one the item degrades to its smaller URL | -| `crates/x-media/src/site/mod.rs` | Dispatcher, `Fetched`/`FetchError`, shared `CLIENT`, `download_media_limited`/`download_media_to_file` (add `Referer: https://www.pixiv.net/` for `pximg.net` hotlink protection), `needs_media_headers` (the same per-site rule, asked by the inline path to skip what Telegram cannot fetch) | +| `crates/x-media/src/site/{mod,download}.rs` | `mod.rs`: dispatcher, `Fetched`/`FetchError`, `needs_media_headers` (the per-site rule, asked by the inline path to skip what Telegram cannot fetch). `download.rs`: the media-download stack — the metadata vs. media HTTP clients, the host-network guard (applied to the start URL and every redirect hop) and `download_media_limited`/`download_media_to_file` (which add the site's headers, e.g. `Referer: https://www.pixiv.net/` for `pximg.net`) | | `crates/x-media/src/site/pixiv/api.rs` | OAuth token exchange (hardcoded app client id/secret), access-token cache, ugoira zip→MP4 via ffmpeg in `spawn_blocking` | | `Dockerfile` | Multi-stage: cached dep layer via stub sources + `touch *.rs` mtime bump (cargo's freshness is mtime-based and `cargo clean -p` removes 0 files — the touch is what forces the real sources to rebuild while deps stay cached), static ffmpeg from ffmpeg.martin-riedl.de (`FFMPEG_URL` arg, optional `FFMPEG_SHA256` checksum, `unzip -t` integrity check), `debian:bookworm-slim` runtime, entrypoint. Runtime ships **no libssl/libcrypto/CA bundle** — rustls webpki-roots handles all TLS, and the static ffmpeg only processes local files (downloads go through reqwest) | | `docker-entrypoint.sh` | Privilege drop: `useradd` with `LOCAL_USER_ID` (default 9001) + `setpriv` (no gosu on bookworm-slim) | @@ -107,7 +107,7 @@ Docker: `docker build -t tgxmb .` then `docker run --rm -d --name tgxmb --env-fi - **~180 tests, all inline `#[cfg(test)] mod tests`** — no `tests/` integration directories. Framework: built-in Rust test + `#[tokio::test]` (dev-deps only in `x-media`: tokio macros/rt-multi-thread, dotenv). - No mocking framework anywhere (no mockito/wiremock/mockall). Conventions: pure-function units (regex parsing, serde round-trips, chunking, retry math) tested synchronously; async tests use real dependencies — file-backed SQLite via `tempfile` (`queue.rs::new_queue()` helper), live network fetches. Tests that must go through a **real `Bot`** (its URL/multipart building, the per-chat limiter and the bot-wide budget) talk to a stand-in API instead (`media_sender::test_support::fake_api::FakeApi`, a `tokio` TCP listener that records every call and answers the smallest result each method needs — teloxide keys methods by payload type, so the recorded name is `SendMediaGroup`, not `sendMediaGroup`): a media group, the edit-before-forward prompt through the real callback path, and `handlers::handle_message` (the context-taking body of `message_handler`, split out for exactly this). -- Live-network tests exist in `site/twitter/interface.rs` (5), `site/bsky/interface.rs` (1), `site/misskey/interface.rs` (1), `site/bilibili/interface.rs` (5), `site/pixiv/api.rs` (1); `photo.rs` adds one `#[ignore = "heavy: …"]` test. `site/mod.rs` also has a **token-gated but not `#[ignore]`d** pixiv download test (`download_media_pixiv_original_with_referer`): it hits `i.pximg.net` whenever `PIXIV_REFRESH_TOKEN` is set, so a local `cargo test --workspace` is not fully offline and can flake on a pixiv CDN body timeout. `disabled_site_is_reported_not_ignored` (same file) is gated the other way round: it asserts `fetch` answers `FetchError::Disabled { site: "pixiv" }` for a pixiv link and early-returns when `PIXIV_REFRESH_TOKEN` **is** set (the site is then enabled). Test gating convention (enforced by `.github/workflows/ci.yml`): pure unit tests always run; live-network tests carry `#[ignore = "live network: ..."]` (run via `cargo test --workspace -- --ignored live`); token-gated pixiv tests early-return when `PIXIV_REFRESH_TOKEN` is absent **or empty** (an unset GitHub secret arrives as `""` — `is_err()` alone would run them tokenless and fail), and the bilibili live tests early-return when the API answers risk control (`-352`, which bilibili applies per IP by request volume). Run the full offline suite with `cargo test --workspace`. +- Live-network tests exist in `site/twitter/interface.rs` (5), `site/bsky/interface.rs` (1), `site/misskey/interface.rs` (1), `site/bilibili/interface.rs` (5), `site/pixiv/api.rs` (1); `photo.rs` adds one `#[ignore = "heavy: …"]` test. `site/download.rs` also has a **token-gated but not `#[ignore]`d** pixiv download test (`download_media_pixiv_original_with_referer`): it hits `i.pximg.net` whenever `PIXIV_REFRESH_TOKEN` is set, so a local `cargo test --workspace` is not fully offline and can flake on a pixiv CDN body timeout. `disabled_site_is_reported_not_ignored` (same file) is gated the other way round: it asserts `fetch` answers `FetchError::Disabled { site: "pixiv" }` for a pixiv link and early-returns when `PIXIV_REFRESH_TOKEN` **is** set (the site is then enabled). Test gating convention (enforced by `.github/workflows/ci.yml`): pure unit tests always run; live-network tests carry `#[ignore = "live network: ..."]` (run via `cargo test --workspace -- --ignored live`); token-gated pixiv tests early-return when `PIXIV_REFRESH_TOKEN` is absent **or empty** (an unset GitHub secret arrives as `""` — `is_err()` alone would run them tokenless and fail), and the bilibili live tests early-return when the API answers risk control (`-352`, which bilibili applies per IP by request volume). Run the full offline suite with `cargo test --workspace`. - Fixtures are inline `serde_json::json!` builder fns (`fixture()`, `thread_json()`, `illust_json()`), not files. The shared `CLIENT` sets `pool_max_idle_per_host(0)` under `#[cfg(test)]` to avoid cross-runtime `DispatchGone`. - **CI** — `.github/workflows/ci.yml` (actions pinned to commit SHAs, `--locked` on every cargo invocation, `concurrency` cancels superseded runs, `RUST_BACKTRACE=1`) runs (behind a `changes` gate job, so a push/PR whose entire diff is markdown skips it instead of burning four minutes on nothing) `cargo fmt --check` + `cargo clippy --workspace --all-targets --locked -- -D warnings` + `cargo test --workspace --locked` + a release-profile `cargo build --release --locked` + an `actions-rust-lang/audit` dependency-vulnerability gate (offline, no secrets, on every push/PR) and a `live` job (schedule/manual/tag only, `-p x-media` since every network/secret-gated test lives there, `continue-on-error`) for the `#[ignore]`d live + token tests. `.github/workflows/docker.yml` builds and pushes the image on master/tag and runs a **build-only check on pull requests touching the build inputs**; its `should-build` gate skips a branch push that is already tagged (`git tag --points-at` — the tag run builds it, so push both refs together) or that touched no build input at all, while a tag push always builds (`Dockerfile`, entrypoint, manifests, `.dockerignore`); a release tag must match both crate versions or the build stops, and `FFMPEG_URL`/`FFMPEG_SHA256` are taken from repository variables when set (a release can pin an exact ffmpeg build). `.github/dependabot.yml` keeps crates, the pinned actions and the Docker base images current. - Untested and hard to test without a mock seam: `config.rs`, `handlers/statics.rs`; `db.rs` is covered for the migration chain but not for pool behaviour under contention; `main.rs` is covered where it was split out (`periodic_sweep`, `sweep_temp_dir`) but not for startup/shutdown or its `dptree` branch tree (the handlers themselves are, through the stand-in API); in `x-media`: `media.rs`, `lib.rs`, all `model.rs`. The `commands.rs` *executor* needs a real `Bot` (only its pure report builder is tested). Everything else — `handlers/{mod,callback,inline,urls}.rs`, `send/*`, `ctx.rs`, `state.rs`, `queue.rs`, `link_cache.rs`, `rate_limit.rs` — is driven through `TestStores`/`ctx::test_support` and the scripted `MockSender`. diff --git a/crates/x-media/src/site/download.rs b/crates/x-media/src/site/download.rs new file mode 100644 index 0000000..8813b92 --- /dev/null +++ b/crates/x-media/src/site/download.rs @@ -0,0 +1,426 @@ +//! The media-download stack: the two HTTP clients (site metadata vs. media, +//! which need different timeouts), the guard that keeps a download out of the +//! host's own network, and the two streaming entry points — a capped body in +//! memory ([`download_media_limited`]) and a large one written as it arrives +//! ([`download_media_to_file`]). +//! +//! Site-specific headers come from each adapter's `Site::media_headers`; no +//! code here knows about a particular site. + +use super::{FetchError, SITES}; +use std::sync::LazyLock; +use std::time::Duration; + +/// How long a download may make no progress: the response head, and then each +/// individual chunk, must arrive within this window. Not a total timeout — see +/// [`DOWNLOAD_TOTAL_TIMEOUT`]. +const DOWNLOAD_IDLE_TIMEOUT: Duration = Duration::from_secs(30); + +/// Absolute ceiling for one media download, on top of the idle window. A server +/// that drips a byte every 29 s keeps [`next_chunk`] satisfied indefinitely, and +/// on the bot's side each such download holds one of the process-wide upload-prep +/// slots (`send::upload`'s `PREP_SLOTS`) for as long as it lasts. Generous on +/// purpose: the legitimate cases are big — an ugoira frame zip runs to hundreds +/// of MB and an HLS remux pulls a whole video — and a slow link is not an error. +/// Checked between chunks, so a transfer that completes just over the budget is +/// kept rather than thrown away. +const DOWNLOAD_TOTAL_TIMEOUT: Duration = Duration::from_secs(600); + +/// The error a download reports when it spends its whole budget without +/// finishing. Retryable: the transfer may simply have been unlucky, and a retry +/// of the post restarts the download. +fn download_too_slow() -> FetchError { + FetchError::Transient(format!( + "download exceeded {}s", + DOWNLOAD_TOTAL_TIMEOUT.as_secs() + )) +} + +/// Builds a client with the shared configuration (browser User-Agent, the +/// Bot API's proxy, per-runtime pools under test). `total_timeout` is what +/// differs between the two clients below. +fn build_client(total_timeout: Option) -> reqwest::Client { + let mut builder = reqwest::Client::builder() + .user_agent("Mozilla/5.0") + .connect_timeout(Duration::from_secs(10)); + // Redirects stay allowed (site CDNs use them), but every hop goes through + // the same guard as the initial URL, and the cap stays reqwest's default: + // a third-party response must not be able to walk the bot into the host's + // own network. + builder = builder.redirect(reqwest::redirect::Policy::custom(|attempt| { + if !media_url_allowed(attempt.url()) { + log::warn!("refusing a media redirect into the host's own network"); + return attempt.error(FetchError::Blocked); + } + if attempt.previous().len() >= 10 { + return attempt.stop(); + } + attempt.follow() + })); + if let Some(total) = total_timeout { + // reqwest has no total timeout by default; a stalled connection + // would otherwise pin a fetch/handler forever. + builder = builder.timeout(total); + } + // Route site fetches through the same proxy the Bot API uses, so a + // network that needs TELOXIDE_PROXY (e.g. behind the GFW) does not + // leave site fetches dead while the bot itself works. + if let Some(proxy) = std::env::var("TELOXIDE_PROXY") + .ok() + .filter(|s| !s.is_empty()) + && let Ok(p) = reqwest::Proxy::all(&proxy) + { + builder = builder.proxy(p); + } + // Each `#[tokio::test]` runs on its own runtime; the connection pool is + // bound to the runtime that created it, so cross-runtime reuse of idle + // connections fails with DispatchGone. In test builds every request uses + // a fresh connection. Production runs on one runtime and keeps pooling. + #[cfg(test)] + let builder = builder.pool_max_idle_per_host(0); + builder.build().expect("failed to build HTTP client") +} + +/// Shared HTTP client (browser User-Agent) for the site fetches — metadata +/// requests, where 30s is generous. +pub(crate) static CLIENT: LazyLock = + LazyLock::new(|| build_client(Some(Duration::from_secs(30)))); + +/// Client for media *downloads*, with no reqwest-level total timeout: a 10 MiB +/// fallback download, or an ugoira frame zip that may be hundreds of MB, +/// legitimately takes minutes on a slow link — a 30s total cap made those posts +/// impossible to deliver at all (the size cap said 512 MiB, the clock said 30s). +/// What a stalled connection cannot do is hang a worker: the head and every +/// chunk are bounded by [`DOWNLOAD_IDLE_TIMEOUT`] (see [`next_chunk`]), and a +/// transfer that keeps trickling but never finishes is bounded by +/// [`DOWNLOAD_TOTAL_TIMEOUT`]. +static MEDIA_CLIENT: LazyLock = LazyLock::new(|| build_client(None)); + +/// The error a download reports when it stops making progress. +fn download_stalled() -> FetchError { + FetchError::Transient(format!( + "download stalled for {}s", + DOWNLOAD_IDLE_TIMEOUT.as_secs() + )) +} + +/// Sends a media-download request: the response head must arrive within the +/// idle window, and a non-2xx status is classified by [`download_status_error`]. +async fn send_download(request: reqwest::RequestBuilder) -> Result { + let response = match tokio::time::timeout(DOWNLOAD_IDLE_TIMEOUT, request.send()).await { + Ok(Ok(response)) => response, + Ok(Err(e)) => return Err(e.into()), + Err(_) => return Err(download_stalled()), + }; + if response.status().is_success() { + Ok(response) + } else { + Err(download_status_error(response.status())) + } +} + +/// One body chunk, or `None` at the end. A body that stops delivering is a +/// transient download error rather than a hang. +async fn next_chunk(response: &mut reqwest::Response) -> Result, FetchError> { + match tokio::time::timeout(DOWNLOAD_IDLE_TIMEOUT, response.chunk()).await { + Ok(Ok(chunk)) => Ok(chunk), + Ok(Err(e)) => Err(e.into()), + Err(_) => Err(download_stalled()), + } +} + +/// Whether an address must never be fetched. Media URLs come from a site's own +/// API response and the bytes are uploaded to Telegram, so following one into +/// the host's own network would turn the bot into a proxy for it: a cloud +/// metadata endpoint read back into a chat. +fn blocked_ip(addr: std::net::IpAddr) -> bool { + use std::net::IpAddr; + match addr { + IpAddr::V4(v4) => { + let [a, b, ..] = v4.octets(); + v4.is_private() // 10/8, 172.16/12, 192.168/16 + || v4.is_loopback() // 127/8 + || v4.is_link_local() // 169.254/16 — the cloud metadata range + || v4.is_unspecified() + || v4.is_broadcast() + || v4.is_documentation() + || v4.is_multicast() + // Ranges the std helpers do not cover: carrier-grade NAT and + // benchmarking. + || (a == 100 && (64..=127).contains(&b)) + || (a == 198 && (18..=19).contains(&b)) + } + IpAddr::V6(v6) => { + let [first, ..] = v6.segments(); + v6.is_loopback() + || v6.is_unspecified() + || v6.is_multicast() + || (first & 0xfe00) == 0xfc00 // unique local fc00::/7 + || (first & 0xffc0) == 0xfe80 // link local fe80::/10 + || v6.to_ipv4_mapped().is_some_and(|v4| blocked_ip(IpAddr::V4(v4))) + } + } +} + +/// `localhost` (and anything under it) plus the mDNS `.local` suffix: names that +/// only ever mean this machine. +fn is_local_name(name: &str) -> bool { + let name = name.trim_end_matches('.').to_ascii_lowercase(); + name == "localhost" || name.ends_with(".localhost") || name.ends_with(".local") +} + +/// Whether a media URL may be requested at all: http(s), and a host that is no +/// address or name of the host's own network. Applied to the URL a download +/// starts from *and* to every redirect hop. +/// +/// The residual gap is DNS rebinding — a name the site controls that resolves to +/// a private address. Closing it needs a `reqwest::dns::Resolve` wrapper +/// filtering resolved addresses; it is deliberately not installed, because the +/// same resolver also resolves the operator's proxy host and `TELOXIDE_PROXY` +/// is routinely a LAN address, so the guard would take down a working +/// deployment to block a far less likely attack. +fn media_url_allowed(url: &url::Url) -> bool { + if !matches!(url.scheme(), "http" | "https") { + return false; + } + match url.host() { + Some(url::Host::Ipv4(v4)) => !blocked_ip(v4.into()), + Some(url::Host::Ipv6(v6)) => !blocked_ip(v6.into()), + Some(url::Host::Domain(name)) => !is_local_name(name), + None => false, + } +} + +/// Prepares a media download: refuses a URL pointing inside the host's own +/// network ([`FetchError::Blocked`], permanent — the same URL would be refused +/// again), then applies the site's media headers. One choke point so every +/// download path gets the guard. +fn media_request(url: &str) -> Result { + let parsed = url::Url::parse(url).map_err(|e| { + log::warn!("media url is not a url: {e}"); + FetchError::Blocked + })?; + if !media_url_allowed(&parsed) { + log::warn!("refusing to fetch media from the host's own network"); + return Err(FetchError::Blocked); + } + Ok(apply_media_headers(MEDIA_CLIENT.get(parsed), url)) +} + +/// Applies every site's media-header rule to a download request (pixiv's +/// `Referer` for pximg.net hotlink protection). Sites contribute via their +/// `media_headers(url)` — the central download code carries no per-site logic. +fn apply_media_headers(mut request: reqwest::RequestBuilder, url: &str) -> reqwest::RequestBuilder { + for site in SITES.iter() { + if let Some(headers) = site.media_headers(url) { + for (name, value) in headers { + request = request.header(name, value); + } + } + } + request +} + +/// Maps a media download's HTTP status onto the same classes the site +/// adapters use, so callers can tell "try again" from "this URL is dead": +/// 4xx is a property of the media (gone, refused by the host), while 429/5xx +/// is a property of the moment. A transport error never reaches this — it +/// fails in `send()` and stays [`FetchError::Http`]. +fn download_status_error(status: reqwest::StatusCode) -> FetchError { + match status.as_u16() { + 401 | 403 => FetchError::Blocked, + 404 | 410 => FetchError::NotFound, + _ => FetchError::Transient(format!("media status {status}")), + } +} + +/// Downloads a media file with a hard size cap: the body is streamed and the +/// download aborts with [`FetchError::TooLarge`] the moment the cap is +/// crossed (or when a declared Content-Length already exceeds it). Keeps the +/// bot from buffering arbitrarily large bodies into memory — the size check +/// the bot's upload fallback needs is the one here, not a probe of its own. +/// +/// This is the bot's download path for the upload fallback: when Telegram +/// cannot fetch a media URL itself (hotlink protection), the bot downloads +/// the file and uploads it via multipart. Site-appropriate headers come from +/// each site's `media_headers` (pixiv image hosts need `Referer`). +pub async fn download_media_limited(url: &str, max_bytes: u64) -> Result { + let response = send_download(media_request(url)?).await?; + if let Some(len) = response.content_length() + && len > max_bytes + { + return Err(FetchError::TooLarge); + } + let mut response = response; + let mut buf = Vec::new(); + let started = std::time::Instant::now(); + while let Some(chunk) = next_chunk(&mut response).await? { + if started.elapsed() > DOWNLOAD_TOTAL_TIMEOUT { + return Err(download_too_slow()); + } + buf.extend_from_slice(&chunk); + if buf.len() as u64 > max_bytes { + return Err(FetchError::TooLarge); + } + } + Ok(bytes::Bytes::from(buf)) +} + +/// Streams a download to `out`, aborting with [`FetchError::TooLarge`] the +/// moment the body crosses `max_bytes` (or when a declared Content-Length +/// already exceeds it). Unlike [`download_media_limited`] the body is never +/// buffered in memory — used for large files (e.g. the pixiv ugoira frame +/// zip, which can be hundreds of MB) that would otherwise spike RAM. +/// Returns the number of bytes written. +pub async fn download_media_to_file( + url: &str, + max_bytes: u64, + out: &mut std::fs::File, +) -> Result { + use std::io::Write; + let response = send_download(media_request(url)?).await?; + if let Some(len) = response.content_length() + && len > max_bytes + { + return Err(FetchError::TooLarge); + } + let mut response = response; + let mut total: u64 = 0; + let started = std::time::Instant::now(); + while let Some(chunk) = next_chunk(&mut response).await? { + if started.elapsed() > DOWNLOAD_TOTAL_TIMEOUT { + return Err(download_too_slow()); + } + total += chunk.len() as u64; + if total > max_bytes { + return Err(FetchError::TooLarge); + } + out.write_all(&chunk).map_err(FetchError::Io)?; + } + Ok(total) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::site::{Fetched, pixiv}; + + #[test] + fn blocked_addresses_are_the_hosts_own_network() { + for addr in [ + "127.0.0.1", + "10.0.0.1", + "172.16.0.1", + "192.168.1.1", + "169.254.169.254", // cloud metadata + "0.0.0.0", + "255.255.255.255", + "100.64.0.1", // carrier-grade NAT + "198.18.0.1", // benchmarking + "::1", + "::", + "fc00::1", + "fe80::1", + "::ffff:127.0.0.1", + ] { + assert!(blocked_ip(addr.parse().unwrap()), "{addr}"); + } + for addr in [ + "1.1.1.1", + "93.184.216.34", + "2606:4700::1111", + "::ffff:1.1.1.1", + ] { + assert!(!blocked_ip(addr.parse().unwrap()), "{addr}"); + } + } + + #[test] + fn media_urls_inside_the_host_are_refused() { + for url in [ + "http://127.0.0.1:9/x", + "http://169.254.169.254/latest/meta-data/", + "http://[::1]:9/x", + "https://localhost/", + "https://prompt.localhost/x", + "https://printer.local/x", + "file:///etc/passwd", + "gopher://example.com/1", + ] { + let parsed = url::Url::parse(url).unwrap(); + assert!(!media_url_allowed(&parsed), "{url}"); + } + // Real media hosts and any public address stay fetchable. + for url in [ + "https://i.pximg.net/img-original/img/1.jpg", + "https://cdn.bsky.app/img/feed_thumbnail/plain/x", + "http://example.com/a", + "https://93.184.216.34/a", + ] { + let parsed = url::Url::parse(url).unwrap(); + assert!(media_url_allowed(&parsed), "{url}"); + } + } + /// The redirect-hop guard, against a public redirector: the initial URL is + /// checked by [`media_request`], but a redirect is the part of the path a + /// third-party response actually controls. + + #[tokio::test] + #[ignore = "live network: requires outbound HTTPS to httpbin.org"] + async fn live_redirect_into_the_hosts_network_is_refused() { + let url = "https://httpbin.org/redirect-to?url=http://169.254.169.254/latest/meta-data/"; + match download_media_limited(url, u64::MAX).await.unwrap_err() { + // A policy refusal reaches the caller wrapped by reqwest. + FetchError::Http(e) => assert!(e.is_redirect(), "got {e}"), + FetchError::Blocked => {} + other => panic!("expected a refusal, got {other:?}"), + } + } + + #[tokio::test] + async fn a_download_into_the_hosts_network_is_refused() { + // Refused on the URL alone: nothing has to be listening (or leaking) at + // the metadata endpoint for this to hold, and the class is permanent so + // the send path does not retry it. + for url in [ + "http://169.254.169.254/latest/meta-data/", + "http://127.0.0.1:9/secret", + ] { + let err = download_media_limited(url, u64::MAX).await.unwrap_err(); + assert!(matches!(err, FetchError::Blocked), "{url}: got {err:?}"); + } + // A malformed URL is refused the same way instead of becoming a + // retryable transport error. + assert!(matches!( + download_media_limited("not a url", u64::MAX) + .await + .unwrap_err(), + FetchError::Blocked + )); + } + + #[tokio::test] + async fn download_media_pixiv_original_with_referer() { + // Proves the Referer header is attached for i.pximg.net: a header-less + // GET to a pixiv original URL is rejected with 403. + // Empty-string check too: an unset CI secret arrives as "" (GitHub + // Actions), which would otherwise run the test tokenless and fail. + if std::env::var("PIXIV_REFRESH_TOKEN") + .ok() + .filter(|s| !s.is_empty()) + .is_none() + { + eprintln!("skipping: no PIXIV_REFRESH_TOKEN"); + return; + } + let illustration = pixiv::fetch(126839080).await.unwrap(); + let fetched: Fetched = illustration.into(); + let url = match fetched.media.first() { + Some(crate::media::Media::Illustration { url, .. }) => url.clone(), + other => panic!("expected illustration media, got {other:?}"), + }; + assert!(url.contains("i.pximg.net")); + let bytes = download_media_limited(&url, u64::MAX).await.unwrap(); + assert!(!bytes.is_empty()); + } +} diff --git a/crates/x-media/src/site/mod.rs b/crates/x-media/src/site/mod.rs index 76907f5..f689913 100644 --- a/crates/x-media/src/site/mod.rs +++ b/crates/x-media/src/site/mod.rs @@ -15,12 +15,16 @@ use thiserror::Error; pub mod bilibili; pub mod bsky; +mod download; pub mod misskey; pub mod pixiv; pub mod twitter; pub use pixiv::PixivError; +pub(crate) use download::CLIENT; +pub use download::{download_media_limited, download_media_to_file}; + /// The result of fetching a post: canonical URL, HTML caption, the post's /// title and body, media list and spoiler flag. Produced by [`fetch`]. #[derive(Debug)] @@ -307,124 +311,6 @@ pub fn status_error(site: &'static str, status: reqwest::StatusCode) -> FetchErr } } -/// How long a download may make no progress: the response head, and then each -/// individual chunk, must arrive within this window. Not a total timeout — see -/// [`DOWNLOAD_TOTAL_TIMEOUT`]. -const DOWNLOAD_IDLE_TIMEOUT: Duration = Duration::from_secs(30); - -/// Absolute ceiling for one media download, on top of the idle window. A server -/// that drips a byte every 29 s keeps [`next_chunk`] satisfied indefinitely, and -/// on the bot's side each such download holds one of the process-wide upload-prep -/// slots (`send::upload`'s `PREP_SLOTS`) for as long as it lasts. Generous on -/// purpose: the legitimate cases are big — an ugoira frame zip runs to hundreds -/// of MB and an HLS remux pulls a whole video — and a slow link is not an error. -/// Checked between chunks, so a transfer that completes just over the budget is -/// kept rather than thrown away. -const DOWNLOAD_TOTAL_TIMEOUT: Duration = Duration::from_secs(600); - -/// The error a download reports when it spends its whole budget without -/// finishing. Retryable: the transfer may simply have been unlucky, and a retry -/// of the post restarts the download. -fn download_too_slow() -> FetchError { - FetchError::Transient(format!( - "download exceeded {}s", - DOWNLOAD_TOTAL_TIMEOUT.as_secs() - )) -} - -/// Builds a client with the shared configuration (browser User-Agent, the -/// Bot API's proxy, per-runtime pools under test). `total_timeout` is what -/// differs between the two clients below. -fn build_client(total_timeout: Option) -> reqwest::Client { - let mut builder = reqwest::Client::builder() - .user_agent("Mozilla/5.0") - .connect_timeout(Duration::from_secs(10)); - // Redirects stay allowed (site CDNs use them), but every hop goes through - // the same guard as the initial URL, and the cap stays reqwest's default: - // a third-party response must not be able to walk the bot into the host's - // own network. - builder = builder.redirect(reqwest::redirect::Policy::custom(|attempt| { - if !media_url_allowed(attempt.url()) { - log::warn!("refusing a media redirect into the host's own network"); - return attempt.error(FetchError::Blocked); - } - if attempt.previous().len() >= 10 { - return attempt.stop(); - } - attempt.follow() - })); - if let Some(total) = total_timeout { - // reqwest has no total timeout by default; a stalled connection - // would otherwise pin a fetch/handler forever. - builder = builder.timeout(total); - } - // Route site fetches through the same proxy the Bot API uses, so a - // network that needs TELOXIDE_PROXY (e.g. behind the GFW) does not - // leave site fetches dead while the bot itself works. - if let Some(proxy) = std::env::var("TELOXIDE_PROXY") - .ok() - .filter(|s| !s.is_empty()) - && let Ok(p) = reqwest::Proxy::all(&proxy) - { - builder = builder.proxy(p); - } - // Each `#[tokio::test]` runs on its own runtime; the connection pool is - // bound to the runtime that created it, so cross-runtime reuse of idle - // connections fails with DispatchGone. In test builds every request uses - // a fresh connection. Production runs on one runtime and keeps pooling. - #[cfg(test)] - let builder = builder.pool_max_idle_per_host(0); - builder.build().expect("failed to build HTTP client") -} - -/// Shared HTTP client (browser User-Agent) for the site fetches — metadata -/// requests, where 30s is generous. -pub(crate) static CLIENT: LazyLock = - LazyLock::new(|| build_client(Some(Duration::from_secs(30)))); - -/// Client for media *downloads*, with no reqwest-level total timeout: a 10 MiB -/// fallback download, or an ugoira frame zip that may be hundreds of MB, -/// legitimately takes minutes on a slow link — a 30s total cap made those posts -/// impossible to deliver at all (the size cap said 512 MiB, the clock said 30s). -/// What a stalled connection cannot do is hang a worker: the head and every -/// chunk are bounded by [`DOWNLOAD_IDLE_TIMEOUT`] (see [`next_chunk`]), and a -/// transfer that keeps trickling but never finishes is bounded by -/// [`DOWNLOAD_TOTAL_TIMEOUT`]. -static MEDIA_CLIENT: LazyLock = LazyLock::new(|| build_client(None)); - -/// The error a download reports when it stops making progress. -fn download_stalled() -> FetchError { - FetchError::Transient(format!( - "download stalled for {}s", - DOWNLOAD_IDLE_TIMEOUT.as_secs() - )) -} - -/// Sends a media-download request: the response head must arrive within the -/// idle window, and a non-2xx status is classified by [`download_status_error`]. -async fn send_download(request: reqwest::RequestBuilder) -> Result { - let response = match tokio::time::timeout(DOWNLOAD_IDLE_TIMEOUT, request.send()).await { - Ok(Ok(response)) => response, - Ok(Err(e)) => return Err(e.into()), - Err(_) => return Err(download_stalled()), - }; - if response.status().is_success() { - Ok(response) - } else { - Err(download_status_error(response.status())) - } -} - -/// One body chunk, or `None` at the end. A body that stops delivering is a -/// transient download error rather than a hang. -async fn next_chunk(response: &mut reqwest::Response) -> Result, FetchError> { - match tokio::time::timeout(DOWNLOAD_IDLE_TIMEOUT, response.chunk()).await { - Ok(Ok(chunk)) => Ok(chunk), - Ok(Err(e)) => Err(e.into()), - Err(_) => Err(download_stalled()), - } -} - /// Whether a usable `ffmpeg` binary is on PATH (probed once). Shared by the /// pixiv ugoira encoder and the bsky HLS remuxer. static FFMPEG_AVAILABLE: LazyLock = LazyLock::new(|| { @@ -501,7 +387,7 @@ type SiteFuture<'a, T, E = FetchError> = Pin>> = LazyLock::new(|| { +pub(crate) static SITES: LazyLock>> = LazyLock::new(|| { vec![ Box::new(twitter::TwitterSite), Box::new(bsky::BskySite), @@ -624,177 +510,6 @@ pub fn needs_media_headers(url: &str) -> bool { SITES.iter().any(|site| site.media_headers(url).is_some()) } -/// Applies every site's media-header rule to a download request (pixiv's -/// `Referer` for pximg.net hotlink protection). Sites contribute via their -/// `media_headers(url)` — the central download code carries no per-site logic. -/// Whether an address must never be fetched. Media URLs come from a site's own -/// API response and the bytes are uploaded to Telegram, so following one into -/// the host's own network would turn the bot into a proxy for it: a cloud -/// metadata endpoint read back into a chat. -fn blocked_ip(addr: std::net::IpAddr) -> bool { - use std::net::IpAddr; - match addr { - IpAddr::V4(v4) => { - let [a, b, ..] = v4.octets(); - v4.is_private() // 10/8, 172.16/12, 192.168/16 - || v4.is_loopback() // 127/8 - || v4.is_link_local() // 169.254/16 — the cloud metadata range - || v4.is_unspecified() - || v4.is_broadcast() - || v4.is_documentation() - || v4.is_multicast() - // Ranges the std helpers do not cover: carrier-grade NAT and - // benchmarking. - || (a == 100 && (64..=127).contains(&b)) - || (a == 198 && (18..=19).contains(&b)) - } - IpAddr::V6(v6) => { - let [first, ..] = v6.segments(); - v6.is_loopback() - || v6.is_unspecified() - || v6.is_multicast() - || (first & 0xfe00) == 0xfc00 // unique local fc00::/7 - || (first & 0xffc0) == 0xfe80 // link local fe80::/10 - || v6.to_ipv4_mapped().is_some_and(|v4| blocked_ip(IpAddr::V4(v4))) - } - } -} - -/// `localhost` (and anything under it) plus the mDNS `.local` suffix: names that -/// only ever mean this machine. -fn is_local_name(name: &str) -> bool { - let name = name.trim_end_matches('.').to_ascii_lowercase(); - name == "localhost" || name.ends_with(".localhost") || name.ends_with(".local") -} - -/// Whether a media URL may be requested at all: http(s), and a host that is no -/// address or name of the host's own network. Applied to the URL a download -/// starts from *and* to every redirect hop. -/// -/// The residual gap is DNS rebinding — a name the site controls that resolves to -/// a private address. Closing it needs a `reqwest::dns::Resolve` wrapper -/// filtering resolved addresses; it is deliberately not installed, because the -/// same resolver also resolves the operator's proxy host and `TELOXIDE_PROXY` -/// is routinely a LAN address, so the guard would take down a working -/// deployment to block a far less likely attack. -fn media_url_allowed(url: &url::Url) -> bool { - if !matches!(url.scheme(), "http" | "https") { - return false; - } - match url.host() { - Some(url::Host::Ipv4(v4)) => !blocked_ip(v4.into()), - Some(url::Host::Ipv6(v6)) => !blocked_ip(v6.into()), - Some(url::Host::Domain(name)) => !is_local_name(name), - None => false, - } -} - -/// Prepares a media download: refuses a URL pointing inside the host's own -/// network ([`FetchError::Blocked`], permanent — the same URL would be refused -/// again), then applies the site's media headers. One choke point so every -/// download path gets the guard. -fn media_request(url: &str) -> Result { - let parsed = url::Url::parse(url).map_err(|e| { - log::warn!("media url is not a url: {e}"); - FetchError::Blocked - })?; - if !media_url_allowed(&parsed) { - log::warn!("refusing to fetch media from the host's own network"); - return Err(FetchError::Blocked); - } - Ok(apply_media_headers(MEDIA_CLIENT.get(parsed), url)) -} - -fn apply_media_headers(mut request: reqwest::RequestBuilder, url: &str) -> reqwest::RequestBuilder { - for site in SITES.iter() { - if let Some(headers) = site.media_headers(url) { - for (name, value) in headers { - request = request.header(name, value); - } - } - } - request -} - -/// Maps a media download's HTTP status onto the same classes the site -/// adapters use, so callers can tell "try again" from "this URL is dead": -/// 4xx is a property of the media (gone, refused by the host), while 429/5xx -/// is a property of the moment. A transport error never reaches this — it -/// fails in `send()` and stays [`FetchError::Http`]. -fn download_status_error(status: reqwest::StatusCode) -> FetchError { - match status.as_u16() { - 401 | 403 => FetchError::Blocked, - 404 | 410 => FetchError::NotFound, - _ => FetchError::Transient(format!("media status {status}")), - } -} - -/// Downloads a media file with a hard size cap: the body is streamed and the -/// download aborts with [`FetchError::TooLarge`] the moment the cap is -/// crossed (or when a declared Content-Length already exceeds it). Keeps the -/// bot from buffering arbitrarily large bodies into memory — the size check -/// the bot's upload fallback needs is the one here, not a probe of its own. -/// -/// This is the bot's download path for the upload fallback: when Telegram -/// cannot fetch a media URL itself (hotlink protection), the bot downloads -/// the file and uploads it via multipart. Site-appropriate headers come from -/// each site's `media_headers` (pixiv image hosts need `Referer`). -pub async fn download_media_limited(url: &str, max_bytes: u64) -> Result { - let response = send_download(media_request(url)?).await?; - if let Some(len) = response.content_length() - && len > max_bytes - { - return Err(FetchError::TooLarge); - } - let mut response = response; - let mut buf = Vec::new(); - let started = std::time::Instant::now(); - while let Some(chunk) = next_chunk(&mut response).await? { - if started.elapsed() > DOWNLOAD_TOTAL_TIMEOUT { - return Err(download_too_slow()); - } - buf.extend_from_slice(&chunk); - if buf.len() as u64 > max_bytes { - return Err(FetchError::TooLarge); - } - } - Ok(bytes::Bytes::from(buf)) -} - -/// Streams a download to `out`, aborting with [`FetchError::TooLarge`] the -/// moment the body crosses `max_bytes` (or when a declared Content-Length -/// already exceeds it). Unlike [`download_media_limited`] the body is never -/// buffered in memory — used for large files (e.g. the pixiv ugoira frame -/// zip, which can be hundreds of MB) that would otherwise spike RAM. -/// Returns the number of bytes written. -pub async fn download_media_to_file( - url: &str, - max_bytes: u64, - out: &mut std::fs::File, -) -> Result { - use std::io::Write; - let response = send_download(media_request(url)?).await?; - if let Some(len) = response.content_length() - && len > max_bytes - { - return Err(FetchError::TooLarge); - } - let mut response = response; - let mut total: u64 = 0; - let started = std::time::Instant::now(); - while let Some(chunk) = next_chunk(&mut response).await? { - if started.elapsed() > DOWNLOAD_TOTAL_TIMEOUT { - return Err(download_too_slow()); - } - total += chunk.len() as u64; - if total > max_bytes { - return Err(FetchError::TooLarge); - } - out.write_all(&chunk).map_err(FetchError::Io)?; - } - Ok(total) -} - #[cfg(test)] mod tests { use super::*; @@ -988,100 +703,6 @@ mod tests { assert!(out.chars().count() <= MAX_CAPTION_CHARS); } - #[test] - fn blocked_addresses_are_the_hosts_own_network() { - for addr in [ - "127.0.0.1", - "10.0.0.1", - "172.16.0.1", - "192.168.1.1", - "169.254.169.254", // cloud metadata - "0.0.0.0", - "255.255.255.255", - "100.64.0.1", // carrier-grade NAT - "198.18.0.1", // benchmarking - "::1", - "::", - "fc00::1", - "fe80::1", - "::ffff:127.0.0.1", - ] { - assert!(blocked_ip(addr.parse().unwrap()), "{addr}"); - } - for addr in [ - "1.1.1.1", - "93.184.216.34", - "2606:4700::1111", - "::ffff:1.1.1.1", - ] { - assert!(!blocked_ip(addr.parse().unwrap()), "{addr}"); - } - } - - #[test] - fn media_urls_inside_the_host_are_refused() { - for url in [ - "http://127.0.0.1:9/x", - "http://169.254.169.254/latest/meta-data/", - "http://[::1]:9/x", - "https://localhost/", - "https://prompt.localhost/x", - "https://printer.local/x", - "file:///etc/passwd", - "gopher://example.com/1", - ] { - let parsed = url::Url::parse(url).unwrap(); - assert!(!media_url_allowed(&parsed), "{url}"); - } - // Real media hosts and any public address stay fetchable. - for url in [ - "https://i.pximg.net/img-original/img/1.jpg", - "https://cdn.bsky.app/img/feed_thumbnail/plain/x", - "http://example.com/a", - "https://93.184.216.34/a", - ] { - let parsed = url::Url::parse(url).unwrap(); - assert!(media_url_allowed(&parsed), "{url}"); - } - } - - /// The redirect-hop guard, against a public redirector: the initial URL is - /// checked by [`media_request`], but a redirect is the part of the path a - /// third-party response actually controls. - #[tokio::test] - #[ignore = "live network: requires outbound HTTPS to httpbin.org"] - async fn live_redirect_into_the_hosts_network_is_refused() { - let url = "https://httpbin.org/redirect-to?url=http://169.254.169.254/latest/meta-data/"; - match download_media_limited(url, u64::MAX).await.unwrap_err() { - // A policy refusal reaches the caller wrapped by reqwest. - FetchError::Http(e) => assert!(e.is_redirect(), "got {e}"), - FetchError::Blocked => {} - other => panic!("expected a refusal, got {other:?}"), - } - } - - #[tokio::test] - async fn a_download_into_the_hosts_network_is_refused() { - // Refused on the URL alone: nothing has to be listening (or leaking) at - // the metadata endpoint for this to hold, and the class is permanent so - // the send path does not retry it. - for url in [ - "http://169.254.169.254/latest/meta-data/", - "http://127.0.0.1:9/secret", - ] { - let err = download_media_limited(url, u64::MAX).await.unwrap_err(); - assert!(matches!(err, FetchError::Blocked), "{url}: got {err:?}"); - } - // A malformed URL is refused the same way instead of becoming a - // retryable transport error. - assert!(matches!( - download_media_limited("not a url", u64::MAX) - .await - .unwrap_err(), - FetchError::Blocked - )); - } - #[tokio::test] async fn unsupported_urls_return_none() { // Neither a URL no site pattern matches nor a string that is no URL at @@ -1136,29 +757,4 @@ mod tests { Some("pixiv:1".into()) ); } - - #[tokio::test] - async fn download_media_pixiv_original_with_referer() { - // Proves the Referer header is attached for i.pximg.net: a header-less - // GET to a pixiv original URL is rejected with 403. - // Empty-string check too: an unset CI secret arrives as "" (GitHub - // Actions), which would otherwise run the test tokenless and fail. - if std::env::var("PIXIV_REFRESH_TOKEN") - .ok() - .filter(|s| !s.is_empty()) - .is_none() - { - eprintln!("skipping: no PIXIV_REFRESH_TOKEN"); - return; - } - let illustration = pixiv::fetch(126839080).await.unwrap(); - let fetched: Fetched = illustration.into(); - let url = match fetched.media.first() { - Some(crate::media::Media::Illustration { url, .. }) => url.clone(), - other => panic!("expected illustration media, got {other:?}"), - }; - assert!(url.contains("i.pximg.net")); - let bytes = download_media_limited(&url, u64::MAX).await.unwrap(); - assert!(!bytes.is_empty()); - } }