Merge pull request #2541 from uploadcare/fix-truncated-png-loading

Fix truncated png loading

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
 

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()

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;
 

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);

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