diff --git a/PIL/ImageFile.py b/PIL/ImageFile.py index 5e4745d76..bdcc769ec 100644 --- a/PIL/ImageFile.py +++ b/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 diff --git a/Tests/check_j2k_leaks.py b/Tests/check_j2k_leaks.py new file mode 100755 index 000000000..9dbb8c1f4 --- /dev/null +++ b/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() diff --git a/decode.c b/decode.c index e9aa6a387..c56f42592 100644 --- a/decode.c +++ b/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); diff --git a/encode.c b/encode.c index 3fa900b1d..34e3b8933 100644 --- a/encode.c +++ b/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 */ diff --git a/libImaging/Jpeg2KDecode.c b/libImaging/Jpeg2KDecode.c index 97ec81003..abf8cebbe 100644 --- a/libImaging/Jpeg2KDecode.c +++ b/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; } diff --git a/libImaging/Jpeg2KEncode.c b/libImaging/Jpeg2KEncode.c old mode 100644 new mode 100755 index e8eef08c1..ea4bca2f2 --- a/libImaging/Jpeg2KEncode.c +++ b/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; }