diff --git a/src/cv_compat.rs b/src/cv_compat.rs index 7aac388..16ddcf2 100644 --- a/src/cv_compat.rs +++ b/src/cv_compat.rs @@ -20,6 +20,10 @@ pub mod imdecode { /// Decode image from byte buffer (equivalent to cv2.imdecode) pub fn imdecode(buf: &[u8], flags: ImreadFlags) -> Result> { + if buf.starts_with(&[0xff, 0xd8]) { + validate_jpeg_complete(buf)?; + } + let img = image::load_from_memory(buf) .map_err(|e| anyhow::anyhow!("Failed to decode image: {}", e))?; @@ -38,6 +42,90 @@ pub mod imdecode { } } + fn validate_jpeg_complete(buf: &[u8]) -> Result<()> { + let mut offset = 2; + let mut in_entropy_data = false; + let mut saw_scan = false; + + while offset < buf.len() { + let marker = if in_entropy_data { + loop { + while offset < buf.len() && buf[offset] != 0xff { + offset += 1; + } + if offset == buf.len() { + anyhow::bail!("JPEG image is truncated: missing end-of-image marker"); + } + + while offset < buf.len() && buf[offset] == 0xff { + offset += 1; + } + if offset == buf.len() { + anyhow::bail!("JPEG image is truncated: incomplete marker"); + } + + let marker = buf[offset]; + offset += 1; + match marker { + 0x00 | 0xd0..=0xd7 => continue, + _ => { + in_entropy_data = false; + break marker; + } + } + } + } else { + if buf[offset] != 0xff { + anyhow::bail!("JPEG image is malformed: expected marker"); + } + while offset < buf.len() && buf[offset] == 0xff { + offset += 1; + } + if offset == buf.len() { + anyhow::bail!("JPEG image is truncated: incomplete marker"); + } + + let marker = buf[offset]; + offset += 1; + marker + }; + + match marker { + 0xd9 if saw_scan => return Ok(()), + 0xd9 => anyhow::bail!("JPEG image is malformed: end marker precedes scan data"), + 0x01 | 0xd0..=0xd8 => continue, + _ => { + let length_end = offset.checked_add(2).ok_or_else(|| { + anyhow::anyhow!("JPEG image is truncated: invalid segment length") + })?; + if length_end > buf.len() { + anyhow::bail!("JPEG image is truncated: incomplete segment length"); + } + + let segment_length = + usize::from(u16::from_be_bytes([buf[offset], buf[offset + 1]])); + if segment_length < 2 { + anyhow::bail!("JPEG image is malformed: invalid segment length"); + } + + offset = offset.checked_add(segment_length).ok_or_else(|| { + anyhow::anyhow!("JPEG image is truncated: invalid segment length") + })?; + if offset > buf.len() { + anyhow::bail!("JPEG image is truncated: incomplete segment"); + } + + if marker == 0xda { + saw_scan = true; + in_entropy_data = true; + } + } + } + } + + anyhow::bail!("JPEG image is truncated: missing end-of-image marker") + } + fn dynamic_image_to_ndarray(img: DynamicImage) -> Result> { match img.color().channel_count() { 1 => { diff --git a/src/tests.rs b/src/tests.rs index e908b77..8aa51d3 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -348,6 +348,22 @@ mod cv_compat_tests { use ndarray::Array3; use std::io::Cursor; + fn patterned_jpeg() -> Vec { + let mut encoded = Cursor::new(Vec::new()); + let rgb = ImageBuffer::from_fn(512, 512, |x, y| { + Rgb([ + (x.wrapping_mul(31) ^ y.wrapping_mul(17)) as u8, + (x.wrapping_mul(13) ^ y.wrapping_mul(29)) as u8, + (x.wrapping_mul(7) ^ y.wrapping_mul(37)) as u8, + ]) + }); + + DynamicImage::ImageRgb8(rgb) + .write_to(&mut encoded, ImageFormat::Jpeg) + .unwrap(); + encoded.into_inner() + } + #[test] fn test_imdecode_detects_jpeg_without_png_hint() { let mut encoded = Cursor::new(Vec::new()); @@ -361,6 +377,39 @@ mod cv_compat_tests { assert_eq!(decoded.dim(), (2, 3, 3)); } + #[test] + fn test_imdecode_rejects_truncated_jpeg() { + let mut truncated = patterned_jpeg(); + truncated.truncate(truncated.len() / 20); + + let error = imdecode(&truncated, ImreadFlags::ImreadColor).unwrap_err(); + assert!(error.to_string().contains("truncated")); + } + + #[test] + fn test_imdecode_does_not_treat_metadata_bytes_as_jpeg_end() { + let mut truncated = patterned_jpeg(); + truncated.splice(2..2, [0xff, 0xe1, 0x00, 0x06, 0xff, 0xd9, 0x12, 0x34]); + truncated.truncate(truncated.len() / 20); + + let error = imdecode(&truncated, ImreadFlags::ImreadColor).unwrap_err(); + assert!(error.to_string().contains("truncated")); + } + + #[test] + fn test_imdecode_accepts_jpeg_with_trailing_bytes() { + let mut encoded = Cursor::new(Vec::new()); + let rgb = ImageBuffer::from_fn(3, 2, |x, y| Rgb([(x * 40) as u8, (y * 80) as u8, 200])); + + DynamicImage::ImageRgb8(rgb) + .write_to(&mut encoded, ImageFormat::Jpeg) + .unwrap(); + encoded.get_mut().extend_from_slice(b"trailing data"); + + let decoded = imdecode(encoded.get_ref(), ImreadFlags::ImreadColor).unwrap(); + assert_eq!(decoded.dim(), (2, 3, 3)); + } + #[test] fn test_imdecode_detects_webp_without_png_hint() { let mut encoded = Cursor::new(Vec::new()); diff --git a/tests/test_python_bindings.py b/tests/test_python_bindings.py index 5c2531d..cf72429 100644 --- a/tests/test_python_bindings.py +++ b/tests/test_python_bindings.py @@ -47,6 +47,24 @@ def test_imdecode_owned_buffer(self, buffer_type): np.testing.assert_array_equal(decoded, image) + def test_imdecode_rejects_truncated_jpeg(self): + """Reject JPEGs whose missing scan data would otherwise be concealed.""" + y, x = np.indices((512, 512), dtype=np.uint32) + image = np.stack( + ( + (x * 31) ^ (y * 17), + (x * 13) ^ (y * 29), + (x * 7) ^ (y * 37), + ), + axis=-1, + ).astype(np.uint8) + encoded = io.BytesIO() + Image.fromarray(image).save(encoded, format="JPEG") + truncated = encoded.getvalue()[: len(encoded.getvalue()) // 20] + + with pytest.raises(RuntimeError, match="truncated"): + tsr.imdecode_py(truncated, 1) + def test_batch_crop_images(self, sample_images): """Test batch cropping with custom coordinates.""" crop_boxes = [(10, 10, 50, 50), (20, 20, 40, 40), (5, 5, 60, 60)]