• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 
29 static 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 
59 static PyTypeObject Comptype;
60 static PyTypeObject Decomptype;
61 
62 static PyObject *ZlibError;
63 
64 typedef struct
65 {
66     PyObject_HEAD
67     z_stream zst;
68     PyObject *unused_data;
69     PyObject *unconsumed_tail;
70     int is_initialised;
71 } compobject;
72 
73 static void
zlib_error(z_stream zst,int err,char * msg)74 zlib_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 
102 PyDoc_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 
107 PyDoc_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 
112 static compobject *
newcompobject(PyTypeObject * type)113 newcompobject(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 
133 static void
arrange_input_buffer(z_stream * zst,Py_ssize_t * remains)134 arrange_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 
144 static Py_ssize_t
arrange_output_buffer_with_maximum(z_stream * zst,PyObject ** buffer,Py_ssize_t length,Py_ssize_t max_length)145 arrange_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 
185 static Py_ssize_t
arrange_output_buffer(z_stream * zst,PyObject ** buffer,Py_ssize_t length)186 arrange_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 
198 PyDoc_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 
203 static PyObject *
PyZlib_compress(PyObject * self,PyObject * args)204 PyZlib_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 
279 PyDoc_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 
285 static PyObject *
PyZlib_decompress(PyObject * self,PyObject * args)286 PyZlib_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 
380 static PyObject *
PyZlib_compressobj(PyObject * selfptr,PyObject * args)381 PyZlib_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 
419 static PyObject *
PyZlib_decompressobj(PyObject * selfptr,PyObject * args)420 PyZlib_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 
455 static void
Comp_dealloc(compobject * self)456 Comp_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 
465 static void
Decomp_dealloc(compobject * self)466 Decomp_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 
475 PyDoc_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 
483 static PyObject *
PyZlib_objcompress(compobject * self,PyObject * args)484 PyZlib_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. */
530 static int
save_unconsumed_input(compobject * self,Byte * input,Py_ssize_t inplen,int err)531 save_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 
574 PyDoc_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 
585 static PyObject *
PyZlib_objdecompress(compobject * self,PyObject * args)586 PyZlib_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 
676 PyDoc_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 
684 static PyObject *
PyZlib_flush(compobject * self,PyObject * args)685 PyZlib_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
756 PyDoc_STRVAR(comp_copy__doc__,
757 "copy() -- Return a copy of the compression object.");
758 
759 static PyObject *
PyZlib_copy(compobject * self)760 PyZlib_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 
799 error:
800     LEAVE_ZLIB
801     Py_XDECREF(retval);
802     return NULL;
803 }
804 
805 PyDoc_STRVAR(decomp_copy__doc__,
806 "copy() -- Return a copy of the decompression object.");
807 
808 static PyObject *
PyZlib_uncopy(compobject * self)809 PyZlib_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 
848 error:
849     LEAVE_ZLIB
850     Py_XDECREF(retval);
851     return NULL;
852 }
853 #endif
854 
855 PyDoc_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 
862 static PyObject *
PyZlib_unflush(compobject * self,PyObject * args)863 PyZlib_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 
937 static 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 
950 static 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 
963 static PyObject *
Comp_getattr(compobject * self,char * name)964 Comp_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 
972 static PyObject *
Decomp_getattr(compobject * self,char * name)973 Decomp_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 
993 PyDoc_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 
999 static PyObject *
PyZlib_adler32(PyObject * self,PyObject * args)1000 PyZlib_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 
1025 PyDoc_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 
1031 static PyObject *
PyZlib_crc32(PyObject * self,PyObject * args)1032 PyZlib_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 
1058 static 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 
1075 static 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 
1091 static 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 
1107 PyDoc_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 
1122 PyMODINIT_FUNC
PyInit_zlib(void)1123 PyInit_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