fix(analysis): cap OCR decoded pixels to stop decompression-bomb OOM

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
MechaCat02
2026-07-01 07:12:40 +02:00
parent 83c2899373
commit 66ae4b221b
8 changed files with 99 additions and 10 deletions

View File

@@ -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

2
backend/Cargo.lock generated
View File

@@ -1558,7 +1558,7 @@ checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4"
[[package]]
name = "mangalord"
version = "0.90.0"
version = "0.90.1"
dependencies = [
"anyhow",
"argon2",

View File

@@ -1,6 +1,6 @@
[package]
name = "mangalord"
version = "0.90.0"
version = "0.90.1"
edition = "2021"
default-run = "mangalord"

View File

@@ -60,13 +60,23 @@ pub fn lines_to_analysis(lines: Vec<String>) -> 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<Self> {
///
/// `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<Self> {
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<image::RgbImage> {
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<Vec<String>> {
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<u8> {
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}"
);
}
}

View File

@@ -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 {

View File

@@ -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),

View File

@@ -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,
})
}
}

View File

@@ -1,6 +1,6 @@
{
"name": "mangalord-frontend",
"version": "0.90.0",
"version": "0.90.1",
"private": true,
"type": "module",
"scripts": {