From 1765936055bb62c7321677d5a19a6252f1020291 Mon Sep 17 00:00:00 2001 From: Ryan Fowler Date: Mon, 25 May 2026 19:43:52 -0700 Subject: [PATCH] Harden image decoding with preallocation limits --- src/image/mod.rs | 63 +++++++++++++++++++++++++++++++++++++++++--- tests/integration.rs | 36 +++++++++++++------------ 2 files changed, 79 insertions(+), 20 deletions(-) diff --git a/src/image/mod.rs b/src/image/mod.rs index b1f2c16..441ff3d 100644 --- a/src/image/mod.rs +++ b/src/image/mod.rs @@ -11,6 +11,10 @@ use ::image::imageops::FilterType; use ::image::{ColorType, DynamicImage, GenericImageView, ImageEncoder, Rgba, RgbaImage}; use base64::Engine; +const IMAGE_DIMENSION_LIMIT: u32 = 8192; +const IMAGE_DECODE_MAX_ALLOC_BYTES: u64 = + IMAGE_DIMENSION_LIMIT as u64 * IMAGE_DIMENSION_LIMIT as u64 * 4; + #[derive(Debug, thiserror::Error)] pub enum ImageError { #[error("{0}")] @@ -128,10 +132,11 @@ fn decode_image(bytes: &[u8], decode_mode: DecodeMode) -> Result Result { - const LIMIT: u32 = 8192; - let img = ::image::load_from_memory(bytes)?; + let mut reader = ::image::ImageReader::new(Cursor::new(bytes)).with_guessed_format()?; + reader.limits(image_decode_limits()); + let img = reader.decode()?; let (width, height) = img.dimensions(); - if width > LIMIT || height > LIMIT { + if width > IMAGE_DIMENSION_LIMIT || height > IMAGE_DIMENSION_LIMIT { return Err(ImageError::Message(format!( "image dimensions are too large {width}x{height}" ))); @@ -139,6 +144,14 @@ fn decode_image_std(bytes: &[u8]) -> Result { Ok(img) } +fn image_decode_limits() -> ::image::Limits { + let mut limits = ::image::Limits::default(); + limits.max_image_width = Some(IMAGE_DIMENSION_LIMIT); + limits.max_image_height = Some(IMAGE_DIMENSION_LIMIT); + limits.max_alloc = Some(IMAGE_DECODE_MAX_ALLOC_BYTES); + limits +} + #[derive(Clone, Copy)] struct Adaptor { name: &'static str, @@ -992,6 +1005,17 @@ mod tests { ); } + #[test] + fn decode_rejects_oversized_png_dimensions_before_full_allocation() { + let bytes = png_with_dimensions(IMAGE_DIMENSION_LIMIT + 1, IMAGE_DIMENSION_LIMIT + 1); + let err = decode_image_std(&bytes).unwrap_err(); + + assert!(matches!( + err, + ImageError::Decode(::image::ImageError::Limits(_)) + )); + } + #[test] fn builtin_decode_mode_rejects_unsupported_bytes_without_external_adaptors() { let err = decode_image(b"not an image", DecodeMode::BuiltIn).unwrap_err(); @@ -1038,6 +1062,39 @@ mod tests { assert_eq!(Emulator::Apple.protocol(), Protocol::Block); } + fn png_with_dimensions(width: u32, height: u32) -> Vec { + let mut bytes = Vec::from(&[137, 80, 78, 71, 13, 10, 26, 10]); + let mut ihdr = Vec::new(); + ihdr.extend_from_slice(&width.to_be_bytes()); + ihdr.extend_from_slice(&height.to_be_bytes()); + ihdr.extend_from_slice(&[8, 6, 0, 0, 0]); + append_png_chunk(&mut bytes, b"IHDR", &ihdr); + append_png_chunk(&mut bytes, b"IEND", &[]); + bytes + } + + fn append_png_chunk(bytes: &mut Vec, kind: &[u8; 4], data: &[u8]) { + bytes.extend_from_slice(&u32::try_from(data.len()).unwrap().to_be_bytes()); + bytes.extend_from_slice(kind); + bytes.extend_from_slice(data); + bytes.extend_from_slice(&png_crc32(kind, data).to_be_bytes()); + } + + fn png_crc32(kind: &[u8; 4], data: &[u8]) -> u32 { + let mut crc = 0xffff_ffff_u32; + for byte in kind.iter().chain(data) { + crc ^= u32::from(*byte); + for _ in 0..8 { + crc = if crc & 1 == 0 { + crc >> 1 + } else { + (crc >> 1) ^ 0xedb8_8320 + }; + } + } + !crc + } + fn jpeg_with_orientation(orientation: u16) -> Vec { let mut bytes = Vec::new(); bytes.extend_from_slice(&0xffd8_u16.to_be_bytes()); diff --git a/tests/integration.rs b/tests/integration.rs index 0d88992..210056d 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -14,7 +14,7 @@ use std::collections::HashMap; use std::env; use std::fs; use std::io::{BufRead, BufReader, Read, Write}; -use std::net::{Ipv4Addr, TcpListener, TcpStream, UdpSocket}; +use std::net::{Ipv4Addr, Shutdown, TcpListener, TcpStream, UdpSocket}; use std::path::{Path, PathBuf}; use std::process::{Command, ExitStatus, Stdio}; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -24,6 +24,7 @@ use std::time::{Duration, Instant}; use tempfile::TempDir; const FAST_RETRY_DELAY: &str = "0.000001"; +const PARTIAL_REPLAY_BODY_PREFIX_BYTES: usize = 1024 * 1024; #[cfg(unix)] use image::ImageEncoder; @@ -330,26 +331,27 @@ impl PartialBodyReplayServer { if first { first = false; let headers = headers.clone(); + let _ = write!(stream, "HTTP/1.1 {status} {reason}\r\n"); + for (name, value) in &headers { + let _ = write!(stream, "{name}: {value}\r\n"); + } + let _ = write!( + stream, + "Content-Length: 1073741824\r\nConnection: close\r\n\r\n" + ); + let body = vec![b'x'; PARTIAL_REPLAY_BODY_PREFIX_BYTES]; + let _ = stream.write_all(&body); + let _ = stream.flush(); thread::spawn(move || { - let _ = write!(stream, "HTTP/1.1 {status} {reason}\r\n"); - for (name, value) in &headers { - let _ = write!(stream, "{name}: {value}\r\n"); - } - let _ = write!( - stream, - "Content-Length: 1073741824\r\nConnection: close\r\n\r\n" - ); - let chunk = vec![b'x'; 16 * 1024]; - for _ in 0..128 { - if stream.write_all(&chunk).is_err() { - break; - } - } - let _ = stream.flush(); thread::sleep(Duration::from_secs(5)); + let _ = stream.shutdown(Shutdown::Both); }); } else { - write_response(&mut stream, TestResponse::ok(final_body)); + write_response( + &mut stream, + TestResponse::ok(final_body).header("Connection", "close"), + ); + let _ = stream.shutdown(Shutdown::Both); break; } }