From c807959d755589d3d579f37556d33e01d0ac3c82 Mon Sep 17 00:00:00 2001
From: Kirill Gladkikh <boombarah@mail.ru>
Date: Sun, 16 Jun 2024 18:16:01 +0500
Subject: [PATCH] patterns/jpeg: Support EOF in the middle of JPEG (#240)

* Support EOF in the middle of file

* Changed to eoi marker searching
---
 patterns/jpeg.hexpat | 56 +++++++++++++++++++++++++++-----------------
 1 file changed, 35 insertions(+), 21 deletions(-)

diff --git a/patterns/jpeg.hexpat b/patterns/jpeg.hexpat
index 8b42db30..43f8e931 100644
--- a/patterns/jpeg.hexpat
+++ b/patterns/jpeg.hexpat
@@ -41,7 +41,8 @@ enum Marker : u8 {
     APP13   = 0xED,
     APP14   = 0xEE,
     APP15   = 0xEF,
-    COM     = 0xFE
+    COM     = 0xFE,
+    UNKNOWN = 0X00
 };
 
 enum DensityUnit : u8 {
@@ -103,6 +104,15 @@ fn sof0_component_read(SOF0Component c) {
     return std::format("({}, {}, {})", c.componentId, c.samplingFactors, c.quantizationTableId);
 };
 
+fn get_eoi_marker_position() {
+    u32 pos = std::mem::find_sequence_in_range(0, $, std::mem::size(), 0xFF);
+    while (std::mem::read_unsigned(pos + 1, 1) != 0xD9) {
+        pos = std::mem::find_sequence_in_range(0, pos + 1, std::mem::size() + 1, 0xFF);
+    }
+    return pos;
+    
+};
+
 struct SOF0 {
     u8 bitsPerSample;
     u16 imageHeight, imageWidth;
@@ -121,30 +131,34 @@ struct SOS {
     u8 startSpectralSelection;
     u8 endSpectral;
     u8 apprBitPos;
-    
-    u8 image_data[std::mem::size() - $ - 2] [[sealed]];
+    u8 image_data[get_eoi_marker_position() - $] [[sealed]];
 };
 
 struct Segment {
-    type::Magic<"\xff"> magic;
-    Marker marker;
-    
-    if (marker == Marker::SOI || marker == Marker::EOI) {
-        
+    if (std::mem::read_unsigned($, 1) != 0xFF) {
+        Marker marker = Marker::UNKNOWN;
+        u8 data[get_eoi_marker_position() - $];
     } else {
-        u16 length;
-        if (marker == Marker::APP0) {
-            APP0 data;
-        } else if (marker == Marker::APP14) {
-            APP14 data;
-        } else if (marker == Marker::COM) {
-            char data[length - sizeof(length)];
-        } else if (marker == Marker::SOF0) {
-            SOF0 data;
-        } else if (marker == Marker::SOS) {
-            SOS data;
+        type::Magic<"\xff"> magic;
+        Marker marker;
+        
+        if (marker == Marker::SOI || marker == Marker::EOI) {
+            
         } else {
-            u8 data[length - sizeof(length)] [[sealed]];
+            u16 length;
+            if (marker == Marker::APP0) {
+                APP0 data;
+            } else if (marker == Marker::APP14) {
+                APP14 data;
+            } else if (marker == Marker::COM) {
+                char data[length - sizeof(length)];
+            } else if (marker == Marker::SOF0) {
+                SOF0 data;
+            } else if (marker == Marker::SOS) {
+                SOS data;
+            } else {
+                u8 data[length - sizeof(length)] [[sealed]];
+            }
         }
     }
 } [[format_read("segment_read")]];
@@ -153,4 +167,4 @@ fn segment_read(Segment s) {
     return std::format("{}", s.marker);
 };
 
-Segment segments[while(!std::mem::eof())] @ 0x00 [[hex::visualize("image", this)]];
\ No newline at end of file
+Segment segments[while(!std::mem::eof())] @ 0x00 [[hex::visualize("image", this)]];