fix(exif): #21 assertion failed: data.len() >= 6 when checking broken…

… exif file
mindeng · Nov 18, 2024 · 193ac5b · 193ac5b
1 parent 02fc32c
commit 193ac5b
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 13 deletions.
diff --git a/src/exif/exif_exif.rs b/src/exif/exif_exif.rs
@@ -1,5 +1,5 @@
 use nom::{
-    branch::alt, bytes::complete::tag, combinator, number::Endianness, sequence, IResult, Needed,
+    branch::alt, bytes::streaming::tag, combinator, number::Endianness, sequence, IResult, Needed,
 };
 
 use crate::{EntryValue, ExifIter, ExifTag, GPSInfo, ParsedExifEntry};
@@ -229,12 +229,8 @@ impl TiffHeader {
     }
 }
 
-/// data.len() MUST >= 6
-pub(crate) fn check_exif_header(data: &[u8]) -> bool {
-    use nom::bytes::complete;
-    assert!(data.len() >= 6);
-
-    complete::tag::<_, _, nom::error::Error<_>>(EXIF_IDENT)(data).is_ok()
+pub(crate) fn check_exif_header(data: &[u8]) -> Result<bool, nom::Err<nom::error::Error<&[u8]>>> {
+    tag::<_, _, nom::error::Error<_>>(EXIF_IDENT)(data).map(|_| true)
 }
 
 pub(crate) fn check_exif_header2(i: &[u8]) -> IResult<&[u8], ()> {

diff --git a/src/heif.rs b/src/heif.rs
@@ -88,7 +88,7 @@ pub(crate) fn extract_exif_with_meta<'a>(
 
     if let Some(data) = data {
         let (remain, _) = be_u32(data)?;
-        if check_exif_header(remain) {
+        if check_exif_header(remain)? {
             Ok((out_remain, Some(&remain[6..]))) // Safe-slice
         } else {
             Ok((out_remain, None))

diff --git a/src/jpeg.rs b/src/jpeg.rs
@@ -74,10 +74,28 @@ impl<'a> Segment<'a> {
 }
 
 fn find_exif_segment(input: &[u8]) -> IResult<&[u8], Option<Segment<'_>>> {
-    let (remain, segment) = travel_until(input, |s| {
-        (s.marker_code == MarkerCode::APP1.code() && check_exif_header(s.payload))
-            || s.marker_code == MarkerCode::Sos.code() // searching stop at SOS
-    })?;
+    let mut remain = input;
+
+    let (remain, segment) = loop {
+        let (rem, (_, code)) = tuple((streaming::tag([0xFF]), number::streaming::u8))(remain)?;
+        let (rem, segment) = parse_segment(code, rem)?;
+        // Sanity check
+        assert!(rem.len() < remain.len());
+        remain = rem;
+        tracing::debug!(
+            marker = format!("0x{:04x}", segment.marker_code),
+            size = format!("0x{:04x}", segment.payload.len()),
+            "got segment"
+        );
+
+        let s = &segment;
+        if (s.marker_code == MarkerCode::APP1.code() && check_exif_header(s.payload)?)
+            || s.marker_code == MarkerCode::Sos.code()
+        // searching stop at SOS
+        {
+            break (remain, segment);
+        }
+    };
 
     if segment.marker_code != MarkerCode::Sos.code() {
         Ok((remain, Some(segment)))
@@ -154,7 +172,7 @@ pub(crate) fn check_jpeg_exif<'a>(input: &'a [u8]) -> IResult<&'a [u8], bool> {
             }
 
             let (rem, header) = streaming::take(EXIF_HEADER_SIZE)(rem)?;
-            let b = check_exif_header(header);
+            let b = check_exif_header(header)?;
             return Ok((rem, b));
         } else {
             // skip to next segment