diff --git a/Modules/_bz2module.c b/Modules/_bz2module.c index 46da5abad27..a67aec505bf 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 b482a7767db..711604de00c 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 30147aeef3b..2e5f4735740 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) {