Apply image decoder limits before allocation

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,17 +132,26 @@ fn decode_image(bytes: &[u8], decode_mode: DecodeMode) -> Result<DynamicImage, I
 }
 
 fn decode_image_std(bytes: &[u8]) -> Result<DynamicImage, ImageError> {
-    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}"
         )));
     }
     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<u8> {
+        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<u8>, 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<u8> {
         let mut bytes = Vec::new();
         bytes.extend_from_slice(&0xffd8_u16.to_be_bytes());

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;
                         }
                     }