From 66ae4b221b76cbc07611226cfa6b1624a284f7ad Mon Sep 17 00:00:00 2001 From: MechaCat02 Date: Wed, 1 Jul 2026 07:12:40 +0200 Subject: [PATCH] fix(analysis): cap OCR decoded pixels to stop decompression-bomb OOM Co-Authored-By: Claude Opus 4.8 --- .env.example | 5 +++ backend/Cargo.lock | 2 +- backend/Cargo.toml | 2 +- backend/src/analysis/ocr.rs | 83 +++++++++++++++++++++++++++++++++---- backend/src/app.rs | 1 + backend/src/config.rs | 12 ++++++ backend/src/settings.rs | 2 + frontend/package.json | 2 +- 8 files changed, 99 insertions(+), 10 deletions(-) diff --git a/.env.example b/.env.example index 7ffb215..f151260 100644 --- a/.env.example +++ b/.env.example @@ -328,6 +328,10 @@ ANALYSIS_API_KEY= # beyond cap). Default 16. # ANALYSIS_MAX_IMAGE_BYTES Per-image byte cap before downscale. # Default 8388608 (8 MiB). +# ANALYSIS_OCR_MAX_DECODE_PIXELS Decompression-bomb guard for the OCR +# backend: hard cap on a page's *decoded* pixel +# count (encoded size is bounded separately). +# Default 100000000 (100 MP). # ANALYSIS_RESPONSE_FORMAT `json_schema` | `json_object` | `none`. # Default `json_schema`. # ANALYSIS_FREQUENCY_PENALTY Discourages repetition loops. Default 0.3. @@ -340,6 +344,7 @@ ANALYSIS_SLICE_OVERLAP=0.12 ANALYSIS_TALL_ASPECT=1.6 ANALYSIS_MAX_SLICES=16 ANALYSIS_MAX_IMAGE_BYTES=8388608 +ANALYSIS_OCR_MAX_DECODE_PIXELS=100000000 ANALYSIS_RESPONSE_FORMAT=json_schema ANALYSIS_FREQUENCY_PENALTY=0.3 ANALYSIS_TEMPERATURE=0.0 diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 34c09f0..c89fd73 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -1558,7 +1558,7 @@ checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4" [[package]] name = "mangalord" -version = "0.90.0" +version = "0.90.1" dependencies = [ "anyhow", "argon2", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 81ba63c..75148f2 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "mangalord" -version = "0.90.0" +version = "0.90.1" edition = "2021" default-run = "mangalord" diff --git a/backend/src/analysis/ocr.rs b/backend/src/analysis/ocr.rs index e60556b..fa8cbba 100644 --- a/backend/src/analysis/ocr.rs +++ b/backend/src/analysis/ocr.rs @@ -60,13 +60,23 @@ pub fn lines_to_analysis(lines: Vec) -> VisionAnalysis { /// recognize path borrows `&self`. pub struct OcrsEngine { engine: ocrs::OcrEngine, + /// Hard cap on decoded pixel count (decompression-bomb guard). See + /// [`decode_rgb8_within`]. + max_decode_pixels: u64, } impl OcrsEngine { /// Load the detection + recognition models from disk and build the engine. /// Fails (at startup) if either model file is missing or unreadable, so a /// misconfigured path is a loud boot error rather than a per-page failure. - pub fn from_model_paths(detection: &str, recognition: &str) -> anyhow::Result { + /// + /// `max_decode_pixels` bounds the decoded image size (see + /// [`decode_rgb8_within`]) — wired from `AnalysisConfig::ocr_max_decode_pixels`. + pub fn from_model_paths( + detection: &str, + recognition: &str, + max_decode_pixels: u64, + ) -> anyhow::Result { use anyhow::Context; let detection_model = rten::Model::load_file(detection) .with_context(|| format!("load ocrs detection model {detection}"))?; @@ -78,17 +88,41 @@ impl OcrsEngine { ..Default::default() }) .context("construct ocrs engine")?; - Ok(Self { engine }) + Ok(Self { engine, max_decode_pixels }) } } +/// Decode encoded image bytes to RGB8 while refusing decompression bombs. +/// +/// `image::load_from_memory` allocates the full decoded buffer up front, so a +/// tiny file declaring enormous dimensions (e.g. a 50000×50000 PNG) inflates +/// to billions of bytes and OOM-kills the process. We cap the decoder's +/// allocation at `max_decode_pixels` worth of RGBA (4 bytes/px headroom over +/// the RGB8 result), which the `image` crate checks against the header +/// *before* allocating — so an over-size image fails fast instead of dying. +fn decode_rgb8_within(image: &[u8], max_decode_pixels: u64) -> anyhow::Result { + use anyhow::Context; + use std::io::Cursor; + + let mut reader = image::ImageReader::new(Cursor::new(image)) + .with_guessed_format() + .context("guess page image format for OCR")?; + let mut limits = image::Limits::default(); + // 4 bytes/px (RGBA) gives headroom over the eventual RGB8 buffer and any + // single intermediate the decoder allocates per pixel. + limits.max_alloc = Some(max_decode_pixels.saturating_mul(4)); + reader.limits(limits); + Ok(reader + .decode() + .context("decode page image for OCR")? + .into_rgb8()) +} + impl OcrEngine for OcrsEngine { fn recognize(&self, image: &[u8]) -> anyhow::Result> { - use anyhow::Context; - // Decode to RGB8 so `ImageSource` gets a known channel layout. - let rgb = image::load_from_memory(image) - .context("decode page image for OCR")? - .into_rgb8(); + // Decode to RGB8 so `ImageSource` gets a known channel layout, + // bounding the decoded size against the decompression-bomb cap. + let rgb = decode_rgb8_within(image, self.max_decode_pixels)?; let source = ocrs::ImageSource::from_bytes(rgb.as_raw(), rgb.dimensions()) .map_err(|e| anyhow::anyhow!("build OCR image source: {e}"))?; let input = self.engine.prepare_input(source)?; @@ -198,4 +232,39 @@ mod tests { let v = lines_to_analysis(Vec::new()); assert!(v.ocr_results.is_empty()); } + + /// A minimal valid PNG of `w`×`h` (single black pixel scaled via IHDR is + /// not valid; instead encode a real tiny image, then we patch the IHDR + /// dimensions for the bomb case). For the happy path we just encode a real + /// small image. + fn encode_png(w: u32, h: u32) -> Vec { + let img = image::RgbImage::new(w, h); + let mut buf = std::io::Cursor::new(Vec::new()); + image::DynamicImage::ImageRgb8(img) + .write_to(&mut buf, image::ImageFormat::Png) + .unwrap(); + buf.into_inner() + } + + #[test] + fn decode_rgb8_within_accepts_normal_page() { + // A perfectly ordinary page decodes fine under a generous cap. + let png = encode_png(64, 96); + let rgb = decode_rgb8_within(&png, 100_000_000).unwrap(); + assert_eq!(rgb.dimensions(), (64, 96)); + } + + #[test] + fn decode_rgb8_within_rejects_oversize_image() { + // The same image, but the cap is set below its pixel count: the + // allocation limit must trip rather than the decode succeeding. This + // is the decompression-bomb guard in miniature — a real bomb declares + // a huge size in a few bytes; here we shrink the budget instead. + let png = encode_png(2000, 2000); // 4 MP + let err = decode_rgb8_within(&png, 1_000).unwrap_err(); + assert!( + err.to_string().contains("decode page image"), + "expected a decode error, got: {err}" + ); + } } diff --git a/backend/src/app.rs b/backend/src/app.rs index aeeeb7e..868a89c 100644 --- a/backend/src/app.rs +++ b/backend/src/app.rs @@ -440,6 +440,7 @@ async fn spawn_analysis_daemon( let engine = crate::analysis::ocr::OcrsEngine::from_model_paths( &cfg.ocr_detection_model, &cfg.ocr_recognition_model, + cfg.ocr_max_decode_pixels, ) .context("build ocrs engine")?; let dispatcher = Arc::new(crate::analysis::ocr::OcrAnalyzeDispatcher { diff --git a/backend/src/config.rs b/backend/src/config.rs index 750f0bc..a084db0 100644 --- a/backend/src/config.rs +++ b/backend/src/config.rs @@ -202,6 +202,13 @@ pub struct AnalysisConfig { /// Hard cap on a page image's stored size; larger pages are skipped /// (`ANALYSIS_MAX_IMAGE_BYTES`). pub max_image_bytes: usize, + /// Hard cap on a page image's **decoded** pixel count for the OCR backend + /// (`ANALYSIS_OCR_MAX_DECODE_PIXELS`). `max_image_bytes` only bounds the + /// *encoded* size; without a decode bound a tiny image declaring + /// 50000×50000 inflates to billions of bytes and OOM-kills the worker + /// (decompression bomb). Generous by default (100 MP) so legitimately + /// tall, un-sliced manga pages still decode. + pub ocr_max_decode_pixels: u64, /// Output-constraint mode (`ANALYSIS_RESPONSE_FORMAT`): /// `json_schema` (default) | `json_object` | `none`. pub response_format: ResponseFormat, @@ -251,6 +258,7 @@ impl Default for AnalysisConfig { tall_aspect_threshold: 1.6, max_slices: 16, max_image_bytes: 8 * 1024 * 1024, + ocr_max_decode_pixels: 100_000_000, response_format: ResponseFormat::JsonSchema, frequency_penalty: 0.3, temperature: 0.0, @@ -327,6 +335,10 @@ impl AnalysisConfig { .max(1.0), max_slices: env_usize("ANALYSIS_MAX_SLICES", d.max_slices).max(1), max_image_bytes: env_usize("ANALYSIS_MAX_IMAGE_BYTES", d.max_image_bytes), + ocr_max_decode_pixels: env_u64( + "ANALYSIS_OCR_MAX_DECODE_PIXELS", + d.ocr_max_decode_pixels, + ), response_format: std::env::var("ANALYSIS_RESPONSE_FORMAT") .map(|s| ResponseFormat::from_str(&s)) .unwrap_or(d.response_format), diff --git a/backend/src/settings.rs b/backend/src/settings.rs index b79c576..417a41c 100644 --- a/backend/src/settings.rs +++ b/backend/src/settings.rs @@ -473,6 +473,8 @@ impl AnalysisSettings { backend: base.backend, ocr_detection_model: base.ocr_detection_model.clone(), ocr_recognition_model: base.ocr_recognition_model.clone(), + // Env-only decompression-bomb decode cap, carried from the base. + ocr_max_decode_pixels: base.ocr_max_decode_pixels, }) } } diff --git a/frontend/package.json b/frontend/package.json index 5478056..637de24 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -1,6 +1,6 @@ { "name": "mangalord-frontend", - "version": "0.90.0", + "version": "0.90.1", "private": true, "type": "module", "scripts": {