Issue #18227: Use PyMem_RawAlloc() in bz2, lzma and zlib modules

3 files changed, 77 insertions, 8 deletions

diff --git a/Modules/_bz2module.c b/Modules/_bz2module.c
index 46da5ab..a67aec5 100644
--- a/Modules/_bz2module.c
+++ b/Modules/_bz2module.c
@@ -248,6 +248,24 @@ BZ2Compressor_flush(BZ2Compressor *self, PyObject *noargs)
     return result;
 }
 
+static void*
+BZ2_Malloc(void* ctx, int items, int size)
+{
+    if (items < 0 || size < 0)
+        return NULL;
+    if ((size_t)items > (size_t)PY_SSIZE_T_MAX / (size_t)size)
+        return NULL;
+    /* PyMem_Malloc() cannot be used: compress() and decompress()
+       release the GIL */
+    return PyMem_RawMalloc(items * size);
+}
+
+static void
+BZ2_Free(void* ctx, void *ptr)
+{
+    return PyMem_RawFree(ptr);
+}
+
 static int
 BZ2Compressor_init(BZ2Compressor *self, PyObject *args, PyObject *kwargs)
 {
@@ -270,6 +288,9 @@ BZ2Compressor_init(BZ2Compressor *self, PyObject *args, PyObject *kwargs)
     }
 #endif
 
+    self->bzs.opaque = NULL;
+    self->bzs.bzalloc = BZ2_Malloc;
+    self->bzs.bzfree = BZ2_Free;
     bzerror = BZ2_bzCompressInit(&self->bzs, compresslevel, 0, 0);
     if (catch_bz2_error(bzerror))
         goto error;
diff --git a/Modules/_lzmamodule.c b/Modules/_lzmamodule.c
index b482a77..711604d 100644
--- a/Modules/_lzmamodule.c
+++ b/Modules/_lzmamodule.c
@@ -51,6 +51,7 @@ enum {
 
 typedef struct {
     PyObject_HEAD
+    lzma_allocator alloc;
     lzma_stream lzs;
     int flushed;
 #ifdef WITH_THREAD
@@ -60,6 +61,7 @@ typedef struct {
 
 typedef struct {
     PyObject_HEAD
+    lzma_allocator alloc;
     lzma_stream lzs;
     int check;
     char eof;
@@ -117,6 +119,22 @@ catch_lzma_error(lzma_ret lzret)
     }
 }
 
+static void*
+PyLzma_Malloc(void *opaque, size_t items, size_t size)
+{
+    if (items > (size_t)PY_SSIZE_T_MAX / size)
+        return NULL;
+    /* PyMem_Malloc() cannot be used:
+       the GIL is not held when lzma_code() is called */
+    return PyMem_RawMalloc(items * size);
+}
+
+static void
+PyLzma_Free(void *opaque, void *ptr)
+{
+    return PyMem_RawFree(ptr);
+}
+
 #if BUFSIZ < 8192
 #define INITIAL_BUFFER_SIZE 8192
 #else
@@ -656,6 +674,11 @@ Compressor_init(Compressor *self, PyObject *args, PyObject *kwargs)
         if (!uint32_converter(preset_obj, &preset))
             return -1;
 
+    self->alloc.opaque = NULL;
+    self->alloc.alloc = PyLzma_Malloc;
+    self->alloc.free = PyLzma_Free;
+    self->lzs.allocator = &self->alloc;
+
 #ifdef WITH_THREAD
     self->lock = PyThread_allocate_lock();
     if (self->lock == NULL) {
@@ -922,6 +945,11 @@ Decompressor_init(Decompressor *self, PyObject *args, PyObject *kwargs)
         return -1;
     }
 
+    self->alloc.opaque = NULL;
+    self->alloc.alloc = PyLzma_Malloc;
+    self->alloc.free = PyLzma_Free;
+    self->lzs.allocator = &self->alloc;
+
 #ifdef WITH_THREAD
     self->lock = PyThread_allocate_lock();
     if (self->lock == NULL) {
diff --git a/Modules/zlibmodule.c b/Modules/zlibmodule.c
index 30147ae..2e5f473 100644
--- a/Modules/zlibmodule.c
+++ b/Modules/zlibmodule.c
@@ -136,6 +136,22 @@ newcompobject(PyTypeObject *type)
     return self;
 }
 
+static void*
+PyZlib_Malloc(voidpf ctx, uInt items, uInt size)
+{
+    if (items > (size_t)PY_SSIZE_T_MAX / size)
+        return NULL;
+    /* PyMem_Malloc() cannot be used: the GIL is not held when
+       inflate() and deflate() are called */
+    return PyMem_RawMalloc(items * size);
+}
+
+static void
+PyZlib_Free(voidpf ctx, void *ptr)
+{
+    return PyMem_RawFree(ptr);
+}
+
 PyDoc_STRVAR(compress__doc__,
 "compress(string[, level]) -- Returned compressed string.\n"
 "\n"
@@ -175,8 +191,9 @@ PyZlib_compress(PyObject *self, PyObject *args)
     /* Past the point of no return.  From here on out, we need to make sure
        we clean up mallocs & INCREFs. */
 
-    zst.zalloc = (alloc_func)NULL;
-    zst.zfree = (free_func)Z_NULL;
+    zst.opaque = NULL;
+    zst.zalloc = PyZlib_Malloc;
+    zst.zfree = PyZlib_Free;
     zst.next_out = (Byte *)output;
     zst.next_in = (Byte *)input;
     zst.avail_in = length;
@@ -262,8 +279,9 @@ PyZlib_decompress(PyObject *self, PyObject *args)
     if (!(result_str = PyBytes_FromStringAndSize(NULL, r_strlen)))
         goto error;
 
-    zst.zalloc = (alloc_func)NULL;
-    zst.zfree = (free_func)Z_NULL;
+    zst.opaque = NULL;
+    zst.zalloc = PyZlib_Malloc;
+    zst.zfree = PyZlib_Free;
     zst.next_out = (Byte *)PyBytes_AS_STRING(result_str);
     zst.next_in = (Byte *)input;
     err = inflateInit2(&zst, wsize);
@@ -356,8 +374,9 @@ PyZlib_compressobj(PyObject *selfptr, PyObject *args, PyObject *kwargs)
     self = newcompobject(&Comptype);
     if (self==NULL)
         goto error;
-    self->zst.zalloc = (alloc_func)NULL;
-    self->zst.zfree = (free_func)Z_NULL;
+    self->zst.opaque = NULL;
+    self->zst.zalloc = PyZlib_Malloc;
+    self->zst.zfree = PyZlib_Free;
     self->zst.next_in = NULL;
     self->zst.avail_in = 0;
     err = deflateInit2(&self->zst, level, method, wbits, memLevel, strategy);
@@ -420,8 +439,9 @@ PyZlib_decompressobj(PyObject *selfptr, PyObject *args, PyObject *kwargs)
     self = newcompobject(&Decomptype);
     if (self == NULL)
         return(NULL);
-    self->zst.zalloc = (alloc_func)NULL;
-    self->zst.zfree = (free_func)Z_NULL;
+    self->zst.opaque = NULL;
+    self->zst.zalloc = PyZlib_Malloc;
+    self->zst.zfree = PyZlib_Free;
     self->zst.next_in = NULL;
     self->zst.avail_in = 0;
     if (zdict != NULL) {