#ifndef Py_BUILD_CORE_BUILTIN

# define Py_BUILD_CORE_MODULE 1

#include "Python.h"

#include "pycore_pyatomic_ft_wrappers.h" // FT_ATOMIC_STORE_CHAR_RELAXED

#include "zlib.h"

#include "stdbool.h"

#include <stddef.h> // offsetof()

#if defined(ZLIB_VERNUM) && ZLIB_VERNUM < 0x1221

#error "At least zlib version 1.2.2.1 is required"

#if (SIZEOF_OFF_T == SIZEOF_SIZE_T)

# define convert_to_z_off_t PyLong_AsSsize_t

#elif (SIZEOF_OFF_T == SIZEOF_LONG_LONG)

# define convert_to_z_off_t PyLong_AsLongLong

#elif (SIZEOF_OFF_T == SIZEOF_LONG)

# define convert_to_z_off_t PyLong_AsLong

# error off_t does not match either size_t, long, or long long!

#include "pycore_blocks_output_buffer.h"

#if OUTPUT_BUFFER_MAX_BLOCK_SIZE > UINT32_MAX

#error "The maximum block size accepted by zlib is UINT32_MAX."

static inline Py_ssize_t

OutputBuffer_InitAndGrow(_BlocksOutputBuffer *buffer, Py_ssize_t max_length,

Bytef **next_out, uint32_t *avail_out)

{

Py_ssize_t allocated;

allocated = _BlocksOutputBuffer_InitAndGrow(

buffer, max_length, (void**) next_out);

*avail_out = (uint32_t) allocated;

return allocated;

}

static inline Py_ssize_t

OutputBuffer_Grow(_BlocksOutputBuffer *buffer,

Bytef **next_out, uint32_t *avail_out)

{

Py_ssize_t allocated;

allocated = _BlocksOutputBuffer_Grow(

buffer, (void**) next_out, (Py_ssize_t) *avail_out);

*avail_out = (uint32_t) allocated;

return allocated;

}

static inline Py_ssize_t

OutputBuffer_GetDataSize(_BlocksOutputBuffer *buffer, uint32_t avail_out)

{

return _BlocksOutputBuffer_GetDataSize(buffer, (Py_ssize_t) avail_out);

}

static inline PyObject *

OutputBuffer_Finish(_BlocksOutputBuffer *buffer, uint32_t avail_out)

{

return _BlocksOutputBuffer_Finish(buffer, (Py_ssize_t) avail_out);

}

static inline void

OutputBuffer_OnError(_BlocksOutputBuffer *buffer)

{

_BlocksOutputBuffer_OnError(buffer);

}

typedef struct {

Py_ssize_t left_bytes;

Bytef *next_posi;

} _Uint32Window;

static inline Py_ssize_t

OutputBuffer_WindowInitWithSize(_BlocksOutputBuffer *buffer, _Uint32Window *window,

Py_ssize_t init_size,

Bytef **next_out, uint32_t *avail_out)

{

Py_ssize_t allocated = _BlocksOutputBuffer_InitWithSize(

buffer, init_size, (void**) next_out);

if (allocated >= 0) {

// the UINT32_MAX sliding window

Py_ssize_t window_size = Py_MIN((size_t)allocated, UINT32_MAX);

*avail_out = (uint32_t) window_size;

window->left_bytes = allocated - window_size;

window->next_posi = *next_out + window_size;

}

return allocated;

}

static inline Py_ssize_t

OutputBuffer_WindowGrow(_BlocksOutputBuffer *buffer, _Uint32Window *window,

Bytef **next_out, uint32_t *avail_out)

{

Py_ssize_t allocated;

/* ensure no gaps in the data.

if (*avail_out != 0) {

PyErr_SetString(PyExc_SystemError,

"*avail_out != 0 in OutputBuffer_WindowGrow().");

return -1;

}

// slide the UINT32_MAX sliding window

if (window->left_bytes > 0) {

Py_ssize_t window_size = Py_MIN((size_t)window->left_bytes, UINT32_MAX);

*next_out = window->next_posi;

*avail_out = (uint32_t) window_size;

window->left_bytes -= window_size;

window->next_posi += window_size;

return window_size;

}

assert(window->left_bytes == 0);

// only the first block may > UINT32_MAX

allocated = _BlocksOutputBuffer_Grow(

buffer, (void**) next_out, (Py_ssize_t) *avail_out);

*avail_out = (uint32_t) allocated;

return allocated;

}

static inline PyObject *

OutputBuffer_WindowFinish(_BlocksOutputBuffer *buffer, _Uint32Window *window,

uint32_t avail_out)

{

Py_ssize_t real_avail_out = (Py_ssize_t) avail_out + window->left_bytes;

return _BlocksOutputBuffer_Finish(buffer, real_avail_out);

}

static inline void

OutputBuffer_WindowOnError(_BlocksOutputBuffer *buffer, _Uint32Window *window)

{

_BlocksOutputBuffer_OnError(buffer);

}

#define DEFLATED 8

#if MAX_MEM_LEVEL >= 8

# define DEF_MEM_LEVEL 8

# define DEF_MEM_LEVEL MAX_MEM_LEVEL

#define DEF_BUF_SIZE (16*1024)

#define DEF_MAX_INITIAL_BUF_SIZE (16 * 1024 * 1024)

static PyModuleDef zlibmodule;

typedef struct {

PyTypeObject *Comptype;

PyTypeObject *Decomptype;

PyTypeObject *ZlibDecompressorType;

PyObject *ZlibError;

} zlibstate;

static inline zlibstate*

get_zlib_state(PyObject *module)

{

void *state = PyModule_GetState(module);

assert(state != NULL);

return (zlibstate *)state;

}

typedef struct

{

PyObject_HEAD

z_stream zst;

PyObject *unused_data;

PyObject *unconsumed_tail;

char eof;

bool is_initialised;

PyObject *zdict;

PyMutex mutex;

} compobject;

#define _compobject_CAST(op) ((compobject *)op)

static void

zlib_error(zlibstate *state, z_stream zst, int err, const char *msg)

{

const char *zmsg = Z_NULL;

/* In case of a version mismatch, zst.msg won't be initialized.

if (err == Z_VERSION_ERROR)

zmsg = "library version mismatch";

if (zmsg == Z_NULL)

zmsg = zst.msg;

if (zmsg == Z_NULL) {

switch (err) {

case Z_BUF_ERROR:

zmsg = "incomplete or truncated stream";

break;

case Z_STREAM_ERROR:

zmsg = "inconsistent stream state";

break;

case Z_DATA_ERROR:

zmsg = "invalid input data";

break;

}

}

if (zmsg == Z_NULL)

PyErr_Format(state->ZlibError, "Error %d %s", err, msg);

else

PyErr_Format(state->ZlibError, "Error %d %s: %.200s", err, msg, zmsg);

}

static compobject *

newcompobject(PyTypeObject *type)

{

compobject *self;

assert(type != NULL);

assert(type->tp_alloc != NULL);

self = _compobject_CAST(type->tp_alloc(type, 0));

if (self == NULL)

return NULL;

self->eof = 0;

self->is_initialised = 0;

self->zdict = NULL;

self->unused_data = Py_GetConstant(Py_CONSTANT_EMPTY_BYTES);

if (self->unused_data == NULL) {

Py_DECREF(self);

return NULL;

}

self->unconsumed_tail = Py_GetConstant(Py_CONSTANT_EMPTY_BYTES);

if (self->unconsumed_tail == NULL) {

Py_DECREF(self);

return NULL;

}

self->mutex = (PyMutex){0};

return self;

}

static void*

PyZlib_Malloc(voidpf ctx, uInt items, uInt size)

{

if (size != 0 && items > (size_t)PY_SSIZE_T_MAX / size)

return NULL;

/* PyMem_Malloc() cannot be used: the GIL is not held when

return PyMem_RawMalloc((size_t)items * (size_t)size);

}

static void

PyZlib_Free(voidpf ctx, void *ptr)

{

PyMem_RawFree(ptr);

}

static void

arrange_input_buffer(z_stream *zst, Py_ssize_t *remains)

{

zst->avail_in = (uInt)Py_MIN((size_t)*remains, UINT_MAX);

*remains -= zst->avail_in;

}

static PyObject *

zlib_compress_impl(PyObject *module, Py_buffer *data, int level, int wbits)

{

PyObject *return_value;

int flush;

z_stream zst;

_BlocksOutputBuffer buffer = {.writer = NULL};

zlibstate *state = get_zlib_state(module);

Byte *ibuf = data->buf;

Py_ssize_t ibuflen = data->len;

if (OutputBuffer_InitAndGrow(&buffer, -1, &zst.next_out, &zst.avail_out) < 0) {

goto error;

}

zst.opaque = NULL;

zst.zalloc = PyZlib_Malloc;

zst.zfree = PyZlib_Free;

zst.next_in = ibuf;

int err = deflateInit2(&zst, level, DEFLATED, wbits, DEF_MEM_LEVEL,

Z_DEFAULT_STRATEGY);

switch (err) {

case Z_OK:

break;

case Z_MEM_ERROR:

PyErr_SetString(PyExc_MemoryError,

"Out of memory while compressing data");

goto error;

case Z_STREAM_ERROR:

PyErr_SetString(state->ZlibError, "Bad compression level");

goto error;

default:

deflateEnd(&zst);

zlib_error(state, zst, err, "while compressing data");

goto error;

}

do {

arrange_input_buffer(&zst, &ibuflen);

flush = ibuflen == 0 ? Z_FINISH : Z_NO_FLUSH;

do {

if (zst.avail_out == 0) {

if (OutputBuffer_Grow(&buffer, &zst.next_out, &zst.avail_out) < 0) {

deflateEnd(&zst);

goto error;

}

}

Py_BEGIN_ALLOW_THREADS

err = deflate(&zst, flush);

Py_END_ALLOW_THREADS

if (err == Z_STREAM_ERROR) {

deflateEnd(&zst);

zlib_error(state, zst, err, "while compressing data");

goto error;

}

} while (zst.avail_out == 0);

assert(zst.avail_in == 0);

} while (flush != Z_FINISH);

assert(err == Z_STREAM_END);

err = deflateEnd(&zst);

if (err == Z_OK) {

return_value = OutputBuffer_Finish(&buffer, zst.avail_out);

if (return_value == NULL) {

goto error;

}

return return_value;

}

else

zlib_error(state, zst, err, "while finishing compression");

error:

OutputBuffer_OnError(&buffer);

return NULL;

}

static PyObject *

zlib_decompress_impl(PyObject *module, Py_buffer *data, int wbits,

Py_ssize_t bufsize)

{

PyObject *return_value;

Byte *ibuf;

Py_ssize_t ibuflen;

int err, flush;

z_stream zst;

_BlocksOutputBuffer buffer = {.writer = NULL};

_Uint32Window window; // output buffer's UINT32_MAX sliding window

zlibstate *state = get_zlib_state(module);

if (bufsize == 0) {

bufsize = 1;

}

if (OutputBuffer_WindowInitWithSize(&buffer, &window, bufsize,

&zst.next_out, &zst.avail_out) < 0) {

goto error;

}

ibuf = data->buf;

ibuflen = data->len;

zst.opaque = NULL;

zst.zalloc = PyZlib_Malloc;

zst.zfree = PyZlib_Free;

zst.avail_in = 0;

zst.next_in = ibuf;

err = inflateInit2(&zst, wbits);

switch (err) {

case Z_OK:

break;

case Z_MEM_ERROR:

PyErr_SetString(PyExc_MemoryError,

"Out of memory while decompressing data");

goto error;

default:

inflateEnd(&zst);

zlib_error(state, zst, err, "while preparing to decompress data");

goto error;

}

do {

arrange_input_buffer(&zst, &ibuflen);

flush = ibuflen == 0 ? Z_FINISH : Z_NO_FLUSH;

do {

if (zst.avail_out == 0) {

if (OutputBuffer_WindowGrow(&buffer, &window,

&zst.next_out, &zst.avail_out) < 0) {

inflateEnd(&zst);

goto error;

}

}

Py_BEGIN_ALLOW_THREADS

err = inflate(&zst, flush);

Py_END_ALLOW_THREADS

switch (err) {

case Z_OK: _Py_FALLTHROUGH;

case Z_BUF_ERROR: _Py_FALLTHROUGH;

case Z_STREAM_END:

break;

case Z_MEM_ERROR:

inflateEnd(&zst);

PyErr_SetString(PyExc_MemoryError,

"Out of memory while decompressing data");

goto error;

default:

inflateEnd(&zst);

zlib_error(state, zst, err, "while decompressing data");

goto error;

}

} while (zst.avail_out == 0);

} while (err != Z_STREAM_END && ibuflen != 0);

if (err != Z_STREAM_END) {

inflateEnd(&zst);

zlib_error(state, zst, err, "while decompressing data");

goto error;

}

err = inflateEnd(&zst);

if (err != Z_OK) {

zlib_error(state, zst, err, "while finishing decompression");

goto error;

}

return_value = OutputBuffer_WindowFinish(&buffer, &window, zst.avail_out);

if (return_value != NULL) {

return return_value;

}

error:

OutputBuffer_WindowOnError(&buffer, &window);

return NULL;

}

static PyObject *

zlib_compressobj_impl(PyObject *module, int level, int method, int wbits,

int memLevel, int strategy, Py_buffer *zdict)

{

zlibstate *state = get_zlib_state(module);

if (zdict->buf != NULL && (size_t)zdict->len > UINT_MAX) {

PyErr_SetString(PyExc_OverflowError,

"zdict length does not fit in an unsigned int");

return NULL;

}

compobject *self = newcompobject(state->Comptype);

if (self == NULL)

goto error;

self->zst.opaque = NULL;

self->zst.zalloc = PyZlib_Malloc;

self->zst.zfree = PyZlib_Free;

self->zst.next_in = NULL;

self->zst.avail_in = 0;

int err = deflateInit2(&self->zst, level, method, wbits, memLevel, strategy);

switch (err) {

case Z_OK:

self->is_initialised = 1;

if (zdict->buf == NULL) {

goto success;

} else {

err = deflateSetDictionary(&self->zst,

zdict->buf, (unsigned int)zdict->len);

switch (err) {

case Z_OK:

goto success;

case Z_STREAM_ERROR:

PyErr_SetString(PyExc_ValueError, "Invalid dictionary");

goto error;

default:

PyErr_SetString(PyExc_ValueError, "deflateSetDictionary()");

goto error;

}

}

case Z_MEM_ERROR:

PyErr_SetString(PyExc_MemoryError,

"Can't allocate memory for compression object");

goto error;

case Z_STREAM_ERROR:

PyErr_SetString(PyExc_ValueError, "Invalid initialization option");

goto error;

default:

zlib_error(state, self->zst, err, "while creating compression object");

goto error;

}

error:

Py_CLEAR(self);

success:

return (PyObject *)self;

}

static int

set_inflate_zdict(zlibstate *state, compobject *self)

{

Py_buffer zdict_buf;

if (PyObject_GetBuffer(self->zdict, &zdict_buf, PyBUF_SIMPLE) == -1) {

return -1;

}

if ((size_t)zdict_buf.len > UINT_MAX) {

PyErr_SetString(PyExc_OverflowError,

"zdict length does not fit in an unsigned int");

PyBuffer_Release(&zdict_buf);

return -1;

}

int err;

err = inflateSetDictionary(&self->zst,

zdict_buf.buf, (unsigned int)zdict_buf.len);

PyBuffer_Release(&zdict_buf);

if (err != Z_OK) {

zlib_error(state, self->zst, err, "while setting zdict");

return -1;

}

return 0;

}

static PyObject *

zlib_decompressobj_impl(PyObject *module, int wbits, PyObject *zdict)

{

zlibstate *state = get_zlib_state(module);

if (zdict != NULL && !PyObject_CheckBuffer(zdict)) {

PyErr_SetString(PyExc_TypeError,

"zdict argument must support the buffer protocol");

return NULL;

}

compobject *self = newcompobject(state->Decomptype);

if (self == NULL)

return 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) {

self->zdict = Py_NewRef(zdict);

}

int err = inflateInit2(&self->zst, wbits);

switch (err) {

case Z_OK:

self->is_initialised = 1;

if (self->zdict != NULL && wbits < 0) {

if (set_inflate_zdict(state, self) < 0) {

Py_DECREF(self);

return NULL;

}

}

return (PyObject *)self;

case Z_STREAM_ERROR:

Py_DECREF(self);

PyErr_SetString(PyExc_ValueError, "Invalid initialization option");

return NULL;

case Z_MEM_ERROR:

Py_DECREF(self);

PyErr_SetString(PyExc_MemoryError,

"Can't allocate memory for decompression object");

return NULL;

default:

zlib_error(state, self->zst, err, "while creating decompression object");

Py_DECREF(self);

return NULL;

}

}

static void

compobject_dealloc_impl(PyObject *op, int (*dealloc)(z_streamp))

{

PyTypeObject *type = Py_TYPE(op);

PyObject_GC_UnTrack(op);

compobject *self = _compobject_CAST(op);

assert(!PyMutex_IsLocked(&self->mutex));

if (self->is_initialised) {

(void)dealloc(&self->zst);

}

Py_XDECREF(self->unused_data);

Py_XDECREF(self->unconsumed_tail);

Py_XDECREF(self->zdict);

type->tp_free(self);

Py_DECREF(type);

}

static int

compobject_traverse(PyObject *op, visitproc visit, void *arg)

{

compobject *self = _compobject_CAST(op);

Py_VISIT(Py_TYPE(op));

Py_VISIT(self->zdict);

return 0;

}

static void

Comp_dealloc(PyObject *op)

{

compobject_dealloc_impl(op, &deflateEnd);

}

static void

Decomp_dealloc(PyObject *op)

{

compobject_dealloc_impl(op, &inflateEnd);

}

static PyObject *

zlib_Compress_compress_impl(compobject *self, PyTypeObject *cls,

Py_buffer *data)

{

PyObject *return_value;

int err;

_BlocksOutputBuffer buffer = {.writer = NULL};

zlibstate *state = PyType_GetModuleState(cls);

PyMutex_Lock(&self->mutex);

self->zst.next_in = data->buf;

Py_ssize_t ibuflen = data->len;

if (OutputBuffer_InitAndGrow(&buffer, -1, &self->zst.next_out, &self->zst.avail_out) < 0) {

goto error;

}

do {

arrange_input_buffer(&self->zst, &ibuflen);

do {

if (self->zst.avail_out == 0) {

if (OutputBuffer_Grow(&buffer, &self->zst.next_out, &self->zst.avail_out) < 0) {

goto error;

}

}

Py_BEGIN_ALLOW_THREADS

err = deflate(&self->zst, Z_NO_FLUSH);

Py_END_ALLOW_THREADS

if (err == Z_STREAM_ERROR) {

zlib_error(state, self->zst, err, "while compressing data");

goto error;

}

} while (self->zst.avail_out == 0);

assert(self->zst.avail_in == 0);

} while (ibuflen != 0);

return_value = OutputBuffer_Finish(&buffer, self->zst.avail_out);

if (return_value != NULL) {

goto success;

}

error:

OutputBuffer_OnError(&buffer);

return_value = NULL;

success:

PyMutex_Unlock(&self->mutex);

return return_value;

}

static int

save_unconsumed_input(compobject *self, Py_buffer *data, int err)

{

if (err == Z_STREAM_END) {

/* The end of the compressed data has been reached. Store the leftover

if (self->zst.avail_in > 0) {

Py_ssize_t old_size = PyBytes_GET_SIZE(self->unused_data);

Py_ssize_t left_size;

left_size = (Byte *)data->buf + data->len - self->zst.next_in;

if (left_size > (PY_SSIZE_T_MAX - old_size)) {

PyErr_NoMemory();

return -1;

}

PyBytesWriter *writer = PyBytesWriter_Create(old_size + left_size);

if (writer == NULL) {

return -1;

}

char *new_data = PyBytesWriter_GetData(writer);

memcpy(new_data, PyBytes_AS_STRING(self->unused_data), old_size);

memcpy(new_data + old_size, self->zst.next_in, left_size);

PyObject *new_unused_data = PyBytesWriter_Finish(writer);

if (new_unused_data == NULL) {

return -1;

}

Py_SETREF(self->unused_data, new_unused_data);

self->zst.avail_in = 0;

}

}

if (self->zst.avail_in > 0 || PyBytes_GET_SIZE(self->unconsumed_tail)) {

/* This code handles two distinct cases:

Py_ssize_t left_size = (Byte *)data->buf + data->len - self->zst.next_in;

PyObject *new_data = PyBytes_FromStringAndSize(

(char *)self->zst.next_in, left_size);

if (new_data == NULL)

return -1;

Py_SETREF(self->unconsumed_tail, new_data);

}

return 0;

}

static PyObject *

zlib_Decompress_decompress_impl(compobject *self, PyTypeObject *cls,

Py_buffer *data, Py_ssize_t max_length)

{

int err = Z_OK;

Py_ssize_t ibuflen;

PyObject *return_value;

_BlocksOutputBuffer buffer = {.writer = NULL};

PyObject *module = PyType_GetModule(cls);

if (module == NULL)

return NULL;

zlibstate *state = get_zlib_state(module);

if (max_length == 0) {

max_length = -1;

}

PyMutex_Lock(&self->mutex);

self->zst.next_in = data->buf;

ibuflen = data->len;

if (OutputBuffer_InitAndGrow(&buffer, max_length, &self->zst.next_out, &self->zst.avail_out) < 0) {

goto abort;

}

do {

arrange_input_buffer(&self->zst, &ibuflen);

do {

if (self->zst.avail_out == 0) {

if (OutputBuffer_GetDataSize(&buffer, self->zst.avail_out) == max_length) {

goto save;

}

if (OutputBuffer_Grow(&buffer, &self->zst.next_out, &self->zst.avail_out) < 0) {

goto abort;

}

}

Py_BEGIN_ALLOW_THREADS

err = inflate(&self->zst, Z_SYNC_FLUSH);

Py_END_ALLOW_THREADS

switch (err) {

case Z_OK: _Py_FALLTHROUGH;

case Z_BUF_ERROR: _Py_FALLTHROUGH;

case Z_STREAM_END:

break;

default:

if (err == Z_NEED_DICT && self->zdict != NULL) {

if (set_inflate_zdict(state, self) < 0) {

goto abort;

}

else

break;

}

goto save;

}

} while (self->zst.avail_out == 0 || err == Z_NEED_DICT);

} while (err != Z_STREAM_END && ibuflen != 0);

save:

if (save_unconsumed_input(self, data, err) < 0)

goto abort;

if (err == Z_STREAM_END) {

/* This is the logical place to call inflateEnd, but the old behaviour

FT_ATOMIC_STORE_CHAR_RELAXED(self->eof, 1);

} else if (err != Z_OK && err != Z_BUF_ERROR) {

/* We will only get Z_BUF_ERROR if the output buffer was full

zlib_error(state, self->zst, err, "while decompressing data");

goto abort;

}

return_value = OutputBuffer_Finish(&buffer, self->zst.avail_out);

if (return_value != NULL) {

goto success;

}

abort:

OutputBuffer_OnError(&buffer);

return_value = NULL;

success:

PyMutex_Unlock(&self->mutex);

return return_value;

}

static PyObject *

zlib_Compress_flush_impl(compobject *self, PyTypeObject *cls, int mode)

{

int err;

PyObject *return_value;

_BlocksOutputBuffer buffer = {.writer = NULL};

zlibstate *state = PyType_GetModuleState(cls);

/* Flushing with Z_NO_FLUSH is a no-op, so there's no point in

if (mode == Z_NO_FLUSH) {