Merge pull request #2 from radarhere/write-jpeg-com

Use jpeg_write_marker to write comment
bigcat88 · Dec 5, 2022 · 399975f · 399975f
2 parents 1ed1a3a + e50ae85
commit 399975f
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 12 deletions.
diff --git a/src/PIL/JpegImagePlugin.py b/src/PIL/JpegImagePlugin.py
@@ -714,13 +714,6 @@ def validate_qtables(qtables):
 
     extra = info.get("extra", b"")
 
-    comment = info.get("comment", im.info.get("comment"))
-    if comment:
-        if isinstance(comment, str):
-            comment = comment.encode()
-        size = o16(2 + len(comment))
-        extra += b"\xFF\xFE%s%s" % (size, comment)
-
     icc_profile = info.get("icc_profile")
     if icc_profile:
         ICC_OVERHEAD_LEN = 14
@@ -743,6 +736,10 @@ def validate_qtables(qtables):
             )
             i += 1
 
+    comment = info.get("comment", im.info.get("comment")) or b""
+    if isinstance(comment, str):
+        comment = comment.encode()
+
     # "progressive" is the official name, but older documentation
     # says "progression"
     # FIXME: issue a warning if the wrong form is used (post-1.1.7)
@@ -765,6 +762,7 @@ def validate_qtables(qtables):
         dpi[1],
         subsampling,
         qtables,
+        comment,
         extra,
         exif,
     )

diff --git a/src/encode.c b/src/encode.c
@@ -1048,14 +1048,16 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
     PyObject *qtables = NULL;
     unsigned int *qarrays = NULL;
     int qtablesLen = 0;
+    char *comment = NULL;
+    Py_ssize_t comment_size;
     char *extra = NULL;
     Py_ssize_t extra_size;
     char *rawExif = NULL;
     Py_ssize_t rawExifLen = 0;
 
     if (!PyArg_ParseTuple(
             args,
-            "ss|nnnnnnnnOy#y#",
+            "ss|nnnnnnnnOy#y#y#",
             &mode,
             &rawmode,
             &quality,
@@ -1067,6 +1069,8 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
             &ydpi,
             &subsampling,
             &qtables,
+            &comment,
+            &comment_size,
             &extra,
             &extra_size,
             &rawExif,
@@ -1090,12 +1094,24 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    // Freed in JpegEncode, Case 5
+    // Freed in JpegEncode, Case 6
     qarrays = get_qtables_arrays(qtables, &qtablesLen);
 
+    if (comment && comment_size > 0) {
+        /* malloc check ok, length is from python parsearg */
+        char *p = malloc(comment_size);  // Freed in JpegEncode, Case 6
+        if (!p) {
+            return ImagingError_MemoryError();
+        }
+        memcpy(p, comment, comment_size);
+        comment = p;
+    } else {
+        comment = NULL;
+    }
+
     if (extra && extra_size > 0) {
         /* malloc check ok, length is from python parsearg */
-        char *p = malloc(extra_size);  // Freed in JpegEncode, Case 5
+        char *p = malloc(extra_size);  // Freed in JpegEncode, Case 6
         if (!p) {
             return ImagingError_MemoryError();
         }
@@ -1107,7 +1123,7 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
 
     if (rawExif && rawExifLen > 0) {
         /* malloc check ok, length is from python parsearg */
-        char *pp = malloc(rawExifLen);  // Freed in JpegEncode, Case 5
+        char *pp = malloc(rawExifLen);  // Freed in JpegEncode, Case 6
         if (!pp) {
             if (extra) {
                 free(extra);
@@ -1134,6 +1150,8 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
     ((JPEGENCODERSTATE *)encoder->state.context)->streamtype = streamtype;
     ((JPEGENCODERSTATE *)encoder->state.context)->xdpi = xdpi;
     ((JPEGENCODERSTATE *)encoder->state.context)->ydpi = ydpi;
+    ((JPEGENCODERSTATE *)encoder->state.context)->comment = comment;
+    ((JPEGENCODERSTATE *)encoder->state.context)->comment_size = comment_size;
     ((JPEGENCODERSTATE *)encoder->state.context)->extra = extra;
     ((JPEGENCODERSTATE *)encoder->state.context)->extra_size = extra_size;
     ((JPEGENCODERSTATE *)encoder->state.context)->rawExif = rawExif;

diff --git a/src/libImaging/Jpeg.h b/src/libImaging/Jpeg.h
@@ -92,6 +92,10 @@ typedef struct {
     /* in factors of DCTSIZE2 */
     int qtablesLen;
 
+    /* Comment */
+    char *comment;
+    size_t comment_size;
+
     /* Extra data (to be injected after header) */
     char *extra;
     int extra_size;

diff --git a/src/libImaging/JpegEncode.c b/src/libImaging/JpegEncode.c
@@ -277,6 +277,13 @@ ImagingJpegEncode(Imaging im, ImagingCodecState state, UINT8 *buf, int bytes) {
             }
 
         case 4:
+
+            if (context->comment_size > 0) {
+                jpeg_write_marker(&context->cinfo, JPEG_COM, (unsigned char *)context->comment, context->comment_size);
+            }
+            state->state++;
+
+        case 5:
             if (1024 > context->destination.pub.free_in_buffer) {
                 break;
             }
@@ -301,7 +308,7 @@ ImagingJpegEncode(Imaging im, ImagingCodecState state, UINT8 *buf, int bytes) {
             state->state++;
             /* fall through */
 
-        case 5:
+        case 6:
 
             /* Finish compression */
             if (context->destination.pub.free_in_buffer < 100) {
@@ -310,6 +317,10 @@ ImagingJpegEncode(Imaging im, ImagingCodecState state, UINT8 *buf, int bytes) {
             jpeg_finish_compress(&context->cinfo);
 
             /* Clean up */
+            if (context->comment) {
+                free(context->comment);
+                context->comment = NULL;
+            }
             if (context->extra) {
                 free(context->extra);
                 context->extra = NULL;