From de43bc99c873fe0c752f5e303c4a40f954b61912 Mon Sep 17 00:00:00 2001
From: Josh Ware <josh.ware@healius.com.au>
Date: Thu, 19 Jan 2023 11:37:14 +1100
Subject: [PATCH] Added support for jpeg2000 comments and PLT marker segments

---
 Tests/test_file_jpeg2k.py            | 44 ++++++++++++++++++++++++++++
 docs/handbook/image-file-formats.rst | 13 ++++++++
 src/PIL/Jpeg2KImagePlugin.py         |  4 +++
 src/libImaging/Jpeg2K.h              |  6 ++++
 src/libImaging/Jpeg2KEncode.c        | 17 +++++++++++
 5 files changed, 84 insertions(+)

diff --git a/Tests/test_file_jpeg2k.py b/Tests/test_file_jpeg2k.py
index 0229b2243..f52c33402 100644
--- a/Tests/test_file_jpeg2k.py
+++ b/Tests/test_file_jpeg2k.py
@@ -1,5 +1,6 @@
 import os
 import re
+import struct
 from io import BytesIO
 
 import pytest
@@ -371,3 +372,46 @@ def test_crashes(test_file):
                 im.load()
             except OSError:
                 pass
+
+
+def test_custom_comment():
+    output_stream = BytesIO()
+    unique_comment = "This is a unique comment, which should be found below"
+    test_card.save(output_stream, "JPEG2000", comment=unique_comment)
+    output_stream.seek(0)
+    data = output_stream.read()
+    # Lazy method to determine if the comment is in the image generated
+    assert(bytes(unique_comment, "utf-8") in data)
+
+
+def test_plt_marker():
+    # Search the start of the codesteam for the PLT box (id 0xFF58)
+    opj_version = re.search(r"(\d+\.\d+)\.\d+$", features.version_codec("jpg_2000"))
+    assert opj_version is not None
+
+    if float(opj_version[1]) >= 2.4:
+        out = BytesIO()
+        test_card.save(out, "JPEG2000", no_jp2=True, add_plt=True)
+        out.seek(0)
+        while True:
+            box_bytes = out.read(2)
+            if len(box_bytes) == 0:
+                # End of steam encounterd and no PLT or SOD
+                break
+            jp2_boxid = struct.unpack(">H", box_bytes)[0]
+
+            if jp2_boxid == 0xff4f:
+                # No length specifier for main header
+                continue
+            elif jp2_boxid == 0xff58:
+                # This is the PLT box we're looking for
+                return
+            elif jp2_boxid == 0xff93:
+                break
+                # SOD box encountered and no PLT, so it wasn't found
+
+            jp2_boxlength = struct.unpack(">H", out.read(2))[0]
+            out.seek(jp2_boxlength - 2, os.SEEK_CUR)
+
+        # The PLT box wasn't found
+        raise ValueError
diff --git a/docs/handbook/image-file-formats.rst b/docs/handbook/image-file-formats.rst
index a41ef7cf8..9128400ac 100644
--- a/docs/handbook/image-file-formats.rst
+++ b/docs/handbook/image-file-formats.rst
@@ -589,6 +589,19 @@ The :py:meth:`~PIL.Image.Image.save` method supports the following options:
 
     .. versionadded:: 9.1.0
 
+**comment**
+    Adds a custom comment to the file, replacing the default 
+    "Created by OpenJPEG version" comment.
+
+    .. versionadded:: 9.5.0
+
+**add_plt**
+    If ``True`` then include a PLT (packet length, tile-part header) marker 
+    segment in the produced file.
+    The default is to not include it.
+
+    .. versionadded:: 9.5.0
+
 .. note::
 
    To enable JPEG 2000 support, you need to build and install the OpenJPEG
diff --git a/src/PIL/Jpeg2KImagePlugin.py b/src/PIL/Jpeg2KImagePlugin.py
index 7457874c1..754010c7c 100644
--- a/src/PIL/Jpeg2KImagePlugin.py
+++ b/src/PIL/Jpeg2KImagePlugin.py
@@ -328,6 +328,8 @@ def _save(im, fp, filename):
     mct = info.get("mct", 0)
     signed = info.get("signed", False)
     fd = -1
+    comment = info.get("comment", None)
+    add_plt = info.get("add_plt", False)
 
     if hasattr(fp, "fileno"):
         try:
@@ -350,6 +352,8 @@ def _save(im, fp, filename):
         mct,
         signed,
         fd,
+        comment,
+        add_plt
     )
 
     ImageFile._save(im, fp, [("jpeg2k", (0, 0) + im.size, 0, kind)])
diff --git a/src/libImaging/Jpeg2K.h b/src/libImaging/Jpeg2K.h
index b28a0440a..65728be5d 100644
--- a/src/libImaging/Jpeg2K.h
+++ b/src/libImaging/Jpeg2K.h
@@ -97,6 +97,12 @@ typedef struct {
     /* PRIVATE CONTEXT (set by decoder) */
     const char *error_msg;
 
+    /* Custom comment */
+    char * comment;
+
+    /* Include PLT marker segment */
+    int add_plt;
+
 } JPEG2KENCODESTATE;
 
 /*
diff --git a/src/libImaging/Jpeg2KEncode.c b/src/libImaging/Jpeg2KEncode.c
index db1c5c0c9..bb280ae94 100644
--- a/src/libImaging/Jpeg2KEncode.c
+++ b/src/libImaging/Jpeg2KEncode.c
@@ -439,6 +439,10 @@ j2k_encode_entry(Imaging im, ImagingCodecState state) {
         params.tcp_mct = context->mct;
     }
 
+    if (context->comment) {
+        params.cp_comment = context->comment;
+    }
+
     params.prog_order = context->progression;
 
     params.cp_cinema = context->cinema_mode;
@@ -492,6 +496,14 @@ j2k_encode_entry(Imaging im, ImagingCodecState state) {
     opj_set_warning_handler(codec, j2k_warn, context);
     opj_setup_encoder(codec, &params, image);
 
+    /* Enabling PLT markers only supported in OpenJPEG 2.4.0 and up */
+#if ((OPJ_VERSION_MAJOR == 2 && OPJ_VERSION_MINOR >= 4) || OPJ_VERSION_MAJOR > 2)
+    if (context->add_plt) {
+        const char * plt_option[2] = {"PLT=YES", NULL};
+        opj_encoder_set_extra_options(codec, plt_option);
+    }
+#endif
+
     /* Start encoding */
     if (!opj_start_compress(codec, image, stream)) {
         state->errcode = IMAGING_CODEC_BROKEN;
@@ -624,7 +636,12 @@ ImagingJpeg2KEncodeCleanup(ImagingCodecState state) {
         free((void *)context->error_msg);
     }
 
+    if (context->comment) {
+        free((void *)context->comment);
+    }
+
     context->error_msg = NULL;
+    context->comment = NULL;
 
     return -1;
 }