From 193ac5baeef35916677e03f4cfab31133bd2f7b8 Mon Sep 17 00:00:00 2001
From: Min Deng <minotes.net@gmail.com>
Date: Mon, 18 Nov 2024 09:47:07 +0800
Subject: [PATCH] fix(exif): #21 assertion failed: data.len() >= 6 when
 checking broken exif file

---
 src/exif/exif_exif.rs | 10 +++-------
 src/heif.rs           |  2 +-
 src/jpeg.rs           | 28 +++++++++++++++++++++++-----
 3 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/src/exif/exif_exif.rs b/src/exif/exif_exif.rs
index 69c6ddb..6f80362 100644
--- 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
index f1103d3..016980f 100644
--- 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
index 9b86bf3..30f07a0 100644
--- 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