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:
@@ -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
2
backend/Cargo.lock
generated
@@ -1558,7 +1558,7 @@ checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4"
|
||||
|
||||
[[package]]
|
||||
name = "mangalord"
|
||||
version = "0.90.0"
|
||||
version = "0.90.1"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"argon2",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "mangalord"
|
||||
version = "0.90.0"
|
||||
version = "0.90.1"
|
||||
edition = "2021"
|
||||
default-run = "mangalord"
|
||||
|
||||
|
||||
@@ -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}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "mangalord-frontend",
|
||||
"version": "0.90.0",
|
||||
"version": "0.90.1",
|
||||
"private": true,
|
||||
"type": "module",
|
||||
"scripts": {
|
||||
|
||||
Reference in New Issue
Block a user