From b1fdc33a229836b0547490dcec23e8c7707ade13 Mon Sep 17 00:00:00 2001
From: Nico Weber <thakis@chromium.org>
Date: Sun, 3 Mar 2024 20:51:38 -0500
Subject: [PATCH] LibGfx/JBIG2: Decode all segment headers

With `#define JBIG2_DEBUG 1` at the top of the file:

    % Build/lagom/bin/image --no-output \
       .../JBIG2_ConformanceData-A20180829/F01_200_TT10.jb2
    JBIG2LoadingContext: Organization: 0 (Sequential)
    JBIG2LoadingContext: Number of pages: 1
    Segment number: 0
    Segment type: 48
    Referred to segment count: 0
    Segment page association: 1
    Segment data length: 19
    Segment number: 1
    Segment type: 39
    Referred to segment count: 0
    Segment page association: 1
    Segment data length: 12666
    Segment number: 2
    Segment type: 49
    Referred to segment count: 0
    Segment page association: 1
    Segment data length: 0
    Runtime error: JBIG2ImageDecoderPlugin: Draw the rest of the owl
---
 .../LibGfx/ImageFormats/JBIG2Loader.cpp       | 25 +++++++++++++++----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.cpp b/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.cpp
index 30c6592ca5..b88c0d2e86 100644
--- a/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.cpp
+++ b/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.cpp
@@ -106,11 +106,8 @@ static ErrorOr<void> decode_jbig2_header(JBIG2LoadingContext& context)
     return {};
 }
 
-static ErrorOr<SegmentHeader> decode_segment_header(JBIG2LoadingContext& context)
+static ErrorOr<SegmentHeader> decode_segment_header(SeekableStream& stream)
 {
-    ReadonlyBytes data = context.data.slice(sizeof(id_string) + sizeof(u8) + (context.number_of_pages.has_value() ? sizeof(u32) : 0));
-    FixedMemoryStream stream(data);
-
     // 7.2.2 Segment number
     u32 segment_number = TRY(stream.read_value<BigEndian<u32>>());
     dbgln_if(JBIG2_DEBUG, "Segment number: {}", segment_number);
@@ -181,6 +178,24 @@ static ErrorOr<SegmentHeader> decode_segment_header(JBIG2LoadingContext& context
     return SegmentHeader { segment_number, type, move(referred_to_segment_numbers), segment_page_association, opt_data_length };
 }
 
+static ErrorOr<void> decode_segment_headers(JBIG2LoadingContext& context)
+{
+    FixedMemoryStream stream(context.data.slice(sizeof(id_string) + sizeof(u8) + (context.number_of_pages.has_value() ? sizeof(u32) : 0)));
+    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));
+        }
+
+        // Required per spec for files with RandomAccess organization.
+        if (segment_header.type == SegmentType::EndOfFile)
+            break;
+    }
+    return {};
+}
+
 JBIG2ImageDecoderPlugin::JBIG2ImageDecoderPlugin(ReadonlyBytes data)
 {
     m_context = make<JBIG2LoadingContext>();
@@ -201,7 +216,7 @@ ErrorOr<NonnullOwnPtr<ImageDecoderPlugin>> JBIG2ImageDecoderPlugin::create(Reado
 {
     auto plugin = TRY(adopt_nonnull_own_or_enomem(new (nothrow) JBIG2ImageDecoderPlugin(data)));
     TRY(decode_jbig2_header(*plugin->m_context));
-    TRY(decode_segment_header(*plugin->m_context));
+    TRY(decode_segment_headers(*plugin->m_context));
     return plugin;
 }