From 379ef456889e7f803fead752a2d8572d00c7b25b Mon Sep 17 00:00:00 2001
From: Nico Weber <thakis@chromium.org>
Date: Sun, 3 Mar 2024 21:14:07 -0500
Subject: [PATCH] LibGfx/JBIG2: Store location of segment data bodies

They're in different places for Sequential/Embedded (right after
the header) and RandomAccess (which has all headers first, followed
by all data bits next).

We don't do anything with the data yet, but now everything's in
place to actually process segment data.
---
 .../LibGfx/ImageFormats/JBIG2Loader.cpp       | 44 ++++++++++++++++---
 1 file changed, 39 insertions(+), 5 deletions(-)

diff --git a/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.cpp b/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.cpp
index 4a3fcbfc6c..135c90d076 100644
--- a/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.cpp
+++ b/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.cpp
@@ -62,6 +62,11 @@ struct SegmentHeader {
     Optional<u32> data_length;
 };
 
+struct SegmentData {
+    SegmentHeader header;
+    ReadonlyBytes data;
+};
+
 struct JBIG2LoadingContext {
     enum class State {
         NotDecoded = 0,
@@ -74,6 +79,8 @@ struct JBIG2LoadingContext {
     IntSize size;
 
     Optional<u32> number_of_pages;
+
+    Vector<SegmentData> segments;
 };
 
 static ErrorOr<void> decode_jbig2_header(JBIG2LoadingContext& context)
@@ -184,18 +191,45 @@ static ErrorOr<void> decode_segment_headers(JBIG2LoadingContext& context)
     if (context.organization != Organization::Embedded)
         data = data.slice(sizeof(id_string) + sizeof(u8) + (context.number_of_pages.has_value() ? sizeof(u32) : 0));
     FixedMemoryStream stream(data);
+
+    Vector<ReadonlyBytes> segment_datas;
+    auto store_and_skip_segment_data = [&](SegmentHeader const& segment_header) -> ErrorOr<void> {
+        if (!segment_header.data_length.has_value())
+            return Error::from_string_literal("JBIG2ImageDecoderPlugin: Can't handle segment without data length yet");
+
+        size_t start_offset = TRY(stream.tell());
+        if (start_offset + segment_header.data_length.value() > data.size())
+            return Error::from_string_literal("JBIG2ImageDecoderPlugin: Segment data length exceeds file size");
+        ReadonlyBytes segment_data = data.slice(start_offset, segment_header.data_length.value());
+        segment_datas.append(segment_data);
+
+        TRY(stream.seek(segment_header.data_length.value(), SeekMode::FromCurrentPosition));
+        return {};
+    };
+
+    Vector<SegmentHeader> segment_headers;
     while (!stream.is_eof()) {
         auto segment_header = TRY(decode_segment_header(stream));
-        if (context.organization != Organization::RandomAccess) {
-            if (!segment_header.data_length.has_value())
-                return Error::from_string_literal("JBIG2ImageDecoderPlugin: Can't handle non-random-access organization segment without data length yet");
-            TRY(stream.seek(segment_header.data_length.value(), SeekMode::FromCurrentPosition));
-        }
+        segment_headers.append(segment_header);
+
+        if (context.organization != Organization::RandomAccess)
+            TRY(store_and_skip_segment_data(segment_header));
 
         // Required per spec for files with RandomAccess organization.
         if (segment_header.type == SegmentType::EndOfFile)
             break;
     }
+
+    if (context.organization == Organization::RandomAccess) {
+        for (auto const& segment_header : segment_headers)
+            TRY(store_and_skip_segment_data(segment_header));
+    }
+
+    if (segment_headers.size() != segment_datas.size())
+        return Error::from_string_literal("JBIG2ImageDecoderPlugin: Segment headers and segment datas have different sizes");
+    for (size_t i = 0; i < segment_headers.size(); ++i)
+        context.segments.append({ segment_headers[i], segment_datas[i] });
+
     return {};
 }