Fix zlib headers for preset dictionary compression

CHANGELOG.rst

@@ -11,6 +11,8 @@ Changelog
 version 1.8.1-dev
 -----------------
 + Restore PyPy wheel builds.
++ Fix zlib streams compressed with a preset dictionary so they can be decoded
+  by standard zlib, including builds linked against older ISA-L releases.
 
 version 1.8.0
 -----------------

src/isal/isal_zlibmodule.c

@@ -570,6 +570,7 @@ typedef struct
     PyObject_HEAD
     uint8_t *level_buf;
     PyObject *zdict;
+    struct isal_zlib_header zlib_header;
     int is_initialised;
     PyThread_type_lock lock;
     // isal_zstream should be at the bottom as it contains buffers inside the struct.
@@ -595,6 +596,7 @@ newcompobject()
         return NULL;
     self->is_initialised = 0;
     self->zdict = NULL;
+    self->zlib_header.dict_flag = 0;
     self->level_buf = NULL;
     self->lock = PyThread_allocate_lock();
     if (self->lock == NULL) {
@@ -755,8 +757,18 @@ isal_zlib_compressobj_impl(PyObject *module, int level, int method, int wbits,
     } else {
         err = isal_deflate_set_dict(&(self->zst),
                                     zdict->buf, (uint32_t)zdict->len);
-        if (err == COMP_OK)
+        if (err == COMP_OK) {
+            if (flag == IGZIP_ZLIB && zdict->len != 0) {
+                isal_zlib_header_init(&self->zlib_header);
+                self->zlib_header.info = hist_bits - 8;
+                self->zlib_header.level = level == 0 ? 0 : 1;
+                self->zlib_header.dict_flag = 1;
+                self->zlib_header.dict_id = isal_adler32(
+                    1, zdict->buf, (uint32_t)zdict->len);
+                self->zst.gzip_flag = IGZIP_ZLIB_NO_HDR;
+            }
             goto success;
+        }
         PyErr_SetString(PyExc_ValueError, "Invalid dictionary");
         goto error;
         }
@@ -817,6 +829,31 @@ isal_zlib_decompressobj_impl(PyObject *module, int wbits, PyObject *zdict)
     return (PyObject *)self;
 }
 
+/* Called with the compressor locked and a newly arranged output buffer.
+ * The six-byte dictionary header is written once, including when flush()
+ * is the first operation. ISA-L still generates the Adler-32 trailer. */
+static int
+write_pending_zlib_header(compobject *self)
+{
+    if (self->zlib_header.dict_flag) {
+        uint8_t *header = self->zst.next_out;
+        if (isal_write_zlib_header(&self->zst, &self->zlib_header) != 0) {
+            PyErr_SetString(PyExc_RuntimeError,
+                            "Insufficient space for the zlib dictionary header");
+            return -1;
+        }
+        /* Older ISA-L releases write DICTID in little-endian order.
+         * RFC 1950 requires network byte order, regardless of ISA-L version. */
+        uint32_t dict_id = self->zlib_header.dict_id;
+        header[2] = (uint8_t)(dict_id >> 24);
+        header[3] = (uint8_t)(dict_id >> 16);
+        header[4] = (uint8_t)(dict_id >> 8);
+        header[5] = (uint8_t)dict_id;
+        self->zlib_header.dict_flag = 0;
+    }
+    return 0;
+}
+
 static PyObject *
 isal_zlib_Compress_compress_impl(compobject *self, Py_buffer *data)
 /*[clinic end generated code: output=5d5cd791cbc6a7f4 input=0d95908d6e64fab8]*/
@@ -838,6 +875,9 @@ isal_zlib_Compress_compress_impl(compobject *self, Py_buffer *data)
             if (obuflen < 0)
                 goto error;
 
+            if (write_pending_zlib_header(self) < 0)
+                goto error;
+
             Py_BEGIN_ALLOW_THREADS
             err = isal_deflate(&self->zst);
             Py_END_ALLOW_THREADS
@@ -1049,6 +1089,11 @@ isal_zlib_Compress_flush_impl(compobject *self, int mode)
             goto error;
         }
 
+        if (write_pending_zlib_header(self) < 0) {
+            Py_CLEAR(RetVal);
+            goto error;
+        }
+
         Py_BEGIN_ALLOW_THREADS
         err = isal_deflate(&self->zst);
         Py_END_ALLOW_THREADS

tests/test_dictionary_header.py

@@ -0,0 +1,70 @@
+"""Cross-library checks for dictionary headers and compressor lifecycle."""
+import random
+import zlib
+
+from isal import isal_zlib
+
+import pytest
+
+
+@pytest.mark.parametrize('level', range(4))
+@pytest.mark.parametrize('wbits', [9, 15, -15])
+@pytest.mark.parametrize('dictionary_size', [11, 32768, 40000])
+@pytest.mark.parametrize(
+    'first_operation', ['compress', 'empty', 'sync', 'finish'])
+def test_dictionary_cross_decode(level, wbits, dictionary_size,
+                                 first_operation):
+    dictionary = random.Random(21).randbytes(dictionary_size)
+    data = dictionary[-256:] * 5 + random.Random(22).randbytes(4096)
+    compressor = isal_zlib.compressobj(level, wbits=wbits, zdict=dictionary)
+    packed = b''
+    if first_operation == 'empty':
+        packed += compressor.compress(b'')
+    elif first_operation == 'sync':
+        packed += compressor.flush(isal_zlib.Z_SYNC_FLUSH)
+    elif first_operation == 'finish':
+        data = b''
+    for pos in range(0, len(data), 17):
+        packed += compressor.compress(data[pos:pos+17])
+    packed += compressor.flush()
+    if wbits > 0:
+        assert packed[0] >> 4 == wbits - 8
+        assert packed[1] & 32
+        assert int.from_bytes(packed[2:6], 'big') == zlib.adler32(dictionary)
+        assert int.from_bytes(packed[:2], 'big') % 31 == 0
+    decompressor = zlib.decompressobj(wbits, zdict=dictionary)
+    assert decompressor.decompress(packed) + decompressor.flush() == data
+    assert decompressor.eof
+    assert decompressor.unused_data == b''
+
+
+@pytest.mark.parametrize('level', range(4))
+@pytest.mark.parametrize(
+    'flush_mode', [isal_zlib.Z_SYNC_FLUSH, isal_zlib.Z_FULL_FLUSH])
+def test_dictionary_multiple_flushes(level, flush_mode):
+    dictionary = b'hello world'
+    compressor = isal_zlib.compressobj(level, zdict=dictionary)
+    data = dictionary * 1000
+    packed = compressor.compress(data[:5500]) + compressor.flush(flush_mode)
+    packed += compressor.compress(data[5500:]) + compressor.flush()
+    for module in (zlib, isal_zlib):
+        decompressor = module.decompressobj(zdict=dictionary)
+        assert decompressor.decompress(packed) + decompressor.flush() == data
+
+
+@pytest.mark.parametrize('dictionary', [None, b''])
+def test_no_dictionary_header(dictionary):
+    kwargs = {} if dictionary is None else {'zdict': dictionary}
+    compressor = isal_zlib.compressobj(**kwargs)
+    data = b'hello world' * 100
+    packed = compressor.compress(data) + compressor.flush()
+    assert not packed[1] & 32
+    assert zlib.decompress(packed) == data
+
+
+def test_wrong_dictionary_rejected_by_zlib():
+    compressor = isal_zlib.compressobj(zdict=b'hello world')
+    packed = compressor.compress(b'hello world' * 10) + compressor.flush()
+    decompressor = zlib.decompressobj(zdict=b'wrong dictionary')
+    with pytest.raises(zlib.error):
+        decompressor.decompress(packed)