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