From 738f707beeaf1f2fa91afa358ed7dac1e1e84239 Mon Sep 17 00:00:00 2001
From: caleb <etemesicaleb@gmail.com>
Date: Tue, 2 Apr 2024 21:18:02 +0300
Subject: [PATCH] crates/jpeg: Add check to not read anything after seeing an
 EOI marker

---
 crates/zune-jpeg/src/bitstream.rs | 13 +++++++++----
 crates/zune-jpeg/src/mcu.rs       |  7 +++++++
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/crates/zune-jpeg/src/bitstream.rs b/crates/zune-jpeg/src/bitstream.rs
index f36e4fce..9f3032a7 100644
--- a/crates/zune-jpeg/src/bitstream.rs
+++ b/crates/zune-jpeg/src/bitstream.rs
@@ -122,7 +122,10 @@ pub(crate) struct BitStream {
     spec_start:          u8,
     spec_end:            u8,
     pub eob_run:         i32,
-    pub overread_by:     usize
+    pub overread_by:     usize,
+    /// True if we have seen end of image marker.
+    /// Don't read anything after that.
+    pub seen_eoi:        bool
 }
 
 impl BitStream {
@@ -138,7 +141,8 @@ impl BitStream {
             spec_start:      0,
             spec_end:        0,
             eob_run:         0,
-            overread_by:     0
+            overread_by:     0,
+            seen_eoi:        false
         }
     }
 
@@ -155,7 +159,8 @@ impl BitStream {
             spec_start:      spec_start,
             spec_end:        spec_end,
             eob_run:         0,
-            overread_by:     0
+            overread_by:     0,
+            seen_eoi:        false
         }
     }
 
@@ -223,7 +228,7 @@ impl BitStream {
 
         // 32 bits is enough for a decode(16 bits) and receive_extend(max 16 bits)
         // If we have less than 32 bits we refill
-        if self.bits_left < 32 && self.marker.is_none() {
+        if self.bits_left < 32 && self.marker.is_none() && !self.seen_eoi {
             // we optimize for the case where we don't have 255 in the stream and have 4 bytes left
             // as it is the common case
             //
diff --git a/crates/zune-jpeg/src/mcu.rs b/crates/zune-jpeg/src/mcu.rs
index f391d29b..fe633851 100644
--- a/crates/zune-jpeg/src/mcu.rs
+++ b/crates/zune-jpeg/src/mcu.rs
@@ -281,6 +281,13 @@ impl<T: ZByteReaderTrait> JpegDecoder<T> {
                     // acknowledge and ignore EOI marker.
                     stream.marker.take();
                     trace!("Found EOI marker");
+                    // Google Introduced the Ultra-HD image format which is basically
+                    // stitching two images into one container.
+                    // They basically separate two images via a EOI and SOI marker
+                    // so let's just ensure if we ever see EOI, we never read past that
+                    // ever.
+                    // https://github.com/google/libultrahdr
+                    stream.seen_eoi = true;
                 } else if let Marker::RST(_) = m {
                     if self.todo == 0 {
                         self.handle_rst(stream)?;