Merge pull request #898 from wiredfool/joshware-j2k-leak

Jpeg2k Decode/Encode Memory Leak Fix

PIL/ImageFile.py

@@ -227,6 +227,8 @@ class ImageFile(Image.Image):
                         break
                     b = b[n:]
                     t = t + n
+                # Need to cleanup here to prevent leaks in PyPy
+                d.cleanup()
 
         self.tile = []
         self.readonly = readonly
@@ -471,6 +473,7 @@ def _save(im, fp, tile, bufsize=0):
                     break
             if s < 0:
                 raise IOError("encoder error %d when writing image file" % s)
+            e.cleanup()
     else:
         # slight speedup: compress to real file object
         for e, b, o, a in tile:
@@ -481,6 +484,7 @@ def _save(im, fp, tile, bufsize=0):
             s = e.encode_to_file(fh, bufsize)
             if s < 0:
                 raise IOError("encoder error %d when writing image file" % s)
+            e.cleanup()
     try:
         fp.flush()
     except: pass

Tests/check_j2k_leaks.py

@@ -0,0 +1,42 @@
+from helper import unittest, PillowTestCase
+import sys
+from PIL import Image
+from io import BytesIO
+
+# Limits for testing the leak
+mem_limit = 1024*1048576
+stack_size = 8*1048576
+iterations = int((mem_limit/stack_size)*2)
+codecs = dir(Image.core)
+test_file = "Tests/images/rgb_trns_ycbc.jp2"
+
+
+@unittest.skipIf(sys.platform.startswith('win32'), "requires Unix or MacOS")
+class TestJpegLeaks(PillowTestCase):
+    def setUp(self):
+        if "jpeg2k_encoder" not in codecs or "jpeg2k_decoder" not in codecs:
+            self.skipTest('JPEG 2000 support not available')
+
+    def test_leak_load(self):
+        from resource import setrlimit, RLIMIT_AS, RLIMIT_STACK
+        setrlimit(RLIMIT_STACK, (stack_size, stack_size))
+        setrlimit(RLIMIT_AS, (mem_limit, mem_limit))
+        for count in range(iterations):
+            with Image.open(test_file) as im:
+                im.load()
+
+    def test_leak_save(self):
+        from resource import setrlimit, RLIMIT_AS, RLIMIT_STACK
+        setrlimit(RLIMIT_STACK, (stack_size, stack_size))
+        setrlimit(RLIMIT_AS, (mem_limit, mem_limit))
+        for count in range(iterations):
+            with Image.open(test_file) as im:
+                im.load()
+                test_output = BytesIO()
+                im.save(test_output, "JPEG2000")
+                test_output.seek(0)
+                output = test_output.read()
+
+
+if __name__ == '__main__':
+    unittest.main()

decode.c

@@ -103,6 +103,8 @@ PyImaging_DecoderNew(int contextsize)
 static void
 _dealloc(ImagingDecoderObject* decoder)
 {
+    if (decoder->cleanup)
+        decoder->cleanup(&decoder->state);
     free(decoder->state.buffer);
     free(decoder->state.context);
     Py_XDECREF(decoder->lock);

encode.c

@@ -99,6 +99,18 @@ _dealloc(ImagingEncoderObject* encoder)
     PyObject_Del(encoder);
 }
 
+static PyObject*
+_encode_cleanup(ImagingEncoderObject* encoder, PyObject* args)
+{
+    int status = 0;
+
+    if (encoder->cleanup){
+        status = encoder->cleanup(&encoder->state);
+    }
+
+    return Py_BuildValue("i", status);
+}
+
 static PyObject*
 _encode(ImagingEncoderObject* encoder, PyObject* args)
 {
@@ -239,6 +251,7 @@ _setimage(ImagingEncoderObject* encoder, PyObject* args)
 
 static struct PyMethodDef methods[] = {
     {"encode", (PyCFunction)_encode, 1},
+    {"cleanup", (PyCFunction)_encode_cleanup, 1},
     {"encode_to_file", (PyCFunction)_encode_to_file, 1},
     {"setimage", (PyCFunction)_setimage, 1},
     {NULL, NULL} /* sentinel */

libImaging/Jpeg2KDecode.c

@@ -800,6 +800,11 @@ ImagingJpeg2KDecodeCleanup(ImagingCodecState state) {
     if (context->decoder)
         ImagingIncrementalCodecDestroy(context->decoder);
 
+    context->error_msg = NULL;
+
+    /* Prevent multiple calls to ImagingIncrementalCodecDestroy */
+    context->decoder = NULL;
+
     return -1;
 }
 

libImaging/Jpeg2KEncode.c

@@ -576,15 +576,20 @@ int
 ImagingJpeg2KEncodeCleanup(ImagingCodecState state) {
     JPEG2KENCODESTATE *context = (JPEG2KENCODESTATE *)state->context;
 
-    if (context->quality_layers)
+    if (context->quality_layers && context->encoder)
         Py_DECREF(context->quality_layers);
 
     if (context->error_msg)
         free ((void *)context->error_msg);
 
+    context->error_msg = NULL;
+
     if (context->encoder)
         ImagingIncrementalCodecDestroy(context->encoder);
 
+    /* Prevent multiple calls to ImagingIncrementalCodecDestroy */
+    context->encoder = NULL;
+
     return -1;
 }