Reject truncated JPEG image buffers

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<Array3<u8>> {
+        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<Array3<u8>> {
         match img.color().channel_count() {
             1 => {

src/tests.rs

@@ -348,6 +348,22 @@ mod cv_compat_tests {
     use ndarray::Array3;
     use std::io::Cursor;
 
+    fn patterned_jpeg() -> Vec<u8> {
+        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());

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)]