From 4d25bef191477a73e8b0808f8df363d6494ef5f3 Mon Sep 17 00:00:00 2001 From: Himanshu Kumar Date: Wed, 14 Oct 2026 06:33:38 +0000 Subject: [PATCH] improve: avoid transforming pixels during overlay image inspection --- apps/desktop-gpui/kernels/Cargo.toml | 4 + .../benches/overlay_image_inspection.rs | 178 ++++++++++++ apps/desktop-gpui/kernels/src/codec.rs | 17 +- .../kernels/src/codec_overlay_tests.rs | 269 ++++++++++++++++++ 4 files changed, 465 insertions(+), 3 deletions(-) create mode 100644 apps/desktop-gpui/kernels/benches/overlay_image_inspection.rs create mode 100644 apps/desktop-gpui/kernels/src/codec_overlay_tests.rs diff --git a/apps/desktop-gpui/kernels/Cargo.toml b/apps/desktop-gpui/kernels/Cargo.toml index 6950a4c77b4..4d0031349ba 100644 --- a/apps/desktop-gpui/kernels/Cargo.toml +++ b/apps/desktop-gpui/kernels/Cargo.toml @@ -15,3 +15,7 @@ tracing = "0.1" [[bench]] name = "screenshot_blit" harness = false + +[[bench]] +name = "overlay_image_inspection" +harness = false diff --git a/apps/desktop-gpui/kernels/benches/overlay_image_inspection.rs b/apps/desktop-gpui/kernels/benches/overlay_image_inspection.rs new file mode 100644 index 00000000000..03d5515cb16 --- /dev/null +++ b/apps/desktop-gpui/kernels/benches/overlay_image_inspection.rs @@ -0,0 +1,178 @@ +use cap_gpui_kernels::codec::inspect_overlay_image; +use image::{ImageDecoder as _, ImageEncoder as _, metadata::Orientation}; +use std::{hint::black_box, time::Instant}; + +type Inspect = fn(&[u8]) -> Result<(&'static str, u32, u32), String>; + +struct Fixture { + name: String, + encoded: Vec, + dimensions: (u32, u32), + orientation: Orientation, + extension: &'static str, +} + +fn exif(orientation: u8) -> Vec { + let mut tiff = b"MM\0*\0\0\0\x08\0\x01\x01\x12\0\x03\0\0\0\x01\0\x01\0\0\0\0\0\0".to_vec(); + tiff[19] = orientation; + tiff +} + +fn pixel(x: u32, y: u32) -> [u8; 3] { + let detail = (x.wrapping_mul(1_664_525) ^ y.wrapping_mul(1_013_904_223)) >> 24; + [ + (x / 17 + y / 11 + detail) as u8, + (x / 29 + y / 7 + detail) as u8, + (x / 13 + y / 31 + detail) as u8, + ] +} + +fn fixtures() -> Vec { + let mut fixtures = Vec::new(); + for (width, height) in [(640, 480), (4032, 3024)] { + let rgb = image::RgbImage::from_fn(width, height, |x, y| image::Rgb(pixel(x, y))); + let mut jpeg = Vec::new(); + image::codecs::jpeg::JpegEncoder::new_with_quality(&mut jpeg, 85) + .encode_image(&rgb) + .unwrap(); + for orientation in [None, Some(1), Some(2), Some(3), Some(6), Some(8)] { + let encoded = if let Some(orientation) = orientation { + let tiff = exif(orientation); + let mut encoded = jpeg[..2].to_vec(); + encoded.extend_from_slice(&[0xff, 0xe1]); + encoded.extend_from_slice(&((tiff.len() + 8) as u16).to_be_bytes()); + encoded.extend_from_slice(b"Exif\0\0"); + encoded.extend_from_slice(&tiff); + encoded.extend_from_slice(&jpeg[2..]); + encoded + } else { + jpeg.clone() + }; + fixtures.push(Fixture { + name: format!( + "jpeg-{width}x{height}-{}", + orientation.map_or_else(|| "no-exif".into(), |value| format!("exif{value}")) + ), + encoded, + dimensions: (width, height), + orientation: Orientation::from_exif(orientation.unwrap_or(1)).unwrap(), + extension: "jpg", + }); + } + } + let (width, height) = (1600, 900); + let rgba = image::RgbaImage::from_fn(width, height, |x, y| { + let [r, g, b] = pixel(x, y); + image::Rgba([r, g, b, (x + y * 7) as u8]) + }); + for orientation in [1, 6] { + let mut encoded = Vec::new(); + let mut encoder = image::codecs::png::PngEncoder::new(&mut encoded); + encoder.set_exif_metadata(exif(orientation)).unwrap(); + encoder + .write_image(&rgba, width, height, image::ExtendedColorType::Rgba8) + .unwrap(); + fixtures.push(Fixture { + name: format!("png-{width}x{height}-exif{orientation}"), + encoded, + dimensions: (width, height), + orientation: Orientation::from_exif(orientation).unwrap(), + extension: "png", + }); + } + fixtures +} + +fn timed(inspect: Inspect, encoded: &[u8]) -> u128 { + let start = Instant::now(); + drop(black_box(inspect(black_box(encoded)))); + start.elapsed().as_nanos() +} + +fn main() { + let fixtures = fixtures(); + println!("workload,width,height,encoded_bytes,sample,before_ns,after_ns"); + for fixture in fixtures { + let reader = image::ImageReader::new(std::io::Cursor::new(&fixture.encoded)) + .with_guessed_format() + .unwrap(); + let mut decoder = reader.into_decoder().unwrap(); + assert_eq!(decoder.dimensions(), fixture.dimensions, "{}", fixture.name); + assert_eq!( + decoder.orientation().unwrap(), + fixture.orientation, + "{}", + fixture.name + ); + drop(decoder); + let (width, height) = fixture.dimensions; + let expected = if matches!( + fixture.orientation, + Orientation::Rotate90 | Orientation::Rotate270 + ) { + (fixture.extension, height, width) + } else { + (fixture.extension, width, height) + }; + assert_eq!(baseline(&fixture.encoded), Ok(expected), "{}", fixture.name); + assert_eq!( + inspect_overlay_image(&fixture.encoded), + Ok(expected), + "{}", + fixture.name + ); + for sample in 0..35 { + let (before, after) = if sample % 2 == 0 { + let before = timed(baseline, &fixture.encoded); + (before, timed(inspect_overlay_image, &fixture.encoded)) + } else { + let after = timed(inspect_overlay_image, &fixture.encoded); + (timed(baseline, &fixture.encoded), after) + }; + if sample >= 5 { + println!( + "{},{width},{height},{},{},{before},{after}", + fixture.name, + fixture.encoded.len(), + sample - 5 + ); + } + } + } +} + +fn baseline(encoded: &[u8]) -> Result<(&'static str, u32, u32), String> { + use image::ImageDecoder as _; + let mut reader = image::ImageReader::new(std::io::Cursor::new(encoded)) + .with_guessed_format() + .map_err(|error| error.to_string())?; + let extension = match reader.format() { + Some(image::ImageFormat::Png) => "png", + Some(image::ImageFormat::Jpeg) => "jpg", + Some(image::ImageFormat::WebP) => "webp", + Some(image::ImageFormat::Gif) => "gif", + Some(image::ImageFormat::Bmp) => "bmp", + _ => return Err("Choose a PNG, JPEG, WebP, GIF or BMP image".into()), + }; + let mut limits = image::Limits::default(); + limits.max_alloc = Some(128 * 1024 * 1024); + limits.max_image_width = Some(32_768); + limits.max_image_height = Some(32_768); + reader.limits(limits); + let mut decoder = reader.into_decoder().map_err(|error| { + format!("Cannot decode image (maximum 32,768 pixels per side): {error}") + })?; + let (source_width, source_height) = decoder.dimensions(); + if source_width == 0 + || source_height == 0 + || u64::from(source_width) * u64::from(source_height) > 16_777_216 + || decoder.total_bytes() > 128 * 1024 * 1024 + { + return Err("Images must have at most 16,777,216 pixels (32,768 per side) and decode to at most 128 MiB".into()); + } + let orientation = decoder.orientation().map_err(|error| error.to_string())?; + let mut decoded = image::DynamicImage::from_decoder(decoder) + .map_err(|error| format!("Cannot decode image: {error}"))?; + decoded.apply_orientation(orientation); + Ok((extension, decoded.width(), decoded.height())) +} diff --git a/apps/desktop-gpui/kernels/src/codec.rs b/apps/desktop-gpui/kernels/src/codec.rs index b7ba7774c36..f4d751ee163 100644 --- a/apps/desktop-gpui/kernels/src/codec.rs +++ b/apps/desktop-gpui/kernels/src/codec.rs @@ -128,10 +128,17 @@ pub fn inspect_overlay_image(encoded: &[u8]) -> Result<(&'static str, u32, u32), return Err("Images must have at most 16,777,216 pixels (32,768 per side) and decode to at most 128 MiB".into()); } let orientation = decoder.orientation().map_err(|error| error.to_string())?; - let mut decoded = image::DynamicImage::from_decoder(decoder) + let decoded = image::DynamicImage::from_decoder(decoder) .map_err(|error| format!("Cannot decode image: {error}"))?; - decoded.apply_orientation(orientation); - Ok((extension, decoded.width(), decoded.height())) + use image::metadata::Orientation; + let (width, height) = match orientation { + Orientation::Rotate90 + | Orientation::Rotate270 + | Orientation::Rotate90FlipH + | Orientation::Rotate270FlipH => (decoded.height(), decoded.width()), + _ => (decoded.width(), decoded.height()), + }; + Ok((extension, width, height)) } pub fn decode_jpeg_to_bgra(bytes: &[u8]) -> Option { @@ -218,3 +225,7 @@ mod tests { ); } } + +#[cfg(test)] +#[path = "codec_overlay_tests.rs"] +mod overlay_inspection_tests; diff --git a/apps/desktop-gpui/kernels/src/codec_overlay_tests.rs b/apps/desktop-gpui/kernels/src/codec_overlay_tests.rs new file mode 100644 index 00000000000..9f622ea19f4 --- /dev/null +++ b/apps/desktop-gpui/kernels/src/codec_overlay_tests.rs @@ -0,0 +1,269 @@ +use super::inspect_overlay_image; +use image::{ImageDecoder as _, ImageEncoder as _, ImageFormat, metadata::Orientation}; + +fn inspect_with_pixel_orientation(encoded: &[u8]) -> Result<(&'static str, u32, u32), String> { + let mut reader = image::ImageReader::new(std::io::Cursor::new(encoded)) + .with_guessed_format() + .map_err(|error| error.to_string())?; + let extension = match reader.format() { + Some(ImageFormat::Png) => "png", + Some(ImageFormat::Jpeg) => "jpg", + Some(ImageFormat::WebP) => "webp", + Some(ImageFormat::Gif) => "gif", + Some(ImageFormat::Bmp) => "bmp", + _ => return Err("Choose a PNG, JPEG, WebP, GIF or BMP image".into()), + }; + let mut limits = image::Limits::default(); + limits.max_alloc = Some(128 * 1024 * 1024); + limits.max_image_width = Some(32_768); + limits.max_image_height = Some(32_768); + reader.limits(limits); + let mut decoder = reader.into_decoder().map_err(|error| { + format!("Cannot decode image (maximum 32,768 pixels per side): {error}") + })?; + let (source_width, source_height) = decoder.dimensions(); + if source_width == 0 + || source_height == 0 + || u64::from(source_width) * u64::from(source_height) > 16_777_216 + || decoder.total_bytes() > 128 * 1024 * 1024 + { + return Err("Images must have at most 16,777,216 pixels (32,768 per side) and decode to at most 128 MiB".into()); + } + let orientation = decoder.orientation().map_err(|error| error.to_string())?; + let mut decoded = image::DynamicImage::from_decoder(decoder) + .map_err(|error| format!("Cannot decode image: {error}"))?; + decoded.apply_orientation(orientation); + Ok((extension, decoded.width(), decoded.height())) +} + +fn assert_matches_original(encoded: &[u8]) -> Result<(&'static str, u32, u32), String> { + let unchanged = encoded.to_vec(); + let expected = inspect_with_pixel_orientation(encoded); + let actual = inspect_overlay_image(encoded); + assert_eq!(actual, expected); + assert_eq!(encoded, unchanged.as_slice()); + actual +} + +fn exif(orientation: u8, big_endian: bool) -> Vec { + if big_endian { + vec![ + b'M', + b'M', + 0, + 42, + 0, + 0, + 0, + 8, + 0, + 1, + 1, + 0x12, + 0, + 3, + 0, + 0, + 0, + 1, + 0, + orientation, + 0, + 0, + 0, + 0, + 0, + 0, + ] + } else { + vec![ + b'I', + b'I', + 42, + 0, + 8, + 0, + 0, + 0, + 1, + 0, + 0x12, + 1, + 3, + 0, + 1, + 0, + 0, + 0, + orientation, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + ] + } +} + +fn encoded_image(format: ImageFormat, width: u32, height: u32, exif: Option>) -> Vec { + let rgba = image::RgbaImage::from_fn(width, height, |x, y| { + image::Rgba([ + (x * 19 + y * 3) as u8, + (y * 23 + x * 7) as u8, + 211, + (x * 13 + y * 17) as u8, + ]) + }); + let mut encoded = Vec::new(); + match format { + ImageFormat::Jpeg => { + let mut encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut encoded, 80); + if let Some(exif) = exif { + encoder.set_exif_metadata(exif).unwrap(); + } + encoder.encode_image(&rgba).unwrap(); + } + ImageFormat::Png => { + let mut encoder = image::codecs::png::PngEncoder::new(&mut encoded); + if let Some(exif) = exif { + encoder.set_exif_metadata(exif).unwrap(); + } + encoder + .write_image( + rgba.as_raw(), + width, + height, + image::ExtendedColorType::Rgba8, + ) + .unwrap(); + } + _ => unreachable!(), + } + encoded +} + +fn decoded_orientation(encoded: &[u8]) -> Orientation { + image::ImageReader::new(std::io::Cursor::new(encoded)) + .with_guessed_format() + .unwrap() + .into_decoder() + .unwrap() + .orientation() + .unwrap() +} + +#[test] +fn all_exif_orientations_match_rotated_pixels_for_jpeg_and_png() { + for format in [ImageFormat::Jpeg, ImageFormat::Png] { + for (width, height) in [(7, 11), (11, 7), (9, 9)] { + for big_endian in [false, true] { + for value in 1..=8 { + let encoded = + encoded_image(format, width, height, Some(exif(value, big_endian))); + assert_eq!( + decoded_orientation(&encoded), + Orientation::from_exif(value).unwrap() + ); + let dimensions = if value >= 5 { + (height, width) + } else { + (width, height) + }; + let extension = if format == ImageFormat::Jpeg { + "jpg" + } else { + "png" + }; + assert_eq!( + assert_matches_original(&encoded).unwrap(), + (extension, dimensions.0, dimensions.1), + "{format:?}, EXIF {value}, big endian {big_endian}" + ); + } + } + } + } +} + +#[test] +fn absent_and_malformed_exif_keep_untransformed_dimensions() { + let truncated_exif = exif(6, false)[..10].to_vec(); + for format in [ImageFormat::Jpeg, ImageFormat::Png] { + for metadata in [ + None, + Some(b"not EXIF".to_vec()), + Some(truncated_exif.clone()), + Some(exif(0, false)), + Some(exif(9, true)), + Some(exif(255, false)), + ] { + let encoded = encoded_image(format, 7, 11, metadata); + assert_eq!(decoded_orientation(&encoded), Orientation::NoTransforms); + let (_, width, height) = assert_matches_original(&encoded).unwrap(); + assert_eq!((width, height), (7, 11)); + } + } +} + +#[test] +fn corrupt_and_truncated_images_keep_decode_errors() { + for encoded in [b"".as_slice(), b"invalid image".as_slice()] { + assert!(assert_matches_original(encoded).is_err()); + } + for format in [ImageFormat::Jpeg, ImageFormat::Png] { + let encoded = encoded_image(format, 17, 11, Some(exif(6, false))); + for end in [2, encoded.len() / 2] { + assert!(assert_matches_original(&encoded[..end]).is_err()); + } + } + let encoded = encoded_image(ImageFormat::Png, 17, 11, Some(exif(6, true))); + let data_start = encoded + .windows(4) + .position(|bytes| bytes == b"IDAT") + .unwrap() + + 4; + let mut corrupt = encoded.clone(); + corrupt[data_start] = 0; + for invalid in [corrupt, encoded[..data_start + 2].to_vec()] { + assert_eq!(decoded_orientation(&invalid), Orientation::Rotate90); + assert!( + assert_matches_original(&invalid) + .unwrap_err() + .starts_with("Cannot decode image: ") + ); + } +} + +fn png_with_dimensions(width: u32, height: u32) -> Vec { + let mut encoded = encoded_image(ImageFormat::Png, 1, 1, None); + assert_eq!(&encoded[12..16], b"IHDR"); + encoded[16..20].copy_from_slice(&width.to_be_bytes()); + encoded[20..24].copy_from_slice(&height.to_be_bytes()); + let mut crc = u32::MAX; + for byte in &encoded[12..29] { + crc ^= u32::from(*byte); + for _ in 0..8 { + crc = if crc & 1 == 1 { + (crc >> 1) ^ 0xedb8_8320 + } else { + crc >> 1 + }; + } + } + encoded[29..33].copy_from_slice(&(!crc).to_be_bytes()); + encoded +} + +#[test] +fn invalid_and_oversized_dimensions_keep_validation_errors() { + for (width, height) in [(0, 1), (1, 0), (32_769, 1), (1, 32_769), (4_097, 4_096)] { + let error = assert_matches_original(&png_with_dimensions(width, height)).unwrap_err(); + if width > 32_768 || height > 32_768 { + assert!(error.starts_with("Cannot decode image (maximum 32,768 pixels per side): ")); + } else if u64::from(width) * u64::from(height) > 16_777_216 { + assert!(error.starts_with("Images must have at most 16,777,216 pixels")); + } + } +}