Merge pull request #136 from pycompression/dicts

Fix bugs when setting zdict

Author: rhpvorderman

CHANGELOG.rst

@@ -9,6 +9,8 @@ Changelog
 
 version 1.2.0-dev
 -----------------
++ Fix bug where zdict's could not be set for ``isal_zlib.decompressobj`` and
+  ``igzip_lib.IgzipDecompressor``.
 + Escape GIL when calling inflate, deflate, crc32 and adler32 functions to
   just like in CPython. This allows for utilising more CPU cores in combination
   with the threading module. This comes with a very slight cost in efficiency

src/isal/igzip_libmodule.c

@@ -49,6 +49,32 @@ IgzipDecompressor_dealloc(IgzipDecompressor *self)
     Py_TYPE(self)->tp_free((PyObject *)self);
 }
 
+
+static int
+set_inflate_zdict_IgzipDecompressor(IgzipDecompressor *self)
+{
+    Py_buffer zdict_buf;
+    if (PyObject_GetBuffer(self->zdict, &zdict_buf, PyBUF_SIMPLE) == -1) {
+        return -1;
+    }
+    if ((size_t)zdict_buf.len > UINT32_MAX) {
+        PyErr_SetString(PyExc_OverflowError,
+                        "zdict length does not fit in an unsigned 32-bit integer");
+        PyBuffer_Release(&zdict_buf);
+        return -1;
+    }
+    int err;
+    err = isal_inflate_set_dict(&self->state,
+                                zdict_buf.buf, (uint32_t)zdict_buf.len);
+    PyBuffer_Release(&zdict_buf);
+    if (err != ISAL_DECOMP_OK) {
+        isal_inflate_error(err);
+        return -1;
+    }
+    return 0;
+}
+
+
 /* Decompress data of length d->bzs_avail_in_real in d->state.next_in.  The output
    buffer is allocated dynamically and returned.  At most max_length bytes are
    returned, so some of the input may not be consumed. d->state.next_in and
@@ -105,12 +131,28 @@ decompress_buf(IgzipDecompressor *self, Py_ssize_t max_length)
             err = isal_inflate(&(self->state));
             Py_END_ALLOW_THREADS;
 
-            if (err != ISAL_DECOMP_OK){
-                isal_inflate_error(err);
-                goto error;
+            switch(err) {
+                case ISAL_DECOMP_OK:
+                    break;
+                case ISAL_NEED_DICT:                
+                    if  (self->zdict != NULL) {
+                        if (set_inflate_zdict_IgzipDecompressor(self) < 0) {
+                            goto error;
+                        }
+                        else 
+                            break;
+                    }
+                    else {
+                        isal_inflate_error(err);
+                        goto error;
+                    }
+                default:
+                    isal_inflate_error(err);
+                    goto error;
             }
-        } while (self->state.avail_out == 0 && 
-                 self->state.block_state != ISAL_BLOCK_FINISH);
+
+        } while (self->state.avail_out == 0 || err == ISAL_NEED_DICT);
+    
     } while(self->avail_in_real != 0 && 
             self->state.block_state != ISAL_BLOCK_FINISH);
 
@@ -414,7 +456,6 @@ igzip_lib_IgzipDecompressor__new__(PyTypeObject *type,
         return NULL;
     }
     IgzipDecompressor *self = PyObject_New(IgzipDecompressor, type); 
-    int err;
     self->eof = 0;
     self->needs_input = 1;
     self->avail_in_real = 0;
@@ -436,26 +477,10 @@ igzip_lib_IgzipDecompressor__new__(PyTypeObject *type,
     self->state.hist_bits = hist_bits;
     self->state.crc_flag = flag;
     if (self->zdict != NULL){
-        Py_buffer zdict_buf;
-        if (PyObject_GetBuffer(self->zdict, &zdict_buf, PyBUF_SIMPLE) == -1) {
-                Py_CLEAR(self);
-                return NULL;
-        }
-        if ((size_t)zdict_buf.len > UINT32_MAX) {
-            PyErr_SetString(PyExc_OverflowError,
-                           "zdict length does not fit in an unsigned 32-bits int");
-            PyBuffer_Release(&zdict_buf);
+        if (set_inflate_zdict_IgzipDecompressor(self) < 0) {
             Py_CLEAR(self);
             return NULL;
         }
-        err = isal_inflate_set_dict(&(self->state), zdict_buf.buf, 
-                                    (uint32_t)zdict_buf.len);
-        PyBuffer_Release(&zdict_buf);
-        if (err != ISAL_DECOMP_OK) {
-            isal_inflate_error(err);
-            Py_CLEAR(self);
-            return NULL;
-        }        
     }
     return (PyObject *)self;
 }

src/isal/isal_zlibmodule.c

@@ -724,14 +724,28 @@ isal_zlib_Decompress_decompress_impl(decompobject *self, Py_buffer *data,
             Py_BEGIN_ALLOW_THREADS
             err = isal_inflate(&self->zst);
             Py_END_ALLOW_THREADS
-
-            if (err != ISAL_DECOMP_OK){
-                isal_inflate_error(err);
-                goto abort;
+            
+            switch(err) {
+                case ISAL_DECOMP_OK:
+                    break;
+                case ISAL_NEED_DICT:                
+                    if  (self->zdict != NULL) {
+                        if (set_inflate_zdict(self) < 0) {
+                            goto abort;
+                        }
+                        else 
+                            break;
+                    }
+                    else {
+                        isal_inflate_error(err);
+                        goto abort;
+                    }
+                default:
+                    isal_inflate_error(err);
+                    goto abort;
             }
 
-        } while (self->zst.avail_out == 0 && 
-                 self->zst.block_state != ISAL_BLOCK_FINISH);
+        } while (self->zst.avail_out == 0 || err == ISAL_NEED_DICT);
 
     } while (self->zst.block_state != ISAL_BLOCK_FINISH && ibuflen != 0);
 

tests/test_compat.py

@@ -165,3 +165,21 @@ def test_unused_data(unused_size, wbits):
     result = decompressor.decompress(compressed + unused_data)
     assert result == data
     assert decompressor.unused_data == unused_data
+
+
+def test_zlib_dictionary_decompress():
+    dictionary = b"bla"
+    data = b"bladiebla"
+    compobj = zlib.compressobj(zdict=dictionary)
+    compressed = compobj.compress(data) + compobj.flush()
+    decompobj = isal_zlib.decompressobj(zdict=dictionary)
+    assert decompobj.decompress(compressed) == data
+
+
+def test_isal_zlib_dictionary_decompress():
+    dictionary = b"bla"
+    data = b"bladiebla"
+    compobj = isal_zlib.compressobj(zdict=dictionary)
+    compressed = compobj.compress(data) + compobj.flush()
+    decompobj = zlib.decompressobj(zdict=dictionary)
+    assert decompobj.decompress(compressed) == data

tests/test_igzip_lib.py

@@ -224,6 +224,12 @@ def test_failure(self):
         with pytest.raises(Exception):
             igzd.decompress(self.BAD_DATA * 30)
 
+    def test_dictionary(self):
+        compressor = zlib.compressobj(zdict=b"bla")
+        data = compressor.compress(b"bladiebla") + compressor.flush()
+        igzd = IgzipDecompressor(flag=DECOMP_ZLIB, zdict=b"bla")
+        assert igzd.decompress(data) == b"bladiebla"
+
 
 @pytest.mark.parametrize("test_offset", range(5))
 def test_igzip_decompressor_raw_deflate_unused_data_zlib(test_offset):