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

1 files changed, 28 insertions, 0 deletions

diff --git a/Modules/_lzmamodule.c b/Modules/_lzmamodule.c
index b482a7767d..711604de00 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) {