From 5cefcad2fe0b428df9b7aeec882d8baa91909d36 Mon Sep 17 00:00:00 2001
From: Nico Weber <thakis@chromium.org>
Date: Sun, 3 Mar 2024 20:05:21 -0500
Subject: [PATCH] LibGfx/JBIG2: Decode the file header

Running `image` with `#define JBIG2_DEBUG 1` now prints number of pages.
---
 .../LibGfx/ImageFormats/JBIG2Loader.cpp       | 90 ++++++++++++++++++-
 .../LibGfx/ImageFormats/JBIG2Loader.h         |  9 ++
 2 files changed, 95 insertions(+), 4 deletions(-)

diff --git a/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.cpp b/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.cpp
index bad5393a61..9708a38ab4 100644
--- a/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.cpp
+++ b/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.cpp
@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: BSD-2-Clause
  */
 
+#include <AK/Debug.h>
 #include <LibGfx/ImageFormats/JBIG2Loader.h>
 
 // Spec: ITU-T_T_88__08_2018.pdf in the zip file here:
@@ -11,16 +12,97 @@
 
 namespace Gfx {
 
+// JBIG2 spec, Annex D, D.4.1 ID string
+static constexpr u8 id_string[] = { 0x97, 0x4A, 0x42, 0x32, 0x0D, 0x0A, 0x1A, 0x0A };
+
+// Annex D
+enum class Organization {
+    // D.1 Sequential organization
+    Sequential,
+
+    // D.2 Random-access organization
+    RandomAccess,
+
+    // D.3 Embedded organization
+    Embedded,
+};
+
+struct JBIG2LoadingContext {
+    enum class State {
+        NotDecoded = 0,
+        Error,
+    };
+    State state { State::NotDecoded };
+    ReadonlyBytes data;
+
+    Organization organization { Organization::Sequential };
+    IntSize size;
+
+    Optional<u32> number_of_pages;
+};
+
+static ErrorOr<void> decode_jbig2_header(JBIG2LoadingContext& context)
+{
+    if (!JBIG2ImageDecoderPlugin::sniff(context.data))
+        return Error::from_string_literal("JBIG2LoadingContext: Invalid JBIG2 header");
+
+    FixedMemoryStream stream(context.data.slice(sizeof(id_string)));
+
+    // D.4.2 File header flags
+    u8 header_flags = TRY(stream.read_value<u8>());
+    if (header_flags & 0b11110000)
+        return Error::from_string_literal("JBIG2LoadingContext: Invalid header flags");
+    context.organization = (header_flags & 1) ? Organization::Sequential : Organization::RandomAccess;
+    bool has_known_number_of_pages = (header_flags & 2) ? false : true;
+    bool uses_templates_with_12_AT_pixels = (header_flags & 4) ? true : false;
+    bool contains_colored_region_segments = (header_flags & 8) ? true : false;
+
+    // FIXME: Do something with these?
+    (void)uses_templates_with_12_AT_pixels;
+    (void)contains_colored_region_segments;
+
+    // D.4.3 Number of pages
+    if (has_known_number_of_pages) {
+        context.number_of_pages = TRY(stream.read_value<BigEndian<u32>>());
+        dbgln_if(JBIG2_DEBUG, "JBIG2LoadingContext: Number of pages: {}", context.number_of_pages.value());
+    }
+
+    return {};
+}
+
+JBIG2ImageDecoderPlugin::JBIG2ImageDecoderPlugin(ReadonlyBytes data)
+{
+    m_context = make<JBIG2LoadingContext>();
+    m_context->data = data;
+}
+
+IntSize JBIG2ImageDecoderPlugin::size()
+{
+    return m_context->size;
+}
+
 bool JBIG2ImageDecoderPlugin::sniff(ReadonlyBytes data)
 {
-    // JBIG2 spec, Annex D, D.4.1 ID string
-    u8 id_string[] = { 0x97, 0x4A, 0x42, 0x32, 0x0D, 0x0A, 0x1A, 0x0A };
     return data.starts_with(id_string);
 }
 
-ErrorOr<NonnullOwnPtr<ImageDecoderPlugin>> JBIG2ImageDecoderPlugin::create(ReadonlyBytes)
+ErrorOr<NonnullOwnPtr<ImageDecoderPlugin>> JBIG2ImageDecoderPlugin::create(ReadonlyBytes data)
 {
-    return Error::from_string_view("FIXME: Draw the rest of the owl"sv);
+    auto plugin = TRY(adopt_nonnull_own_or_enomem(new (nothrow) JBIG2ImageDecoderPlugin(data)));
+    TRY(decode_jbig2_header(*plugin->m_context));
+    return plugin;
+}
+
+ErrorOr<ImageFrameDescriptor> JBIG2ImageDecoderPlugin::frame(size_t index, Optional<IntSize>)
+{
+    // FIXME: Use this for multi-page JBIG2 files?
+    if (index != 0)
+        return Error::from_string_literal("JBIG2ImageDecoderPlugin: Invalid frame index");
+
+    if (m_context->state == JBIG2LoadingContext::State::Error)
+        return Error::from_string_literal("JBIG2ImageDecoderPlugin: Decoding failed");
+
+    return Error::from_string_literal("JBIG2ImageDecoderPlugin: Draw the rest of the owl");
 }
 
 }
diff --git a/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.h b/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.h
index 6faf4f3e59..94779c095a 100644
--- a/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.h
+++ b/Userland/Libraries/LibGfx/ImageFormats/JBIG2Loader.h
@@ -20,6 +20,15 @@ public:
     static ErrorOr<NonnullOwnPtr<ImageDecoderPlugin>> create(ReadonlyBytes);
 
     virtual ~JBIG2ImageDecoderPlugin() override = default;
+
+    virtual IntSize size() override;
+
+    virtual ErrorOr<ImageFrameDescriptor> frame(size_t index, Optional<IntSize> ideal_size = {}) override;
+
+private:
+    JBIG2ImageDecoderPlugin(ReadonlyBytes);
+
+    OwnPtr<JBIG2LoadingContext> m_context;
 };
 
 }