Repository: pauldmccarthy/indexed_gzip (language: python)

Issue:
Fix memory leak in cases when inflate returned an error.
Apparently `inflateEnd` needs to be called regardless of the
`inflate` return value. See https://github.com/python/cpython/blob/main/Modules/zlibmodule.c for the reference.


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 15099f86..be91fb4f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,16 +1,21 @@
 # `indexed_gzip` changelog
 
+## 1.6.2 (September 2nd 2021)
+
+
+* Fixed a memory leak when initialising decompression / inflation (#82, #83).
+
 
 ## 1.6.1 (May 25th 2021)
 
 
-* Tests requiring `nibabel` are now skipped, rather than causing failure.
+* Tests requiring `nibabel` are now skipped, rather than causing failure (#78).
 
 
 ## 1.6.0 (May 23rd 2021)
 
 
-* Python 2.7 wheels for Windows are no longer being built (#71, 73).
+* Python 2.7 wheels for Windows are no longer being built (#71, #73).
 * A backwards-compatible change to the index file format, to accommodate seek
   points at stream boundaries. Index files created with older versions of
   `indexed_gzip` can still be loaded, but index files created with
diff --git a/indexed_gzip/zran.c b/indexed_gzip/zran.c
index 229a95a1..3ad1f90d 100644
--- a/indexed_gzip/zran.c
+++ b/indexed_gzip/zran.c
@@ -763,7 +763,13 @@ void zran_free(zran_index_t *index) {
     for (i = 0; i < index->npoints; i++) {
         pt = &(index->list[i]);
 
-        free(pt->data);
+        /*
+         * points at compression stream boundaries
+         * have no data associated with them
+         */
+        if (pt->data != NULL) {
+            free(pt->data);
+        }
     }
 
     free(index->list);
@@ -1185,19 +1191,19 @@ int _zran_add_point(zran_index_t  *index,
 
 /* Initialise the given z_stream struct for decompression/inflation. */
 int _zran_init_zlib_inflate(zran_index_t *index,
-                            z_stream     *stream,
+                            z_stream     *strm,
                             zran_point_t *point) {
 
     int           ret;
-    int           windowBits;
+    int           window;
     int64_t       seek_loc;
     unsigned long bytes_read;
 
-    bytes_read     = stream->avail_in;
-    windowBits     = index->log_window_size;
-    stream->zalloc = Z_NULL;
-    stream->zfree  = Z_NULL;
-    stream->opaque = Z_NULL;
+    bytes_read   = strm->avail_in;
+    window       = index->log_window_size;
+    strm->zalloc = Z_NULL;
+    strm->zfree  = Z_NULL;
+    strm->opaque = Z_NULL;
 
     /*
      * If we're starting from the the current location in
@@ -1211,9 +1217,9 @@ int _zran_init_zlib_inflate(zran_index_t *index,
 
         zran_log("_zran_init_zlib_inflate from current "
                  "seek location (expecting GZIP header)\n");
-        if (inflateInit2(stream, windowBits + 32) != Z_OK) { goto fail; }
-        if (inflate(stream, Z_BLOCK)              != Z_OK) { goto fail; }
-        if (inflateEnd(stream)                    != Z_OK) { goto fail; }
+        if (inflateInit2(strm, window + 32) != Z_OK) { goto fail; }
+        if (inflate(strm, Z_BLOCK)          != Z_OK) { goto fail_free_strm; }
+        if (inflateEnd(strm)                != Z_OK) { goto fail; }
     }
 
     /*
@@ -1246,7 +1252,7 @@ int _zran_init_zlib_inflate(zran_index_t *index,
      * in _zran_inflate).
      */
 
-    if (inflateInit2(stream, -windowBits) != Z_OK) {
+    if (inflateInit2(strm, -window) != Z_OK) {
         goto fail;
     }
 
@@ -1268,22 +1274,22 @@ int _zran_init_zlib_inflate(zran_index_t *index,
             ret = getc_(index->fd, index->f);
 
             if (ret == -1 && ferror_(index->fd, index->f)) {
-                goto fail;
+                goto fail_free_strm;
             }
 
-            if (inflatePrime(stream,
+            if (inflatePrime(strm,
                              point->bits, ret >> (8 - point->bits)) != Z_OK)
-                goto fail;
+                goto fail_free_strm;
         }
 
         /*
          * Initialise the inflate stream
          * with the index point data.
          */
-        if (inflateSetDictionary(stream,
+        if (inflateSetDictionary(strm,
                                  point->data,
                                  index->window_size) != Z_OK)
-            goto fail;
+            goto fail_free_strm;
     }
 
     /*
@@ -1295,14 +1301,24 @@ int _zran_init_zlib_inflate(zran_index_t *index,
     index->stream_crc32 = 0;
 
     zran_log("_zran_zlib_init_inflate: initialised, read %i bytes\n",
-             bytes_read - stream->avail_in);
+             bytes_read - strm->avail_in);
 
     /*
      * Return the number of bytes of compressed
      * data, if any that were read over
      */
-    return bytes_read - stream->avail_in;
+    return bytes_read - strm->avail_in;
 
+/*
+ * Something has gone wrong, but
+ * inflateInit2 has been called,
+ * so we need to call inflateEnd.
+ * Falls through to the fail:
+ * clause.
+ */
+fail_free_strm:
+    inflateEnd(strm);
+/* Something has gone wrong */
 fail:
     return -1;
 }


Test patch:
diff --git a/indexed_gzip/tests/ctest_zran.pyx b/indexed_gzip/tests/ctest_zran.pyx
index 76381f4c..ce20993a 100644
--- a/indexed_gzip/tests/ctest_zran.pyx
+++ b/indexed_gzip/tests/ctest_zran.pyx
@@ -23,6 +23,14 @@ import                    contextlib
 
 import numpy as np
 
+# The stdlib resource module is only
+# available on unix-like platforms.
+try:
+    import resource
+except ImportError:
+    resource = None
+
+
 cimport numpy as np
 
 from posix.types cimport  off_t
@@ -1399,3 +1407,41 @@ def test_standard_usage_with_null_padding(concat):
 
     pybuf = <bytes>(<char *>buffer)[:dsize]
     assert np.all(np.frombuffer(pybuf, dtype=np.uint32) == data)
+
+
+# pauldmccarthy/indexed_gzip#82
+def test_inflateInit_leak_on_error():
+    """Make sure memory is not leaked after a successful call to
+    inflateInit2(), but then a failure on subsequent zlib calls.
+    """
+
+    cdef zran.zran_index_t index
+
+    # inflateInit2 is called twice in the _zran_zlib_init_inflate function.
+    # We can target the first call by passing a file containing random noise.
+    # I haven't yet figured out a reliable way to target the second call.
+    f = BytesIO(np.arange(1, 100).tobytes())
+
+    iters = np.arange(1, 10000)
+    mem   = np.zeros(10000, dtype=np.uint64)
+
+    for i in iters:
+        assert not zran.zran_init(&index,
+                                  NULL,
+                                  <PyObject*>f,
+                                  1048576,
+                                  32768,
+                                  131072,
+                                  zran.ZRAN_AUTO_BUILD)
+        assert zran.zran_seek(&index, 20, SEEK_SET, NULL) == \
+            zran.ZRAN_SEEK_FAIL
+        zran.zran_free(&index)
+
+        if resource is not None:
+            mem[i] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+
+    # We expect to see some small growth in memory
+    # usage for the first few iterations, but then
+    # it should remain stable
+    mem = mem[5:]
+    assert np.all(mem == mem[0])
diff --git a/indexed_gzip/tests/test_zran.py b/indexed_gzip/tests/test_zran.py
index 5802f808..2cb597d9 100644
--- a/indexed_gzip/tests/test_zran.py
+++ b/indexed_gzip/tests/test_zran.py
@@ -127,3 +127,6 @@ def test_crc_validation(concat):
 
     def test_standard_usage_with_null_padding(concat):
         ctest_zran.test_standard_usage_with_null_padding(concat)
+
+    def test_inflateInit_leak_on_error():
+        ctest_zran.test_inflateInit_leak_on_error()