From f98195a087f482dfc6e096f235135573fcd0a428 Mon Sep 17 00:00:00 2001
From: Codex <codex@openai.com>
Date: Sat, 1 Aug 2026 15:15:36 +0300
Subject: [PATCH] pdf-image: validate decoded samples during decode

Move decoded sample validation into DecodedSamples and apply it before
returning from every decoding path.

Keep ImageXObject focused on pixel assembly while preserving component
count and truncated data errors.

Co-authored-by: Codex <codex@openai.com>
---
 crates/pdf-image/src/decoded_samples.rs | 41 +++++++++++++++++++------
 crates/pdf-image/src/image_xobject.rs   | 22 -------------
 2 files changed, 31 insertions(+), 32 deletions(-)

diff --git a/crates/pdf-image/src/decoded_samples.rs b/crates/pdf-image/src/decoded_samples.rs
index 2b22c57..5ab7363 100644
--- a/crates/pdf-image/src/decoded_samples.rs
+++ b/crates/pdf-image/src/decoded_samples.rs
@@ -24,20 +24,41 @@ impl DecodedSamples {
         objects: &dyn ObjectResolver,
         metadata: &ImageMetadata,
     ) -> Result<Self, PdfImageError> {
-        if let Some(decoded_samples) = Self::decode_preconverted_jpx(raw_data, metadata) {
-            return Ok(decoded_samples);
-        }
+        let decoded_samples = if let Some(decoded_samples) =
+            Self::decode_preconverted_jpx(raw_data, metadata)
+        {
+            decoded_samples
+        } else if let Some(decoded_samples) = Self::decode_preconverted_dct(raw_data, metadata) {
+            decoded_samples
+        } else {
+            match metadata.color_space.as_ref() {
+                Some(ColorSpace::Indexed(indexed)) => {
+                    Self::decode_indexed(dictionary, raw_data, objects, metadata, indexed)
+                }
+                _ => Self::decode_direct(dictionary, raw_data, objects, metadata),
+            }?
+        };
 
-        if let Some(decoded_samples) = Self::decode_preconverted_dct(raw_data, metadata) {
-            return Ok(decoded_samples);
+        decoded_samples.validate(metadata)?;
+        Ok(decoded_samples)
+    }
+
+    /// Ensures the decoded component stream is large enough for the declared dimensions.
+    fn validate(&self, metadata: &ImageMetadata) -> Result<(), PdfImageError> {
+        if self.num_color_components == 0 {
+            return Err(PdfImageError::InvalidColorComponentCount);
         }
 
-        match metadata.color_space.as_ref() {
-            Some(ColorSpace::Indexed(indexed)) => {
-                Self::decode_indexed(dictionary, raw_data, objects, metadata, indexed)
-            }
-            _ => Self::decode_direct(dictionary, raw_data, objects, metadata),
+        let num_pixels = metadata.width.saturating_mul(metadata.height);
+        let expected_bytes = num_pixels.saturating_mul(self.num_color_components);
+        if self.image_data.len() < expected_bytes {
+            return Err(PdfImageError::TruncatedImageData {
+                expected_bytes,
+                actual_bytes: self.image_data.len(),
+            });
         }
+
+        Ok(())
     }
 
     /// Uses DCT decoder output as display samples when the JPEG decoder already converted color.
diff --git a/crates/pdf-image/src/image_xobject.rs b/crates/pdf-image/src/image_xobject.rs
index 921edcf..cf5fc51 100644
--- a/crates/pdf-image/src/image_xobject.rs
+++ b/crates/pdf-image/src/image_xobject.rs
@@ -95,7 +95,6 @@ impl ImageXObject {
         metadata: &ImageMetadata,
     ) -> Result<Self, PdfImageError> {
         let decoded_samples = DecodedSamples::decode(dictionary, raw_data, objects, metadata)?;
-        Self::validate_decoded_samples(metadata, &decoded_samples)?;
         let (data, pixel_format) = Self::assemble_pixel_data(metadata, &decoded_samples, soft_mask);
 
         Ok(Self {
@@ -108,27 +107,6 @@ impl ImageXObject {
         })
     }
 
-    /// Ensures the decoded component stream is large enough for the declared dimensions.
-    fn validate_decoded_samples(
-        metadata: &ImageMetadata,
-        decoded_samples: &DecodedSamples,
-    ) -> Result<(), PdfImageError> {
-        if decoded_samples.num_color_components == 0 {
-            return Err(PdfImageError::InvalidColorComponentCount);
-        }
-
-        let num_pixels = metadata.width.saturating_mul(metadata.height);
-        let expected_bytes = num_pixels.saturating_mul(decoded_samples.num_color_components);
-        if decoded_samples.image_data.len() < expected_bytes {
-            return Err(PdfImageError::TruncatedImageData {
-                expected_bytes,
-                actual_bytes: decoded_samples.image_data.len(),
-            });
-        }
-
-        Ok(())
-    }
-
     /// Builds the final pixel buffer and pixel format after optional soft-mask application.
     fn assemble_pixel_data(
         metadata: &ImageMetadata,