diff --git a/PIL/PngImagePlugin.py b/PIL/PngImagePlugin.py index 5e5eb14c7..e329a6c0d 100644 --- a/PIL/PngImagePlugin.py +++ b/PIL/PngImagePlugin.py @@ -113,7 +113,8 @@ class ChunkStream(object): length = i32(s) if not is_cid(cid): - raise SyntaxError("broken PNG file (chunk %s)" % repr(cid)) + if not ImageFile.LOAD_TRUNCATED_IMAGES: + raise SyntaxError("broken PNG file (chunk %s)" % repr(cid)) return cid, pos, length diff --git a/Tests/test_file_png.py b/Tests/test_file_png.py index f9b9a1eb4..8e9bdae3b 100644 --- a/Tests/test_file_png.py +++ b/Tests/test_file_png.py @@ -1,14 +1,16 @@ from helper import unittest, PillowTestCase, hopper +from PIL import Image, ImageFile, PngImagePlugin from io import BytesIO - -from PIL import Image -from PIL import ImageFile -from PIL import PngImagePlugin import zlib +import sys codecs = dir(Image.core) +# For Truncated phng memory leak +MEM_LIMIT = 1 # max increase in MB +ITERATIONS = 100 + # sample png stream TEST_PNG_FILE = "Tests/images/hopper.png" @@ -530,5 +532,33 @@ class TestFilePng(PillowTestCase): self.assertLess(chunks.index(b"pHYs"), chunks.index(b"IDAT")) +@unittest.skipIf(sys.platform.startswith('win32'), "requires Unix or MacOS") +class TestTruncatedPngPLeaks(PillowTestCase): + + def setUp(self): + if "zip_encoder" not in codecs or "zip_decoder" not in codecs: + self.skipTest("zip/deflate support not available") + + def _get_mem_usage(self): + from resource import getpagesize, getrusage, RUSAGE_SELF + mem = getrusage(RUSAGE_SELF).ru_maxrss + return mem * getpagesize() / 1024 / 1024 + + def test_leak_load(self): + with open('Tests/images/hopper.png', 'rb') as f: + DATA = BytesIO(f.read(16 * 1024)) + + ImageFile.LOAD_TRUNCATED_IMAGES = True + start_mem = self._get_mem_usage() + try: + for _ in range(ITERATIONS): + with Image.open(DATA) as im: + im.load() + mem = (self._get_mem_usage() - start_mem) + self.assertLess(mem, MEM_LIMIT, msg='memory usage limit exceeded') + finally: + ImageFile.LOAD_TRUNCATED_IMAGES = False + + if __name__ == '__main__': unittest.main() diff --git a/decode.c b/decode.c index d7fe02fae..f749a40a7 100644 --- a/decode.c +++ b/decode.c @@ -779,6 +779,7 @@ PyImaging_ZipDecoderNew(PyObject* self, PyObject* args) return NULL; decoder->decode = ImagingZipDecode; + decoder->cleanup = ImagingZipDecodeCleanup; ((ZIPSTATE*)decoder->state.context)->interlaced = interlaced; diff --git a/libImaging/Imaging.h b/libImaging/Imaging.h index 2030821f2..99fff7f67 100644 --- a/libImaging/Imaging.h +++ b/libImaging/Imaging.h @@ -455,6 +455,7 @@ extern int ImagingXbmEncode(Imaging im, ImagingCodecState state, #ifdef HAVE_LIBZ extern int ImagingZipDecode(Imaging im, ImagingCodecState state, UINT8* buffer, int bytes); +extern int ImagingZipDecodeCleanup(ImagingCodecState state); extern int ImagingZipEncode(Imaging im, ImagingCodecState state, UINT8* buffer, int bytes); extern int ImagingZipEncodeCleanup(ImagingCodecState state); diff --git a/libImaging/ZipDecode.c b/libImaging/ZipDecode.c index 37cb2866c..e96e3200c 100644 --- a/libImaging/ZipDecode.c +++ b/libImaging/ZipDecode.c @@ -85,6 +85,8 @@ ImagingZipDecode(Imaging im, ImagingCodecState state, UINT8* buf, int bytes) err = inflateInit(&context->z_stream); if (err < 0) { state->errcode = IMAGING_CODEC_CONFIG; + free(context->previous); + context->previous = NULL; return -1; } @@ -126,6 +128,7 @@ ImagingZipDecode(Imaging im, ImagingCodecState state, UINT8* buf, int bytes) else state->errcode = IMAGING_CODEC_CONFIG; free(context->previous); + context->previous = NULL; inflateEnd(&context->z_stream); return -1; } @@ -191,6 +194,7 @@ ImagingZipDecode(Imaging im, ImagingCodecState state, UINT8* buf, int bytes) default: state->errcode = IMAGING_CODEC_UNKNOWN; free(context->previous); + context->previous = NULL; inflateEnd(&context->z_stream); return -1; } @@ -258,6 +262,7 @@ ImagingZipDecode(Imaging im, ImagingCodecState state, UINT8* buf, int bytes) state->errcode = IMAGING_CODEC_BROKEN; */ free(context->previous); + context->previous = NULL; inflateEnd(&context->z_stream); return -1; /* end of file (errcode=0) */ @@ -274,4 +279,20 @@ ImagingZipDecode(Imaging im, ImagingCodecState state, UINT8* buf, int bytes) } + +int ImagingZipDecodeCleanup(ImagingCodecState state){ + /* called to free the decompression engine when the decode terminates + due to a corrupt or truncated image + */ + ZIPSTATE* context = (ZIPSTATE*) state->context; + + /* Clean up */ + if (context->previous) { + inflateEnd(&context->z_stream); + free(context->previous); + context->previous = NULL; + } + return -1; +} + #endif