Added cleanup method to prevent memory leaks from truncated jpeg decodes

PIL/ImageFile.py

@@ -209,6 +209,10 @@ class ImageFile(Image.Image):
                     if not s: # truncated jpeg
                         self.tile = []
 
+                        # JpegDecode needs to clean things up here either way
+                        # If we don't destroy the decompressor, we have a memory leak.
+                        #d.cleanup()
+                        
                         if LOAD_TRUNCATED_IMAGES:
                             break
                         else:

decode.c

@@ -47,7 +47,8 @@
 typedef struct {
     PyObject_HEAD
     int (*decode)(Imaging im, ImagingCodecState state,
-		  UINT8* buffer, int bytes);
+                  UINT8* buffer, int bytes);
+    int (*cleanup)(ImagingCodecState state);
     struct ImagingCodecStateInstance state;
     Imaging im;
     PyObject* lock;
@@ -88,6 +89,9 @@ PyImaging_DecoderNew(int contextsize)
     /* Target image */
     decoder->lock = NULL;
     decoder->im = NULL;
+    
+    /* Initialize the cleanup function pointer */
+    decoder->cleanup = NULL;
 
     return decoder;
 }
@@ -115,6 +119,20 @@ _decode(ImagingDecoderObject* decoder, PyObject* args)
     return Py_BuildValue("ii", status, decoder->state.errcode);
 }
 
+static PyObject* 
+_decode_cleanup(ImagingDecoderObject* decoder, PyObject* args)
+{
+    int status = 0;
+
+    if (decoder->cleanup){
+        status = decoder->cleanup(&decoder->state);
+    }
+
+    return Py_BuildValue("i", status);
+}
+
+
+
 extern Imaging PyImaging_AsImaging(PyObject *op);
 
 static PyObject*
@@ -756,6 +774,7 @@ PyImaging_JpegDecoderNew(PyObject* self, PyObject* args)
         return NULL;
 
     decoder->decode = ImagingJpegDecode;
+    decoder->cleanup = ImagingJpegDecodeCleanup;
 
     strncpy(((JPEGSTATE*)decoder->state.context)->rawmode, rawmode, 8);
     strncpy(((JPEGSTATE*)decoder->state.context)->jpegmode, jpegmode, 8);

libImaging/Imaging.h

@@ -417,6 +417,8 @@ extern int ImagingHexDecode(Imaging im, ImagingCodecState state,
 #ifdef	HAVE_LIBJPEG
 extern int ImagingJpegDecode(Imaging im, ImagingCodecState state,
 			     UINT8* buffer, int bytes);
+extern int ImagingJpegDecodeCleanup(ImagingCodecState state);
+
 extern int ImagingJpegEncode(Imaging im, ImagingCodecState state,
 			     UINT8* buffer, int bytes);
 #endif

libImaging/JpegDecode.c

@@ -263,5 +263,20 @@ ImagingJpegDecode(Imaging im, ImagingCodecState state, UINT8* buf, int bytes)
 
 }
 
+/* -------------------------------------------------------------------- */
+/* Cleanup                                                              */
+/* -------------------------------------------------------------------- */
+
+int ImagingJpegDecodeCleanup(ImagingCodecState state){
+	/* called to fee the decompression engine when the decode terminates 
+	   due to a corrupt or truncated image 
+	*/
+    JPEGSTATE* context = (JPEGSTATE*) state->context;
+
+	/* Clean up */
+	jpeg_destroy_decompress(&context->cinfo);
+	return -1;
+}
+
 #endif