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