1 files changed, 115 insertions, 82 deletions

diff --git a/erts/emulator/nifs/common/zlib_nif.c b/erts/emulator/nifs/common/zlib_nif.c
index a1a65e1946..b709ed5a6f 100644
--- a/erts/emulator/nifs/common/zlib_nif.c
+++ b/erts/emulator/nifs/common/zlib_nif.c
@@ -69,11 +69,27 @@ typedef enum {
     ST_DEFLATE = 1,
     ST_INFLATE = 2,
     ST_CLOSED = 3
-} zlib_state;
+} zlib_state_t;
+
+/* Controls what to do when the user attempts to decompress more data after
+ * Z_STREAM_END has been returned:
+ *
+ * - 'cut' wipes all further input and returns empty results until reset by
+ * the user. This is the default behavior, matching that of the old driver.
+ * - 'reset' resets the state without discarding any input, making it possible
+ * to decompress blindly concatenated streams.
+ * - 'error' crashes with a data error. */
+typedef enum {
+    EOS_BEHAVIOR_ERROR = 0,
+    EOS_BEHAVIOR_RESET = 1,
+    EOS_BEHAVIOR_CUT = 2
+} zlib_eos_behavior_t;
 
 typedef struct {
     z_stream s;
-    zlib_state state;
+    zlib_state_t state;
+
+    zlib_eos_behavior_t eos_behavior;
 
     /* These refer to the plaintext CRC, and are only needed for zlib:crc32/1
      * which is deprecated. */
@@ -90,6 +106,7 @@ typedef struct {
     int inflateChunk_buffer_size;
 
     ErlNifPid controlling_process;
+    ErlNifMutex *controller_lock;
 
     ErlNifIOQueue *input_queue;
 
@@ -101,8 +118,9 @@ typedef struct {
 
 static ERL_NIF_TERM zlib_open(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM zlib_close(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
+static ERL_NIF_TERM zlib_set_controller(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
+
 static ERL_NIF_TERM zlib_deflateInit(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
-static ERL_NIF_TERM zlib_deflateInit2(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM zlib_deflateSetDictionary(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM zlib_deflateReset(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM zlib_deflateEnd(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
@@ -110,7 +128,6 @@ static ERL_NIF_TERM zlib_deflateParams(ErlNifEnv *env, int argc, const ERL_NIF_T
 static ERL_NIF_TERM zlib_deflate(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
 
 static ERL_NIF_TERM zlib_inflateInit(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
-static ERL_NIF_TERM zlib_inflateInit2(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM zlib_inflateSetDictionary(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM zlib_inflateGetDictionary(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM zlib_inflateReset(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
@@ -129,9 +146,13 @@ static ERL_NIF_TERM zlib_setBufSize(ErlNifEnv *env, int argc, const ERL_NIF_TERM
 static ERL_NIF_TERM zlib_enqueue_input(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
 
 static ErlNifFunc nif_funcs[] = {
+    {"close_nif", 1, zlib_close},
+    {"open_nif", 0, zlib_open},
+
+    {"set_controller_nif", 2, zlib_set_controller},
+
     /* deflate */
-    {"deflateInit_nif", 2, zlib_deflateInit},
-    {"deflateInit_nif", 6, zlib_deflateInit2},
+    {"deflateInit_nif", 6, zlib_deflateInit},
     {"deflateSetDictionary_nif", 2, zlib_deflateSetDictionary},
     {"deflateReset_nif", 1, zlib_deflateReset},
     {"deflateEnd_nif", 1, zlib_deflateEnd},
@@ -139,8 +160,7 @@ static ErlNifFunc nif_funcs[] = {
     {"deflate_nif", 4, zlib_deflate},
 
     /* inflate */
-    {"inflateInit_nif", 1, zlib_inflateInit},
-    {"inflateInit_nif", 2, zlib_inflateInit2},
+    {"inflateInit_nif", 3, zlib_inflateInit},
     {"inflateSetDictionary_nif", 2, zlib_inflateSetDictionary},
     {"inflateGetDictionary_nif", 1, zlib_inflateGetDictionary},
     {"inflateReset_nif", 1, zlib_inflateReset},
@@ -150,10 +170,6 @@ static ErlNifFunc nif_funcs[] = {
     /* running checksum */
     {"crc32_nif", 1, zlib_crc32},
 
-    /* open & close */
-    {"close_nif", 1, zlib_close},
-    {"open_nif", 0, zlib_open},
-
     /* The stash keeps a single term alive across calls, and is used in
      * exception_on_need_dict/1 to retain the old error behavior, and for
      * saving data flushed through deflateParams/3. */
@@ -269,9 +285,7 @@ static ERL_NIF_TERM zlib_return(ErlNifEnv *env, int code) {
     return reason;
 }
 
-static void gc_zlib(ErlNifEnv *env, void* data) {
-    zlib_data_t *d = (zlib_data_t*)data;
-
+static void zlib_internal_close(zlib_data_t *d) {
     if(d->state == ST_DEFLATE) {
         deflateEnd(&d->s);
     } else if(d->state == ST_INFLATE) {
@@ -279,8 +293,6 @@ static void gc_zlib(ErlNifEnv *env, void* data) {
     }
 
     if(d->state != ST_CLOSED) {
-        enif_ioq_destroy(d->input_queue);
-
         if(d->stash_env != NULL) {
             enif_free_env(d->stash_env);
         }
@@ -289,17 +301,36 @@ static void gc_zlib(ErlNifEnv *env, void* data) {
     }
 }
 
+static void gc_zlib(ErlNifEnv *env, void* data) {
+    zlib_data_t *d = (zlib_data_t*)data;
+
+    enif_mutex_destroy(d->controller_lock);
+    enif_ioq_destroy(d->input_queue);
+
+    zlib_internal_close(d);
+
+    (void)env;
+}
+
 static int get_zlib_data(ErlNifEnv *env, ERL_NIF_TERM opaque, zlib_data_t **d) {
     return enif_get_resource(env, opaque, rtype_zlib, (void **)d);
 }
 
 static int zlib_process_check(ErlNifEnv *env, zlib_data_t *d) {
+    int is_controlling_process;
     ErlNifPid current_process;
 
     enif_self(env, &current_process);
 
-    return enif_is_identical(enif_make_pid(env, &current_process),
+    enif_mutex_lock(d->controller_lock);
+
+    is_controlling_process = enif_is_identical(
+        enif_make_pid(env, &current_process),
         enif_make_pid(env, &d->controlling_process));
+
+    enif_mutex_unlock(d->controller_lock);
+
+    return is_controlling_process;
 }
 
 static void zlib_reset_input(zlib_data_t *d) {
@@ -504,14 +535,18 @@ static ERL_NIF_TERM zlib_open(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[
 
     d->input_queue = enif_ioq_create(ERL_NIF_IOQ_NORMAL);
 
+    d->controller_lock = enif_mutex_create("zlib_controller_lock");
+
     d->s.zalloc = zlib_alloc;
     d->s.zfree  = zlib_free;
     d->s.opaque = d;
     d->s.data_type = Z_BINARY;
 
-    d->state = ST_NONE;
+    d->eos_behavior = EOS_BEHAVIOR_CUT;
     d->eos_seen = 0;
 
+    d->state = ST_NONE;
+
     d->want_output_crc = 0;
     d->want_input_crc = 0;
     d->is_raw_stream = 0;
@@ -542,50 +577,35 @@ static ERL_NIF_TERM zlib_close(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv
         return enif_raise_exception(env, am_not_initialized);
     }
 
-    gc_zlib(env, d);
+    zlib_internal_close(d);
 
     return am_ok;
 }
 
-/* deflate */
-
-static ERL_NIF_TERM zlib_deflateInit(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
+static ERL_NIF_TERM zlib_set_controller(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
     zlib_data_t *d;
-    int level, res;
 
-    if(argc != 2 || !get_zlib_data(env, argv[0], &d) ||
-                    !enif_get_int(env, argv[1], &level)) {
+    ErlNifPid new_owner;
+
+    if(argc != 2 || !get_zlib_data(env, argv[0], &d)
+                 || !enif_get_local_pid(env, argv[1], &new_owner)) {
         return enif_make_badarg(env);
     } else if(!zlib_process_check(env, d)) {
         return enif_raise_exception(env, am_not_on_controlling_process);
-    } else if(d->state != ST_NONE) {
-        return enif_raise_exception(env, am_already_initialized);
     }
 
-    res = deflateInit(&d->s, level);
+    enif_mutex_lock(d->controller_lock);
 
-    if(res == Z_OK) {
-        d->state = ST_DEFLATE;
-        d->eos_seen = 0;
+    d->controlling_process = new_owner;
 
-        /* FIXME: crc32/1 is documented as returning "the current calculated
-         * checksum," but failed to mention that the old implementation only
-         * calculated it when WindowBits < 0 (See zlib_deflateInit2).
-         *
-         * We could fix this behavior by setting d->want_input_crc to 1 here,
-         * but we've decided to retain this quirk since the performance hit is
-         * quite significant. */
-        d->want_output_crc = 0;
-        d->want_input_crc = 0;
-
-        d->output_crc = crc32(0L, Z_NULL, 0);
-        d->input_crc = crc32(0L, Z_NULL, 0);
-    }
+    enif_mutex_unlock(d->controller_lock);
 
-    return zlib_return(env, res);
+    return am_ok;
 }
 
-static ERL_NIF_TERM zlib_deflateInit2(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
+/* deflate */
+
+static ERL_NIF_TERM zlib_deflateInit(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
     zlib_data_t *d;
     int level, method, windowBits, memLevel, strategy, res;
 
@@ -697,7 +717,9 @@ static ERL_NIF_TERM zlib_deflateEnd(ErlNifEnv *env, int argc, const ERL_NIF_TERM
 
 static ERL_NIF_TERM zlib_deflateParams(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
     zlib_data_t *d;
+
     int res, level, strategy;
+    Bytef dummy_buffer;
 
     if(argc != 3 || !get_zlib_data(env, argv[0], &d)
                  || !enif_get_int(env, argv[1], &level)
@@ -709,12 +731,27 @@ static ERL_NIF_TERM zlib_deflateParams(ErlNifEnv *env, int argc, const ERL_NIF_T
         return enif_raise_exception(env, am_not_initialized);
     }
 
-    /* deflateParams will flush everything currently in the stream, corrupting
-     * the heap unless it's empty. We therefore pretend to have a full output
-     * buffer, forcing a Z_BUF_ERROR if there's anything left to be flushed. */
-    d->s.avail_out = 0;
+    /* This is a bit of a hack; deflateParams flushes with Z_BLOCK which won't
+     * stop at a byte boundary, so we can't split this operation up, and we
+     * can't allocate a buffer large enough to fit it in one go since we have
+     * to support zlib versions that lack deflatePending.
+     *
+     * We therefore flush everything prior to this call to ensure that we are
+     * stopped on a byte boundary and have no pending data. We then hand it a
+     * dummy buffer to detect when this assumption doesn't hold (Hopefully
+     * never), and to smooth over an issue with zlib 1.2.11 which always
+     * returns Z_BUF_ERROR when d->s.avail_out is 0, regardless of whether
+     * there's any pending data or not. */
+
+    d->s.next_out = &dummy_buffer;
+    d->s.avail_out = 1;
+
     res = deflateParams(&d->s, level, strategy);
 
+    if(d->s.avail_out == 0) {
+        return zlib_return(env, Z_STREAM_ERROR);
+    }
+
     return zlib_return(env, res);
 }
 
@@ -741,39 +778,12 @@ static ERL_NIF_TERM zlib_deflate(ErlNifEnv *env, int argc, const ERL_NIF_TERM ar
 
 static ERL_NIF_TERM zlib_inflateInit(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
     zlib_data_t *d;
-    int res;
-
-    if(argc != 1 || !get_zlib_data(env, argv[0], &d)) {
-        return enif_make_badarg(env);
-    } else if(!zlib_process_check(env, d)) {
-        return enif_raise_exception(env, am_not_on_controlling_process);
-    } else if(d->state != ST_NONE) {
-        return enif_raise_exception(env, am_already_initialized);
-    }
 
-    res = inflateInit(&d->s);
+    int windowBits, eosBehavior, res;
 
-    if(res == Z_OK) {
-        d->state = ST_INFLATE;
-        d->eos_seen = 0;
-
-        d->want_output_crc = 0;
-        d->want_input_crc = 0;
-        d->is_raw_stream = 0;
-
-        d->output_crc = crc32(0L, Z_NULL, 0);
-        d->input_crc = crc32(0L, Z_NULL, 0);
-    }
-
-    return zlib_return(env, res);
-}
-
-static ERL_NIF_TERM zlib_inflateInit2(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
-    zlib_data_t *d;
-    int windowBits, res;
-
-    if(argc != 2 || !get_zlib_data(env, argv[0], &d)
-                 || !enif_get_int(env, argv[1], &windowBits)) {
+    if(argc != 3 || !get_zlib_data(env, argv[0], &d)
+                 || !enif_get_int(env, argv[1], &windowBits)
+                 || !enif_get_int(env, argv[2], &eosBehavior)) {
         return enif_make_badarg(env);
     } else if(!zlib_process_check(env, d)) {
         return enif_raise_exception(env, am_not_on_controlling_process);
@@ -785,6 +795,8 @@ static ERL_NIF_TERM zlib_inflateInit2(ErlNifEnv *env, int argc, const ERL_NIF_TE
 
     if(res == Z_OK) {
         d->state = ST_INFLATE;
+
+        d->eos_behavior = eosBehavior;
         d->eos_seen = 0;
 
         d->is_raw_stream = (windowBits < 0);
@@ -934,6 +946,27 @@ static ERL_NIF_TERM zlib_inflate(ErlNifEnv *env, int argc, const ERL_NIF_TERM ar
         return enif_raise_exception(env, am_not_initialized);
     }
 
+    if(d->eos_seen && enif_ioq_size(d->input_queue) > 0) {
+        int res;
+
+        switch(d->eos_behavior) {
+        case EOS_BEHAVIOR_ERROR:
+            return zlib_return(env, Z_DATA_ERROR);
+        case EOS_BEHAVIOR_RESET:
+            res = inflateReset(&d->s);
+
+            if(res != Z_OK) {
+                return zlib_return(env, res);
+            }
+
+            d->eos_seen = 0;
+
+            break;
+        case EOS_BEHAVIOR_CUT:
+            zlib_reset_input(d);
+        }
+    }
+
     return zlib_codec(&inflate, env, d, input_chunk_size, output_chunk_size, flush);
 }