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