Memory Leak: Freeing malloc'd memory in Jpeg Encode

encode.c

@@ -535,12 +535,12 @@ PyImaging_ZipEncoderNew(PyObject* self, PyObject* args)
 
 #include "Jpeg.h"
 
-static unsigned int** get_qtables_arrays(PyObject* qtables, int* qtablesLen) {
+static unsigned int* get_qtables_arrays(PyObject* qtables, int* qtablesLen) {
     PyObject* tables;
     PyObject* table;
     PyObject* table_data;
     int i, j, num_tables;
-    unsigned int **qarrays;
+    unsigned int *qarrays;
 
     if ((qtables ==  NULL) || (qtables == Py_None)) {
         return NULL;
@@ -554,10 +554,12 @@ static unsigned int** get_qtables_arrays(PyObject* qtables, int* qtablesLen) {
     tables = PySequence_Fast(qtables, "expected a sequence");
     num_tables = PySequence_Size(qtables);
     if (num_tables < 1 || num_tables > NUM_QUANT_TBLS) {
-        PyErr_SetString(PyExc_ValueError, "Not a valid numbers of quantization tables. Should be between 1 and 4.");
+        PyErr_SetString(PyExc_ValueError, 
+            "Not a valid number of quantization tables. Should be between 1 and 4.");
+        Py_DECREF(tables);
         return NULL;
     }
-    qarrays = (unsigned int**) PyMem_Malloc(num_tables * sizeof(unsigned int*));
+    qarrays = (unsigned int*) malloc(num_tables * DCTSIZE2 * sizeof(unsigned int));
     if (!qarrays) {
         Py_DECREF(tables);
         PyErr_NoMemory();
@@ -566,39 +568,32 @@ static unsigned int** get_qtables_arrays(PyObject* qtables, int* qtablesLen) {
     for (i = 0; i < num_tables; i++) {
         table = PySequence_Fast_GET_ITEM(tables, i);
         if (!PySequence_Check(table)) {
-            Py_DECREF(tables);
             PyErr_SetString(PyExc_ValueError, "Invalid quantization tables");
-            return NULL;
+            goto JPEG_QTABLES_ERR;
         }
         if (PySequence_Size(table) != DCTSIZE2) {
-            Py_DECREF(tables);
+            PyErr_SetString(PyExc_ValueError, "Invalid quantization table size");
-            PyErr_SetString(PyExc_ValueError, "Invalid quantization tables");
+            goto JPEG_QTABLES_ERR;
-            return NULL;
         }
         table_data = PySequence_Fast(table, "expected a sequence");
-        qarrays[i] = (unsigned int*) PyMem_Malloc(DCTSIZE2 * sizeof(unsigned int));
-        if (!qarrays[i]) {
-            Py_DECREF(tables);
-            PyErr_NoMemory();
-            return NULL;
-        }
         for (j = 0; j < DCTSIZE2; j++) {
-            qarrays[i][j] = PyInt_AS_LONG(PySequence_Fast_GET_ITEM(table_data, j));
+            qarrays[i * DCTSIZE2 + j] = PyInt_AS_LONG(PySequence_Fast_GET_ITEM(table_data, j));
         }
     }
 
-    Py_DECREF(tables);
     *qtablesLen = num_tables; 
 
+JPEG_QTABLES_ERR:
+    Py_DECREF(tables);  // Run on both error and not error
     if (PyErr_Occurred()) {
-        PyMem_Free(qarrays);
+        free(qarrays);
         qarrays = NULL;
+        return NULL;
     }
 
     return qarrays;
 }
 
-
 PyObject*
 PyImaging_JpegEncoderNew(PyObject* self, PyObject* args)
 {
@@ -614,7 +609,7 @@ PyImaging_JpegEncoderNew(PyObject* self, PyObject* args)
     int xdpi = 0, ydpi = 0;
     int subsampling = -1; /* -1=default, 0=none, 1=medium, 2=high */
     PyObject* qtables=NULL;
-    unsigned int **qarrays = NULL;
+    unsigned int *qarrays = NULL;
     int qtablesLen = 0;
     char* extra = NULL;
     int extra_size;
@@ -638,7 +633,7 @@ PyImaging_JpegEncoderNew(PyObject* self, PyObject* args)
     qarrays = get_qtables_arrays(qtables, &qtablesLen);
 
     if (extra && extra_size > 0) {
-        char* p = malloc(extra_size);
+        char* p = malloc(extra_size); // Freed in JpegEncode, Case 5
         if (!p)
             return PyErr_NoMemory();
         memcpy(p, extra, extra_size);
@@ -647,7 +642,7 @@ PyImaging_JpegEncoderNew(PyObject* self, PyObject* args)
         extra = NULL;
 
     if (rawExif && rawExifLen > 0) {
-        char* pp = malloc(rawExifLen);
+        char* pp = malloc(rawExifLen); // Freed in JpegEncode, Case 5
         if (!pp)
             return PyErr_NoMemory();
         memcpy(pp, rawExif, rawExifLen);

libImaging/Jpeg.h

@@ -89,7 +89,9 @@ typedef struct {
     int subsampling;
 
     /* Custom quantization tables () */
-    unsigned int **qtables;
+    unsigned int *qtables;
+
+    /* in factors of DCTSIZE2 */
     int qtablesLen;
 
     /* Extra data (to be injected after header) */

libImaging/JpegEncode.c

@@ -153,7 +153,7 @@ ImagingJpegEncode(Imaging im, ImagingCodecState state, UINT8* buf, int bytes)
 	    }
 	    for (i = 0; i < context->qtablesLen; i++) {
 		// TODO: Should add support for none baseline
-		jpeg_add_quant_table(&context->cinfo, i, context->qtables[i],
+		jpeg_add_quant_table(&context->cinfo, i, &context->qtables[i * DCTSIZE2],
 				     quality, TRUE);
 	    }
 	} else if (context->quality > 0) {
@@ -289,8 +289,19 @@ ImagingJpegEncode(Imaging im, ImagingCodecState state, UINT8* buf, int bytes)
 	jpeg_finish_compress(&context->cinfo);
 
 	/* Clean up */
-        if (context->extra)
+        if (context->extra) {
             free(context->extra);
+            context->extra = NULL;
+        }
+        if (context->rawExif) {
+            free(context->rawExif);
+            context->rawExif = NULL;
+        }
+        if (context->qtables) {
+            free(context->qtables);
+            context->qtables = NULL;
+        }            
+
 	jpeg_destroy_compress(&context->cinfo);
 	/* if (jerr.pub.num_warnings) return BROKEN; */
 	state->errcode = IMAGING_CODEC_END;