1/* zlibmodule.c -- gzip-compatible data compression */
2/* See http://zlib.net/ */
3
4/* Windows users:  read Python's PCbuild\readme.txt */
5
6#define PY_SSIZE_T_CLEAN
7
8#include "Python.h"
9#include "zlib.h"
10
11#ifdef WITH_THREAD
12#include "pythread.h"
13
14/* #defs ripped off from _tkinter.c, even though the situation here is much
15   simpler, because we don't have to worry about waiting for Tcl
16   events!  And, since zlib itself is threadsafe, we don't need to worry
17   about re-entering zlib functions.
18
19   N.B.
20
21   Since ENTER_ZLIB and LEAVE_ZLIB only need to be called on functions
22   that modify the components of preexisting de/compress objects, it
23   could prove to be a performance gain on multiprocessor machines if
24   there was an de/compress object-specific lock.  However, for the
25   moment the ENTER_ZLIB and LEAVE_ZLIB calls are global for ALL
26   de/compress objects.
27 */
28
29static PyThread_type_lock zlib_lock = NULL; /* initialized on module load */
30
31#define ENTER_ZLIB \
32        Py_BEGIN_ALLOW_THREADS \
33        PyThread_acquire_lock(zlib_lock, 1); \
34        Py_END_ALLOW_THREADS
35
36#define LEAVE_ZLIB \
37        PyThread_release_lock(zlib_lock);
38
39#else
40
41#define ENTER_ZLIB
42#define LEAVE_ZLIB
43
44#endif
45
46/* The following parameters are copied from zutil.h, version 0.95 */
47#define DEFLATED   8
48#if MAX_MEM_LEVEL >= 8
49#  define DEF_MEM_LEVEL 8
50#else
51#  define DEF_MEM_LEVEL  MAX_MEM_LEVEL
52#endif
53#define DEF_WBITS MAX_WBITS
54
55/* The output buffer will be increased in chunks of DEFAULTALLOC bytes. */
56#define DEFAULTALLOC (16*1024)
57#define PyInit_zlib initzlib
58
59static PyTypeObject Comptype;
60static PyTypeObject Decomptype;
61
62static PyObject *ZlibError;
63
64typedef struct
65{
66    PyObject_HEAD
67    z_stream zst;
68    PyObject *unused_data;
69    PyObject *unconsumed_tail;
70    int is_initialised;
71} compobject;
72
73static void
74zlib_error(z_stream zst, int err, char *msg)
75{
76    const char *zmsg = Z_NULL;
77    /* In case of a version mismatch, zst.msg won't be initialized.
78       Check for this case first, before looking at zst.msg. */
79    if (err == Z_VERSION_ERROR)
80        zmsg = "library version mismatch";
81    if (zmsg == Z_NULL)
82        zmsg = zst.msg;
83    if (zmsg == Z_NULL) {
84        switch (err) {
85        case Z_BUF_ERROR:
86            zmsg = "incomplete or truncated stream";
87            break;
88        case Z_STREAM_ERROR:
89            zmsg = "inconsistent stream state";
90            break;
91        case Z_DATA_ERROR:
92            zmsg = "invalid input data";
93            break;
94        }
95    }
96    if (zmsg == Z_NULL)
97        PyErr_Format(ZlibError, "Error %d %s", err, msg);
98    else
99        PyErr_Format(ZlibError, "Error %d %s: %.200s", err, msg, zmsg);
100}
101
102PyDoc_STRVAR(compressobj__doc__,
103"compressobj([level]) -- Return a compressor object.\n"
104"\n"
105"Optional arg level is the compression level, in 0-9 or -1.");
106
107PyDoc_STRVAR(decompressobj__doc__,
108"decompressobj([wbits]) -- Return a decompressor object.\n"
109"\n"
110"Optional arg wbits indicates the window buffer size and container format.");
111
112static compobject *
113newcompobject(PyTypeObject *type)
114{
115    compobject *self;
116    self = PyObject_New(compobject, type);
117    if (self == NULL)
118        return NULL;
119    self->is_initialised = 0;
120    self->unused_data = PyString_FromString("");
121    if (self->unused_data == NULL) {
122        Py_DECREF(self);
123        return NULL;
124    }
125    self->unconsumed_tail = PyString_FromString("");
126    if (self->unconsumed_tail == NULL) {
127        Py_DECREF(self);
128        return NULL;
129    }
130    return self;
131}
132
133static void
134arrange_input_buffer(z_stream *zst, Py_ssize_t *remains)
135{
136    if ((size_t)*remains > UINT_MAX) {
137        zst->avail_in = UINT_MAX;
138    } else {
139        zst->avail_in = *remains;
140    }
141    *remains -= zst->avail_in;
142}
143
144static Py_ssize_t
145arrange_output_buffer_with_maximum(z_stream *zst, PyObject **buffer,
146                                   Py_ssize_t length,
147                                   Py_ssize_t max_length)
148{
149    Py_ssize_t occupied;
150
151    if (*buffer == NULL) {
152        if (!(*buffer = PyBytes_FromStringAndSize(NULL, length)))
153            return -1;
154        occupied = 0;
155    }
156    else {
157        occupied = zst->next_out - (Byte *)PyBytes_AS_STRING(*buffer);
158
159        if (length == occupied) {
160            Py_ssize_t new_length;
161            assert(length <= max_length);
162            /* can not scale the buffer over max_length */
163            if (length == max_length)
164                return -2;
165            if (length <= (max_length >> 1))
166                new_length = length << 1;
167            else
168                new_length = max_length;
169            if (_PyBytes_Resize(buffer, new_length) < 0)
170                return -1;
171            length = new_length;
172        }
173    }
174
175    if ((size_t)(length - occupied) > UINT_MAX) {
176        zst->avail_out = UINT_MAX;
177    } else {
178        zst->avail_out = length - occupied;
179    }
180    zst->next_out = (Byte *)PyBytes_AS_STRING(*buffer) + occupied;
181
182    return length;
183}
184
185static Py_ssize_t
186arrange_output_buffer(z_stream *zst, PyObject **buffer, Py_ssize_t length)
187{
188    Py_ssize_t ret;
189
190    ret = arrange_output_buffer_with_maximum(zst, buffer, length,
191                                             PY_SSIZE_T_MAX);
192    if (ret == -2)
193        PyErr_NoMemory();
194
195    return ret;
196}
197
198PyDoc_STRVAR(compress__doc__,
199"compress(string[, level]) -- Returned compressed string.\n"
200"\n"
201"Optional arg level is the compression level, in 0-9.");
202
203static PyObject *
204PyZlib_compress(PyObject *self, PyObject *args)
205{
206    PyObject *RetVal = NULL;
207    Byte *ibuf;
208    Py_ssize_t ibuflen, obuflen = DEFAULTALLOC;
209    int level=Z_DEFAULT_COMPRESSION;
210    int err, flush;
211    z_stream zst;
212
213    /* require Python string object, optional 'level' arg */
214    if (!PyArg_ParseTuple(args, "s#|i:compress", &ibuf, &ibuflen, &level))
215        return NULL;
216
217    zst.zalloc = (alloc_func)NULL;
218    zst.zfree = (free_func)Z_NULL;
219    zst.next_in = ibuf;
220    err = deflateInit(&zst, level);
221
222    switch (err) {
223    case Z_OK:
224        break;
225    case Z_MEM_ERROR:
226        PyErr_SetString(PyExc_MemoryError,
227                        "Out of memory while compressing data");
228        goto error;
229    case Z_STREAM_ERROR:
230        PyErr_SetString(ZlibError, "Bad compression level");
231        goto error;
232    default:
233        deflateEnd(&zst);
234        zlib_error(zst, err, "while compressing data");
235        goto error;
236    }
237
238    do {
239        arrange_input_buffer(&zst, &ibuflen);
240        flush = ibuflen == 0 ? Z_FINISH : Z_NO_FLUSH;
241
242        do {
243            obuflen = arrange_output_buffer(&zst, &RetVal, obuflen);
244            if (obuflen < 0) {
245                deflateEnd(&zst);
246                goto error;
247            }
248
249            Py_BEGIN_ALLOW_THREADS
250            err = deflate(&zst, flush);
251            Py_END_ALLOW_THREADS
252
253            if (err == Z_STREAM_ERROR) {
254                deflateEnd(&zst);
255                zlib_error(zst, err, "while compressing data");
256                goto error;
257            }
258
259        } while (zst.avail_out == 0);
260        assert(zst.avail_in == 0);
261
262    } while (flush != Z_FINISH);
263    assert(err == Z_STREAM_END);
264
265    err = deflateEnd(&zst);
266    if (err == Z_OK) {
267        if (_PyBytes_Resize(&RetVal, zst.next_out -
268                            (Byte *)PyBytes_AS_STRING(RetVal)) < 0)
269            goto error;
270        return RetVal;
271    }
272    else
273        zlib_error(zst, err, "while finishing compression");
274 error:
275    Py_XDECREF(RetVal);
276    return NULL;
277}
278
279PyDoc_STRVAR(decompress__doc__,
280"decompress(string[, wbits[, bufsize]]) -- Return decompressed string.\n"
281"\n"
282"Optional arg wbits indicates the window buffer size and container format.\n"
283"Optional arg bufsize is the initial output buffer size.");
284
285static PyObject *
286PyZlib_decompress(PyObject *self, PyObject *args)
287{
288    PyObject *RetVal = NULL;
289    Byte *ibuf;
290    Py_ssize_t ibuflen;
291    int err, flush;
292    int wsize=DEF_WBITS;
293    Py_ssize_t r_strlen=DEFAULTALLOC;
294    z_stream zst;
295
296    if (!PyArg_ParseTuple(args, "s#|in:decompress",
297                          &ibuf, &ibuflen, &wsize, &r_strlen))
298        return NULL;
299
300    if (r_strlen <= 0) {
301        r_strlen = 1;
302    }
303
304    zst.zalloc = (alloc_func)NULL;
305    zst.zfree = (free_func)Z_NULL;
306    zst.avail_in = 0;
307    zst.next_in = ibuf;
308    err = inflateInit2(&zst, wsize);
309
310    switch (err) {
311    case Z_OK:
312        break;
313    case Z_MEM_ERROR:
314        PyErr_SetString(PyExc_MemoryError,
315                        "Out of memory while decompressing data");
316        goto error;
317    default:
318        inflateEnd(&zst);
319        zlib_error(zst, err, "while preparing to decompress data");
320        goto error;
321    }
322
323    do {
324        arrange_input_buffer(&zst, &ibuflen);
325        flush = ibuflen == 0 ? Z_FINISH : Z_NO_FLUSH;
326
327        do {
328            r_strlen = arrange_output_buffer(&zst, &RetVal, r_strlen);
329            if (r_strlen < 0) {
330                inflateEnd(&zst);
331                goto error;
332            }
333
334            Py_BEGIN_ALLOW_THREADS
335            err = inflate(&zst, flush);
336            Py_END_ALLOW_THREADS
337
338            switch (err) {
339            case Z_OK:            /* fall through */
340            case Z_BUF_ERROR:     /* fall through */
341            case Z_STREAM_END:
342                break;
343            case Z_MEM_ERROR:
344                inflateEnd(&zst);
345                PyErr_SetString(PyExc_MemoryError,
346                                "Out of memory while decompressing data");
347                goto error;
348            default:
349                inflateEnd(&zst);
350                zlib_error(zst, err, "while decompressing data");
351                goto error;
352            }
353
354        } while (zst.avail_out == 0);
355
356    } while (err != Z_STREAM_END && ibuflen != 0);
357
358
359    if (err != Z_STREAM_END) {
360        inflateEnd(&zst);
361        zlib_error(zst, err, "while decompressing data");
362        goto error;
363    }
364
365    err = inflateEnd(&zst);
366    if (err != Z_OK) {
367        zlib_error(zst, err, "while finishing data decompression");
368        goto error;
369    }
370
371    _PyString_Resize(&RetVal, zst.next_out -
372                        (Byte *)PyBytes_AS_STRING(RetVal));
373    return RetVal;
374
375 error:
376    Py_XDECREF(RetVal);
377    return NULL;
378}
379
380static PyObject *
381PyZlib_compressobj(PyObject *selfptr, PyObject *args)
382{
383    compobject *self;
384    int level=Z_DEFAULT_COMPRESSION, method=DEFLATED;
385    int wbits=MAX_WBITS, memLevel=DEF_MEM_LEVEL, strategy=0, err;
386
387    if (!PyArg_ParseTuple(args, "|iiiii:compressobj", &level, &method, &wbits,
388                          &memLevel, &strategy))
389        return NULL;
390
391    self = newcompobject(&Comptype);
392    if (self == NULL)
393        return(NULL);
394    self->zst.zalloc = (alloc_func)NULL;
395    self->zst.zfree = (free_func)Z_NULL;
396    self->zst.next_in = NULL;
397    self->zst.avail_in = 0;
398    err = deflateInit2(&self->zst, level, method, wbits, memLevel, strategy);
399    switch (err) {
400    case Z_OK:
401        self->is_initialised = 1;
402        return (PyObject *)self;
403    case Z_MEM_ERROR:
404        Py_DECREF(self);
405        PyErr_SetString(PyExc_MemoryError,
406                        "Can't allocate memory for compression object");
407        return NULL;
408    case Z_STREAM_ERROR:
409        Py_DECREF(self);
410        PyErr_SetString(PyExc_ValueError, "Invalid initialization option");
411        return NULL;
412    default:
413        zlib_error(self->zst, err, "while creating compression object");
414        Py_DECREF(self);
415        return NULL;
416    }
417}
418
419static PyObject *
420PyZlib_decompressobj(PyObject *selfptr, PyObject *args)
421{
422    int wbits=DEF_WBITS, err;
423    compobject *self;
424    if (!PyArg_ParseTuple(args, "|i:decompressobj", &wbits))
425        return NULL;
426
427    self = newcompobject(&Decomptype);
428    if (self == NULL)
429        return NULL;
430    self->zst.zalloc = (alloc_func)NULL;
431    self->zst.zfree = (free_func)Z_NULL;
432    self->zst.next_in = NULL;
433    self->zst.avail_in = 0;
434    err = inflateInit2(&self->zst, wbits);
435    switch (err) {
436    case Z_OK:
437        self->is_initialised = 1;
438        return (PyObject *)self;
439    case Z_STREAM_ERROR:
440        Py_DECREF(self);
441        PyErr_SetString(PyExc_ValueError, "Invalid initialization option");
442        return NULL;
443    case Z_MEM_ERROR:
444        Py_DECREF(self);
445        PyErr_SetString(PyExc_MemoryError,
446                        "Can't allocate memory for decompression object");
447        return NULL;
448    default:
449        zlib_error(self->zst, err, "while creating decompression object");
450        Py_DECREF(self);
451        return NULL;
452    }
453}
454
455static void
456Comp_dealloc(compobject *self)
457{
458    if (self->is_initialised)
459        deflateEnd(&self->zst);
460    Py_XDECREF(self->unused_data);
461    Py_XDECREF(self->unconsumed_tail);
462    PyObject_Del(self);
463}
464
465static void
466Decomp_dealloc(compobject *self)
467{
468    if (self->is_initialised)
469        inflateEnd(&self->zst);
470    Py_XDECREF(self->unused_data);
471    Py_XDECREF(self->unconsumed_tail);
472    PyObject_Del(self);
473}
474
475PyDoc_STRVAR(comp_compress__doc__,
476"compress(data) -- Return a string containing data compressed.\n"
477"\n"
478"After calling this function, some of the input data may still\n"
479"be stored in internal buffers for later processing.\n"
480"Call the flush() method to clear these buffers.");
481
482
483static PyObject *
484PyZlib_objcompress(compobject *self, PyObject *args)
485{
486    PyObject *RetVal = NULL;
487    Py_ssize_t ibuflen, obuflen = DEFAULTALLOC;
488    int err;
489
490    if (!PyArg_ParseTuple(args, "s#:compress", &self->zst.next_in, &ibuflen))
491        return NULL;
492
493    ENTER_ZLIB
494
495    do {
496        arrange_input_buffer(&self->zst, &ibuflen);
497
498        do {
499            obuflen = arrange_output_buffer(&self->zst, &RetVal, obuflen);
500            if (obuflen < 0)
501                goto error;
502
503            Py_BEGIN_ALLOW_THREADS
504            err = deflate(&self->zst, Z_NO_FLUSH);
505            Py_END_ALLOW_THREADS
506
507            if (err == Z_STREAM_ERROR) {
508                zlib_error(self->zst, err, "while compressing data");
509                goto error;
510            }
511
512        } while (self->zst.avail_out == 0);
513        assert(self->zst.avail_in == 0);
514
515    } while (ibuflen != 0);
516
517    _PyString_Resize(&RetVal, self->zst.next_out -
518                        (Byte *)PyBytes_AS_STRING(RetVal));
519    goto success;
520
521 error:
522    Py_CLEAR(RetVal);
523 success:
524    LEAVE_ZLIB
525    return RetVal;
526}
527
528/* Helper for objdecompress() and unflush(). Saves any unconsumed input data in
529   self->unused_data or self->unconsumed_tail, as appropriate. */
530static int
531save_unconsumed_input(compobject *self, Byte *input, Py_ssize_t inplen,
532    int err)
533{
534    if (err == Z_STREAM_END) {
535        /* The end of the compressed data has been reached. Store the leftover
536           input data in self->unused_data. */
537        if (self->zst.avail_in > 0) {
538            Py_ssize_t old_size = PyString_GET_SIZE(self->unused_data);
539            Py_ssize_t new_size, left_size;
540            PyObject *new_data;
541            left_size = input + inplen - self->zst.next_in;
542            if (left_size > (PY_SSIZE_T_MAX - old_size)) {
543                PyErr_NoMemory();
544                return -1;
545            }
546            new_size = old_size + left_size;
547            new_data = PyString_FromStringAndSize(NULL, new_size);
548            if (new_data == NULL)
549                return -1;
550            Py_MEMCPY(PyString_AS_STRING(new_data),
551                      PyString_AS_STRING(self->unused_data), old_size);
552            Py_MEMCPY(PyString_AS_STRING(new_data) + old_size,
553                      self->zst.next_in, left_size);
554            Py_SETREF(self->unused_data, new_data);
555            self->zst.avail_in = 0;
556        }
557    }
558
559    if (self->zst.avail_in > 0 || PyString_GET_SIZE(self->unconsumed_tail)) {
560        /* This code handles two distinct cases:
561           1. Output limit was reached. Save leftover input in unconsumed_tail.
562           2. All input data was consumed. Clear unconsumed_tail. */
563        Py_ssize_t left_size = input + inplen - self->zst.next_in;
564        PyObject *new_data = PyString_FromStringAndSize(
565                (char *)self->zst.next_in, left_size);
566        if (new_data == NULL)
567            return -1;
568        Py_SETREF(self->unconsumed_tail, new_data);
569    }
570
571    return 0;
572}
573
574PyDoc_STRVAR(decomp_decompress__doc__,
575"decompress(data, max_length) -- Return a string containing the decompressed\n"
576"version of the data.\n"
577"\n"
578"After calling this function, some of the input data may still be stored in\n"
579"internal buffers for later processing.\n"
580"Call the flush() method to clear these buffers.\n"
581"If the max_length parameter is specified then the return value will be\n"
582"no longer than max_length.  Unconsumed input data will be stored in\n"
583"the unconsumed_tail attribute.");
584
585static PyObject *
586PyZlib_objdecompress(compobject *self, PyObject *args)
587{
588    int err = Z_OK;
589    Py_ssize_t inplen, max_length = 0;
590    Py_ssize_t ibuflen, obuflen = DEFAULTALLOC, hard_limit;
591    PyObject *RetVal = NULL;
592    Byte *input;
593
594    if (!PyArg_ParseTuple(args, "s#|n:decompress", &input,
595                          &inplen, &max_length))
596        return NULL;
597    if (max_length < 0) {
598        PyErr_SetString(PyExc_ValueError,
599                        "max_length must be greater than zero");
600        return NULL;
601    } else if (max_length == 0)
602        hard_limit = PY_SSIZE_T_MAX;
603    else
604        hard_limit = max_length;
605
606    self->zst.next_in = input;
607    ibuflen = inplen;
608
609    /* limit amount of data allocated to max_length */
610    if (max_length && obuflen > max_length)
611        obuflen = max_length;
612
613    ENTER_ZLIB
614
615    do {
616        arrange_input_buffer(&self->zst, &ibuflen);
617
618        do {
619            obuflen = arrange_output_buffer_with_maximum(&self->zst, &RetVal,
620                                                         obuflen, hard_limit);
621            if (obuflen == -2) {
622                if (max_length > 0) {
623                    goto save;
624                }
625                PyErr_NoMemory();
626            }
627            if (obuflen < 0) {
628                goto abort;
629            }
630
631            Py_BEGIN_ALLOW_THREADS
632            err = inflate(&self->zst, Z_SYNC_FLUSH);
633            Py_END_ALLOW_THREADS
634
635            switch (err) {
636            case Z_OK:            /* fall through */
637            case Z_BUF_ERROR:     /* fall through */
638            case Z_STREAM_END:
639                break;
640            default:
641                goto save;
642            }
643
644        } while (self->zst.avail_out == 0);
645
646    } while (err != Z_STREAM_END && ibuflen != 0);
647
648 save:
649    if (save_unconsumed_input(self, input, inplen, err) < 0)
650        goto abort;
651
652    /* This is the logical place to call inflateEnd, but the old behaviour of
653       only calling it on flush() is preserved. */
654
655    if (err != Z_STREAM_END && err != Z_OK && err != Z_BUF_ERROR) {
656        /* We will only get Z_BUF_ERROR if the output buffer was full
657           but there wasn't more output when we tried again, so it is
658           not an error condition.
659        */
660        zlib_error(self->zst, err, "while decompressing");
661        goto abort;
662    }
663
664    _PyString_Resize(&RetVal, self->zst.next_out -
665                        (Byte *)PyBytes_AS_STRING(RetVal));
666    goto success;
667
668 abort:
669    Py_CLEAR(RetVal);
670 success:
671    LEAVE_ZLIB
672
673    return RetVal;
674}
675
676PyDoc_STRVAR(comp_flush__doc__,
677"flush( [mode] ) -- Return a string containing any remaining compressed data.\n"
678"\n"
679"mode can be one of the constants Z_SYNC_FLUSH, Z_FULL_FLUSH, Z_FINISH; the\n"
680"default value used when mode is not specified is Z_FINISH.\n"
681"If mode == Z_FINISH, the compressor object can no longer be used after\n"
682"calling the flush() method.  Otherwise, more data can still be compressed.");
683
684static PyObject *
685PyZlib_flush(compobject *self, PyObject *args)
686{
687    int err;
688    Py_ssize_t length = DEFAULTALLOC;
689    PyObject *RetVal = NULL;
690    int flushmode = Z_FINISH;
691
692    if (!PyArg_ParseTuple(args, "|i:flush", &flushmode))
693        return NULL;
694
695    /* Flushing with Z_NO_FLUSH is a no-op, so there's no point in
696       doing any work at all; just return an empty string. */
697    if (flushmode == Z_NO_FLUSH) {
698        return PyString_FromStringAndSize(NULL, 0);
699    }
700
701    ENTER_ZLIB
702
703    self->zst.avail_in = 0;
704
705    do {
706        length = arrange_output_buffer(&self->zst, &RetVal, length);
707        if (length < 0) {
708            Py_CLEAR(RetVal);
709            goto error;
710        }
711
712        Py_BEGIN_ALLOW_THREADS
713        err = deflate(&self->zst, flushmode);
714        Py_END_ALLOW_THREADS
715
716        if (err == Z_STREAM_ERROR) {
717            zlib_error(self->zst, err, "while flushing");
718            Py_CLEAR(RetVal);
719            goto error;
720        }
721    } while (self->zst.avail_out == 0);
722    assert(self->zst.avail_in == 0);
723
724    /* If flushmode is Z_FINISH, we also have to call deflateEnd() to free
725       various data structures. Note we should only get Z_STREAM_END when
726       flushmode is Z_FINISH, but checking both for safety*/
727    if (err == Z_STREAM_END && flushmode == Z_FINISH) {
728        err = deflateEnd(&self->zst);
729        if (err != Z_OK) {
730            zlib_error(self->zst, err, "from deflateEnd()");
731            Py_CLEAR(RetVal);
732            goto error;
733        }
734        else
735            self->is_initialised = 0;
736
737        /* We will only get Z_BUF_ERROR if the output buffer was full
738           but there wasn't more output when we tried again, so it is
739           not an error condition.
740        */
741    } else if (err != Z_OK && err != Z_BUF_ERROR) {
742        zlib_error(self->zst, err, "while flushing");
743        Py_CLEAR(RetVal);
744        goto error;
745    }
746
747    _PyString_Resize(&RetVal, self->zst.next_out -
748                        (Byte *)PyBytes_AS_STRING(RetVal));
749
750 error:
751    LEAVE_ZLIB
752    return RetVal;
753}
754
755#ifdef HAVE_ZLIB_COPY
756PyDoc_STRVAR(comp_copy__doc__,
757"copy() -- Return a copy of the compression object.");
758
759static PyObject *
760PyZlib_copy(compobject *self)
761{
762    compobject *retval = NULL;
763    int err;
764
765    retval = newcompobject(&Comptype);
766    if (!retval) return NULL;
767
768    /* Copy the zstream state
769     * We use ENTER_ZLIB / LEAVE_ZLIB to make this thread-safe
770     */
771    ENTER_ZLIB
772    err = deflateCopy(&retval->zst, &self->zst);
773    switch (err) {
774    case Z_OK:
775        break;
776    case Z_STREAM_ERROR:
777        PyErr_SetString(PyExc_ValueError, "Inconsistent stream state");
778        goto error;
779    case Z_MEM_ERROR:
780        PyErr_SetString(PyExc_MemoryError,
781                        "Can't allocate memory for compression object");
782        goto error;
783    default:
784        zlib_error(self->zst, err, "while copying compression object");
785        goto error;
786    }
787
788    Py_INCREF(self->unused_data);
789    Py_XSETREF(retval->unused_data, self->unused_data);
790    Py_INCREF(self->unconsumed_tail);
791    Py_XSETREF(retval->unconsumed_tail, self->unconsumed_tail);
792
793    /* Mark it as being initialized */
794    retval->is_initialised = 1;
795
796    LEAVE_ZLIB
797    return (PyObject *)retval;
798
799error:
800    LEAVE_ZLIB
801    Py_XDECREF(retval);
802    return NULL;
803}
804
805PyDoc_STRVAR(decomp_copy__doc__,
806"copy() -- Return a copy of the decompression object.");
807
808static PyObject *
809PyZlib_uncopy(compobject *self)
810{
811    compobject *retval = NULL;
812    int err;
813
814    retval = newcompobject(&Decomptype);
815    if (!retval) return NULL;
816
817    /* Copy the zstream state
818     * We use ENTER_ZLIB / LEAVE_ZLIB to make this thread-safe
819     */
820    ENTER_ZLIB
821    err = inflateCopy(&retval->zst, &self->zst);
822    switch (err) {
823    case Z_OK:
824        break;
825    case Z_STREAM_ERROR:
826        PyErr_SetString(PyExc_ValueError, "Inconsistent stream state");
827        goto error;
828    case Z_MEM_ERROR:
829        PyErr_SetString(PyExc_MemoryError,
830                        "Can't allocate memory for decompression object");
831        goto error;
832    default:
833        zlib_error(self->zst, err, "while copying decompression object");
834        goto error;
835    }
836
837    Py_INCREF(self->unused_data);
838    Py_XSETREF(retval->unused_data, self->unused_data);
839    Py_INCREF(self->unconsumed_tail);
840    Py_XSETREF(retval->unconsumed_tail, self->unconsumed_tail);
841
842    /* Mark it as being initialized */
843    retval->is_initialised = 1;
844
845    LEAVE_ZLIB
846    return (PyObject *)retval;
847
848error:
849    LEAVE_ZLIB
850    Py_XDECREF(retval);
851    return NULL;
852}
853#endif
854
855PyDoc_STRVAR(decomp_flush__doc__,
856"flush( [length] ) -- Return a string containing any remaining\n"
857"decompressed data. length, if given, is the initial size of the\n"
858"output buffer.\n"
859"\n"
860"The decompressor object can no longer be used after this call.");
861
862static PyObject *
863PyZlib_unflush(compobject *self, PyObject *args)
864{
865    Py_ssize_t length = DEFAULTALLOC;
866    int err, flush;
867    PyObject *RetVal = NULL;
868    Py_ssize_t ibuflen;
869
870    if (!PyArg_ParseTuple(args, "|n:flush", &length))
871        return NULL;
872    if (length <= 0) {
873        PyErr_SetString(PyExc_ValueError, "length must be greater than zero");
874        return NULL;
875    }
876
877    ENTER_ZLIB
878
879    self->zst.next_in = (Byte *)PyString_AS_STRING(self->unconsumed_tail);
880    ibuflen = PyString_GET_SIZE(self->unconsumed_tail);
881
882    do {
883        arrange_input_buffer(&self->zst, &ibuflen);
884        flush = ibuflen == 0 ? Z_FINISH : Z_NO_FLUSH;
885
886        do {
887            length = arrange_output_buffer(&self->zst, &RetVal, length);
888            if (length < 0)
889                goto abort;
890
891            Py_BEGIN_ALLOW_THREADS
892            err = inflate(&self->zst, flush);
893            Py_END_ALLOW_THREADS
894
895            switch (err) {
896            case Z_OK:            /* fall through */
897            case Z_BUF_ERROR:     /* fall through */
898            case Z_STREAM_END:
899                break;
900            default:
901                goto save;
902            }
903
904        } while (self->zst.avail_out == 0);
905
906    } while (err != Z_STREAM_END && ibuflen != 0);
907
908 save:
909    if (save_unconsumed_input(self,
910            (Byte *)PyString_AS_STRING(self->unconsumed_tail),
911            PyString_GET_SIZE(self->unconsumed_tail), err) < 0)
912        goto abort;
913
914    /* If flushmode is Z_FINISH, we also have to call deflateEnd() to free
915       various data structures. Note we should only get Z_STREAM_END when
916       flushmode is Z_FINISH */
917    if (err == Z_STREAM_END) {
918        err = inflateEnd(&self->zst);
919        self->is_initialised = 0;
920        if (err != Z_OK) {
921            zlib_error(self->zst, err, "from inflateEnd()");
922            goto abort;
923        }
924    }
925
926    _PyString_Resize(&RetVal, self->zst.next_out -
927                        (Byte *)PyBytes_AS_STRING(RetVal));
928    goto success;
929
930 abort:
931    Py_CLEAR(RetVal);
932 success:
933    LEAVE_ZLIB
934    return RetVal;
935}
936
937static PyMethodDef comp_methods[] =
938{
939    {"compress", (binaryfunc)PyZlib_objcompress, METH_VARARGS,
940                 comp_compress__doc__},
941    {"flush", (binaryfunc)PyZlib_flush, METH_VARARGS,
942              comp_flush__doc__},
943#ifdef HAVE_ZLIB_COPY
944    {"copy",  (PyCFunction)PyZlib_copy, METH_NOARGS,
945              comp_copy__doc__},
946#endif
947    {NULL, NULL}
948};
949
950static PyMethodDef Decomp_methods[] =
951{
952    {"decompress", (binaryfunc)PyZlib_objdecompress, METH_VARARGS,
953                   decomp_decompress__doc__},
954    {"flush", (binaryfunc)PyZlib_unflush, METH_VARARGS,
955              decomp_flush__doc__},
956#ifdef HAVE_ZLIB_COPY
957    {"copy",  (PyCFunction)PyZlib_uncopy, METH_NOARGS,
958              decomp_copy__doc__},
959#endif
960    {NULL, NULL}
961};
962
963static PyObject *
964Comp_getattr(compobject *self, char *name)
965{
966  /* No ENTER/LEAVE_ZLIB is necessary because this fn doesn't touch
967     internal data. */
968
969  return Py_FindMethod(comp_methods, (PyObject *)self, name);
970}
971
972static PyObject *
973Decomp_getattr(compobject *self, char *name)
974{
975    PyObject * retval;
976
977    ENTER_ZLIB
978
979    if (strcmp(name, "unused_data") == 0) {
980        Py_INCREF(self->unused_data);
981        retval = self->unused_data;
982    } else if (strcmp(name, "unconsumed_tail") == 0) {
983        Py_INCREF(self->unconsumed_tail);
984        retval = self->unconsumed_tail;
985    } else
986        retval = Py_FindMethod(Decomp_methods, (PyObject *)self, name);
987
988    LEAVE_ZLIB
989
990    return retval;
991}
992
993PyDoc_STRVAR(adler32__doc__,
994"adler32(string[, start]) -- Compute an Adler-32 checksum of string.\n"
995"\n"
996"An optional starting value can be specified.  The returned checksum is\n"
997"a signed integer.");
998
999static PyObject *
1000PyZlib_adler32(PyObject *self, PyObject *args)
1001{
1002    unsigned int adler32val = 1;  /* adler32(0L, Z_NULL, 0) */
1003    Byte *buf;
1004    Py_ssize_t len;
1005    int signed_val;
1006
1007    if (!PyArg_ParseTuple(args, "s#|I:adler32", &buf, &len, &adler32val))
1008        return NULL;
1009
1010    /* Avoid truncation of length for very large buffers. adler32() takes
1011       length as an unsigned int, which may be narrower than Py_ssize_t. */
1012    while ((size_t)len > UINT_MAX) {
1013        adler32val = adler32(adler32val, buf, UINT_MAX);
1014        buf += (size_t) UINT_MAX;
1015        len -= (size_t) UINT_MAX;
1016    }
1017    /* In Python 2.x we return a signed integer regardless of native platform
1018     * long size (the 32bit unsigned long is treated as 32-bit signed and sign
1019     * extended into a 64-bit long inside the integer object).  3.0 does the
1020     * right thing and returns unsigned. http://bugs.python.org/issue1202 */
1021    signed_val = adler32(adler32val, buf, (unsigned int)len);
1022    return PyInt_FromLong(signed_val);
1023}
1024
1025PyDoc_STRVAR(crc32__doc__,
1026"crc32(string[, start]) -- Compute a CRC-32 checksum of string.\n"
1027"\n"
1028"An optional starting value can be specified.  The returned checksum is\n"
1029"a signed integer.");
1030
1031static PyObject *
1032PyZlib_crc32(PyObject *self, PyObject *args)
1033{
1034    unsigned int crc32val = 0;  /* crc32(0L, Z_NULL, 0) */
1035    Byte *buf;
1036    Py_ssize_t len;
1037    int signed_val;
1038
1039    if (!PyArg_ParseTuple(args, "s#|I:crc32", &buf, &len, &crc32val))
1040        return NULL;
1041
1042    /* Avoid truncation of length for very large buffers. crc32() takes
1043       length as an unsigned int, which may be narrower than Py_ssize_t. */
1044    while ((size_t)len > UINT_MAX) {
1045        crc32val = crc32(crc32val, buf, UINT_MAX);
1046        buf += (size_t) UINT_MAX;
1047        len -= (size_t) UINT_MAX;
1048    }
1049    /* In Python 2.x we return a signed integer regardless of native platform
1050     * long size (the 32bit unsigned long is treated as 32-bit signed and sign
1051     * extended into a 64-bit long inside the integer object).  3.0 does the
1052     * right thing and returns unsigned. http://bugs.python.org/issue1202 */
1053    signed_val = crc32(crc32val, buf, (unsigned int)len);
1054    return PyInt_FromLong(signed_val);
1055}
1056
1057
1058static PyMethodDef zlib_methods[] =
1059{
1060    {"adler32", (PyCFunction)PyZlib_adler32, METH_VARARGS,
1061                adler32__doc__},
1062    {"compress", (PyCFunction)PyZlib_compress,  METH_VARARGS,
1063                 compress__doc__},
1064    {"compressobj", (PyCFunction)PyZlib_compressobj, METH_VARARGS,
1065                    compressobj__doc__},
1066    {"crc32", (PyCFunction)PyZlib_crc32, METH_VARARGS,
1067              crc32__doc__},
1068    {"decompress", (PyCFunction)PyZlib_decompress, METH_VARARGS,
1069                   decompress__doc__},
1070    {"decompressobj", (PyCFunction)PyZlib_decompressobj, METH_VARARGS,
1071                   decompressobj__doc__},
1072    {NULL, NULL}
1073};
1074
1075static PyTypeObject Comptype = {
1076    PyVarObject_HEAD_INIT(0, 0)
1077    "zlib.Compress",
1078    sizeof(compobject),
1079    0,
1080    (destructor)Comp_dealloc,       /*tp_dealloc*/
1081    0,                              /*tp_print*/
1082    (getattrfunc)Comp_getattr,      /*tp_getattr*/
1083    0,                              /*tp_setattr*/
1084    0,                              /*tp_compare*/
1085    0,                              /*tp_repr*/
1086    0,                              /*tp_as_number*/
1087    0,                              /*tp_as_sequence*/
1088    0,                              /*tp_as_mapping*/
1089};
1090
1091static PyTypeObject Decomptype = {
1092    PyVarObject_HEAD_INIT(0, 0)
1093    "zlib.Decompress",
1094    sizeof(compobject),
1095    0,
1096    (destructor)Decomp_dealloc,     /*tp_dealloc*/
1097    0,                              /*tp_print*/
1098    (getattrfunc)Decomp_getattr,    /*tp_getattr*/
1099    0,                              /*tp_setattr*/
1100    0,                              /*tp_compare*/
1101    0,                              /*tp_repr*/
1102    0,                              /*tp_as_number*/
1103    0,                              /*tp_as_sequence*/
1104    0,                              /*tp_as_mapping*/
1105};
1106
1107PyDoc_STRVAR(zlib_module_documentation,
1108"The functions in this module allow compression and decompression using the\n"
1109"zlib library, which is based on GNU zip.\n"
1110"\n"
1111"adler32(string[, start]) -- Compute an Adler-32 checksum.\n"
1112"compress(string[, level]) -- Compress string, with compression level in 0-9.\n"
1113"compressobj([level]) -- Return a compressor object.\n"
1114"crc32(string[, start]) -- Compute a CRC-32 checksum.\n"
1115"decompress(string,[wbits],[bufsize]) -- Decompresses a compressed string.\n"
1116"decompressobj([wbits]) -- Return a decompressor object.\n"
1117"\n"
1118"'wbits' is window buffer size and container format.\n"
1119"Compressor objects support compress() and flush() methods; decompressor\n"
1120"objects support decompress() and flush().");
1121
1122PyMODINIT_FUNC
1123PyInit_zlib(void)
1124{
1125    PyObject *m, *ver;
1126    Py_TYPE(&Comptype) = &PyType_Type;
1127    Py_TYPE(&Decomptype) = &PyType_Type;
1128    m = Py_InitModule4("zlib", zlib_methods,
1129                       zlib_module_documentation,
1130                       (PyObject*)NULL,PYTHON_API_VERSION);
1131    if (m == NULL)
1132        return;
1133
1134    ZlibError = PyErr_NewException("zlib.error", NULL, NULL);
1135    if (ZlibError != NULL) {
1136        Py_INCREF(ZlibError);
1137        PyModule_AddObject(m, "error", ZlibError);
1138    }
1139    PyModule_AddIntConstant(m, "MAX_WBITS", MAX_WBITS);
1140    PyModule_AddIntConstant(m, "DEFLATED", DEFLATED);
1141    PyModule_AddIntConstant(m, "DEF_MEM_LEVEL", DEF_MEM_LEVEL);
1142    PyModule_AddIntConstant(m, "Z_BEST_SPEED", Z_BEST_SPEED);
1143    PyModule_AddIntConstant(m, "Z_BEST_COMPRESSION", Z_BEST_COMPRESSION);
1144    PyModule_AddIntConstant(m, "Z_DEFAULT_COMPRESSION", Z_DEFAULT_COMPRESSION);
1145    PyModule_AddIntConstant(m, "Z_FILTERED", Z_FILTERED);
1146    PyModule_AddIntConstant(m, "Z_HUFFMAN_ONLY", Z_HUFFMAN_ONLY);
1147    PyModule_AddIntConstant(m, "Z_DEFAULT_STRATEGY", Z_DEFAULT_STRATEGY);
1148
1149    PyModule_AddIntConstant(m, "Z_FINISH", Z_FINISH);
1150    PyModule_AddIntConstant(m, "Z_NO_FLUSH", Z_NO_FLUSH);
1151    PyModule_AddIntConstant(m, "Z_SYNC_FLUSH", Z_SYNC_FLUSH);
1152    PyModule_AddIntConstant(m, "Z_FULL_FLUSH", Z_FULL_FLUSH);
1153
1154    ver = PyString_FromString(ZLIB_VERSION);
1155    if (ver != NULL)
1156        PyModule_AddObject(m, "ZLIB_VERSION", ver);
1157
1158    PyModule_AddStringConstant(m, "__version__", "1.0");
1159
1160#ifdef WITH_THREAD
1161    zlib_lock = PyThread_allocate_lock();
1162#endif /* WITH_THREAD */
1163}
1164