diff --git a/Tests/test_file_jpeg2k.py b/Tests/test_file_jpeg2k.py index a869d74f0ce..9b4774a8229 100644 --- a/Tests/test_file_jpeg2k.py +++ b/Tests/test_file_jpeg2k.py @@ -373,13 +373,13 @@ def test_comment(): def test_save_comment(): - comment = "Created by Pillow" - out = BytesIO() - test_card.save(out, "JPEG2000", comment=comment) - out.seek(0) + for comment in ("Created by Pillow", b"Created by Pillow"): + out = BytesIO() + test_card.save(out, "JPEG2000", comment=comment) + out.seek(0) - with Image.open(out) as im: - assert im.info["comment"] == b"Created by Pillow" + with Image.open(out) as im: + assert im.info["comment"] == b"Created by Pillow" too_long_comment = " " * 65532 with pytest.raises(ValueError): diff --git a/src/PIL/Jpeg2KImagePlugin.py b/src/PIL/Jpeg2KImagePlugin.py index 4249fe7141c..980c299dbce 100644 --- a/src/PIL/Jpeg2KImagePlugin.py +++ b/src/PIL/Jpeg2KImagePlugin.py @@ -351,10 +351,12 @@ def _save(im, fp, filename): cinema_mode = info.get("cinema_mode", "no") mct = info.get("mct", 0) signed = info.get("signed", False) - fd = -1 comment = info.get("comment") + if isinstance(comment, str): + comment = comment.encode() add_plt = info.get("add_plt", False) + fd = -1 if hasattr(fp, "fileno"): try: fd = fp.fileno() diff --git a/src/encode.c b/src/encode.c index e8946dbaef1..2f824ef80fb 100644 --- a/src/encode.c +++ b/src/encode.c @@ -1215,11 +1215,12 @@ PyImaging_Jpeg2KEncoderNew(PyObject *self, PyObject *args) { int sgnd = 0; Py_ssize_t fd = -1; char *comment = NULL; + Py_ssize_t comment_size; int add_plt = 0; if (!PyArg_ParseTuple( args, - "ss|OOOsOnOOOssbbnzp", + "ss|OOOsOnOOOssbbnz#p", &mode, &format, &offset, @@ -1237,6 +1238,7 @@ PyImaging_Jpeg2KEncoderNew(PyObject *self, PyObject *args) { &sgnd, &fd, &comment, + &comment_size, &add_plt)) { return NULL; } @@ -1319,9 +1321,9 @@ PyImaging_Jpeg2KEncoderNew(PyObject *self, PyObject *args) { } } - if (comment != NULL && strlen(comment) > 0) { + if (comment && comment_size > 0) { /* Size is stored as as an uint16, subtract 4 bytes for the header */ - if (strlen(comment) >= 65531) { + if (comment_size >= 65531) { PyErr_SetString( PyExc_ValueError, "JPEG 2000 comment is too long"); @@ -1329,15 +1331,13 @@ PyImaging_Jpeg2KEncoderNew(PyObject *self, PyObject *args) { return NULL; } - context->comment = strdup(comment); - - if (context->comment == NULL) { - PyErr_SetString( - PyExc_MemoryError, - "Couldn't allocate memory for JPEG 2000 comment"); + char *p = malloc(comment_size); + if (!p) { Py_DECREF(encoder); - return NULL; + return ImagingError_MemoryError(); } + memcpy(p, comment, comment_size + 1); + context->comment = p; } if (quality_layers && PySequence_Check(quality_layers)) {