Refactor parse_sbit to prevent error triggers

image-rs · Nov 15, 2024 · 194605e · 194605e
1 parent 69c4993
commit 194605e
Showing 1 changed file with 58 additions and 54 deletions.
diff --git a/src/decoder/stream.rs b/src/decoder/stream.rs
@@ -1108,70 +1108,74 @@ impl StreamingDecoder {
     }
 
     fn parse_sbit(&mut self) -> Result<Decoded, DecodingError> {
-        let info = self.info.as_mut().unwrap();
-        if info.palette.is_some() {
-            return Err(DecodingError::Format(
-                FormatErrorInner::AfterPlte { kind: chunk::sBIT }.into(),
-            ));
-        }
-
-        if self.have_idat {
-            return Err(DecodingError::Format(
-                FormatErrorInner::AfterIdat { kind: chunk::sBIT }.into(),
-            ));
-        }
-
-        if info.sbit.is_some() {
-            return Err(DecodingError::Format(
-                FormatErrorInner::DuplicateChunk { kind: chunk::sBIT }.into(),
-            ));
-        }
+        let mut parse = || {
+            let info = self.info.as_mut().unwrap();
+            if info.palette.is_some() {
+                return Err(DecodingError::Format(
+                    FormatErrorInner::AfterPlte { kind: chunk::sBIT }.into(),
+                ));
+            }
 
-        let (color_type, bit_depth) = { (info.color_type, info.bit_depth) };
-        // The sample depth for color type 3 is fixed at eight bits.
-        let sample_depth = if color_type == ColorType::Indexed {
-            BitDepth::Eight
-        } else {
-            bit_depth
-        };
-        self.limits
-            .reserve_bytes(self.current_chunk.raw_bytes.len())?;
-        let vec = self.current_chunk.raw_bytes.clone();
-        let len = vec.len();
+            if self.have_idat {
+                return Err(DecodingError::Format(
+                    FormatErrorInner::AfterIdat { kind: chunk::sBIT }.into(),
+                ));
+            }
 
-        // expected lenth of the chunk
-        let expected = match color_type {
-            ColorType::Grayscale => 1,
-            ColorType::Rgb | ColorType::Indexed => 3,
-            ColorType::GrayscaleAlpha => 2,
-            ColorType::Rgba => 4,
-        };
+            if info.sbit.is_some() {
+                return Err(DecodingError::Format(
+                    FormatErrorInner::DuplicateChunk { kind: chunk::sBIT }.into(),
+                ));
+            }
 
-        // Check if the sbit chunk size is valid.
-        if expected != len {
-            return Err(DecodingError::Format(
-                FormatErrorInner::InvalidSbitChunkSize {
-                    color_type,
-                    expected,
-                    len,
-                }
-                .into(),
-            ));
-        }
+            let (color_type, bit_depth) = { (info.color_type, info.bit_depth) };
+            // The sample depth for color type 3 is fixed at eight bits.
+            let sample_depth = if color_type == ColorType::Indexed {
+                BitDepth::Eight
+            } else {
+                bit_depth
+            };
+            self.limits
+                .reserve_bytes(self.current_chunk.raw_bytes.len())?;
+            let vec = self.current_chunk.raw_bytes.clone();
+            let len = vec.len();
+
+            // expected lenth of the chunk
+            let expected = match color_type {
+                ColorType::Grayscale => 1,
+                ColorType::Rgb | ColorType::Indexed => 3,
+                ColorType::GrayscaleAlpha => 2,
+                ColorType::Rgba => 4,
+            };
 
-        for sbit in &vec {
-            if *sbit < 1 || *sbit > sample_depth as u8 {
+            // Check if the sbit chunk size is valid.
+            if expected != len {
                 return Err(DecodingError::Format(
-                    FormatErrorInner::InvalidSbit {
-                        sample_depth,
-                        sbit: *sbit,
+                    FormatErrorInner::InvalidSbitChunkSize {
+                        color_type,
+                        expected,
+                        len,
                     }
                     .into(),
                 ));
             }
-        }
 
-        info.sbit = Some(Cow::Owned(vec));
+            for sbit in &vec {
+                if *sbit < 1 || *sbit > sample_depth as u8 {
+                    return Err(DecodingError::Format(
+                        FormatErrorInner::InvalidSbit {
+                            sample_depth,
+                            sbit: *sbit,
+                        }
+                        .into(),
+                    ));
+                }
+            }
+            info.sbit = Some(Cow::Owned(vec));
+            Ok(Decoded::Nothing)
+        };
+
+        parse().ok();
         Ok(Decoded::Nothing)
     }