From 3f0ee83f14f2ce2c05360f5ed326844388353b6f Mon Sep 17 00:00:00 2001
From: Martin Panter <vadmium+py@gmail.com>
Date: Sun, 5 Jun 2016 10:48:34 +0000
Subject: [PATCH] Issue #27164: Allow decompressing raw Deflate streams with
 predefined zdict

Based on patch by Xiang Zhang.
---
 Lib/test/test_zlib.py |  9 +++++++
 Misc/NEWS             |  3 +++
 Modules/zlibmodule.c  | 62 ++++++++++++++++++++++++++++++-------------
 3 files changed, 56 insertions(+), 18 deletions(-)

diff --git a/Lib/test/test_zlib.py b/Lib/test/test_zlib.py
index b6bebbced11..78ecade950a 100644
--- a/Lib/test/test_zlib.py
+++ b/Lib/test/test_zlib.py
@@ -553,6 +553,15 @@ class CompressObjectTestCase(BaseCompressTestCase, unittest.TestCase):
                 self.assertEqual(dco.unconsumed_tail, b'')
                 self.assertEqual(dco.unused_data, remainder)
 
+    # issue27164
+    def test_decompress_raw_with_dictionary(self):
+        zdict = b'abcdefghijklmnopqrstuvwxyz'
+        co = zlib.compressobj(wbits=-zlib.MAX_WBITS, zdict=zdict)
+        comp = co.compress(zdict) + co.flush()
+        dco = zlib.decompressobj(wbits=-zlib.MAX_WBITS, zdict=zdict)
+        uncomp = dco.decompress(comp) + dco.flush()
+        self.assertEqual(zdict, uncomp)
+
     def test_flush_with_freed_input(self):
         # Issue #16411: decompressor accesses input to last decompress() call
         # in flush(), even if this object has been freed in the meanwhile.
diff --git a/Misc/NEWS b/Misc/NEWS
index 5e7d1edfa9a..44920842883 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -131,6 +131,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #27164: In the zlib module, allow decompressing raw Deflate streams
+  with a predefined zdict.  Based on patch by Xiang Zhang.
+
 - Issue #24291: Fix wsgiref.simple_server.WSGIRequestHandler to completely
   write data to the client.  Previously it could do partial writes and
   truncate data.  Also, wsgiref.handler.ServerHandler can now handle stdout
diff --git a/Modules/zlibmodule.c b/Modules/zlibmodule.c
index 02c747ee822..16cc17886bf 100644
--- a/Modules/zlibmodule.c
+++ b/Modules/zlibmodule.c
@@ -22,6 +22,10 @@
     #define LEAVE_ZLIB(obj)
 #endif
 
+#if defined(ZLIB_VERNUM) && ZLIB_VERNUM >= 0x1221
+#define AT_LEAST_ZLIB_1_2_2_1
+#endif
+
 /* The following parameters are copied from zutil.h, version 0.95 */
 #define DEFLATED   8
 #if MAX_MEM_LEVEL >= 8
@@ -474,6 +478,31 @@ zlib_compressobj_impl(PyModuleDef *module, int level, int method, int wbits,
     return (PyObject*)self;
 }
 
+static int
+set_inflate_zdict(compobject *self)
+{
+    Py_buffer zdict_buf;
+    int err;
+
+    if (PyObject_GetBuffer(self->zdict, &zdict_buf, PyBUF_SIMPLE) == -1) {
+        return -1;
+    }
+    if ((size_t)zdict_buf.len > UINT_MAX) {
+        PyErr_SetString(PyExc_OverflowError,
+                        "zdict length does not fit in an unsigned int");
+        PyBuffer_Release(&zdict_buf);
+        return -1;
+    }
+    err = inflateSetDictionary(&(self->zst),
+                               zdict_buf.buf, (unsigned int)zdict_buf.len);
+    PyBuffer_Release(&zdict_buf);
+    if (err != Z_OK) {
+        zlib_error(self->zst, err, "while setting zdict");
+        return -1;
+    }
+    return 0;
+}
+
 /*[clinic input]
 zlib.decompressobj
 
@@ -515,6 +544,20 @@ zlib_decompressobj_impl(PyModuleDef *module, int wbits, PyObject *zdict)
     switch(err) {
     case (Z_OK):
         self->is_initialised = 1;
+        if (self->zdict != NULL && wbits < 0) {
+#ifdef AT_LEAST_ZLIB_1_2_2_1
+            if (set_inflate_zdict(self) < 0) {
+                Py_DECREF(self);
+                return NULL;
+            }
+#else
+            PyErr_Format(ZlibError,
+                         "zlib version %s does not allow raw inflate with dictionary",
+                         ZLIB_VERSION);
+            Py_DECREF(self);
+            return NULL;
+#endif
+        }
         return (PyObject*)self;
     case(Z_STREAM_ERROR):
         Py_DECREF(self);
@@ -741,29 +784,12 @@ zlib_Decompress_decompress_impl(compobject *self, Py_buffer *data,
     Py_END_ALLOW_THREADS
 
     if (err == Z_NEED_DICT && self->zdict != NULL) {
-        Py_buffer zdict_buf;
-        if (PyObject_GetBuffer(self->zdict, &zdict_buf, PyBUF_SIMPLE) == -1) {
+        if (set_inflate_zdict(self) < 0) {
             Py_DECREF(RetVal);
             RetVal = NULL;
             goto error;
         }
 
-        if ((size_t)zdict_buf.len > UINT_MAX) {
-            PyErr_SetString(PyExc_OverflowError,
-                    "zdict length does not fit in an unsigned int");
-            PyBuffer_Release(&zdict_buf);
-            Py_CLEAR(RetVal);
-            goto error;
-        }
-
-        err = inflateSetDictionary(&(self->zst),
-                                   zdict_buf.buf, (unsigned int)zdict_buf.len);
-        PyBuffer_Release(&zdict_buf);
-        if (err != Z_OK) {
-            zlib_error(self->zst, err, "while decompressing data");
-            Py_CLEAR(RetVal);
-            goto error;
-        }
         /* Repeat the call to inflate. */
         Py_BEGIN_ALLOW_THREADS
         err = inflate(&(self->zst), Z_SYNC_FLUSH);