zlib: reset() method for deflate/inflate streams

indutny · indutny · commit 71ae1753196f · 2012-01-12T03:19:19.000+06:00
* ammended test-zlib-dictionary to cover reusing streams
diff --git a/lib/zlib.js b/lib/zlib.js
@@ -309,6 +309,10 @@ Zlib.prototype.write = function write(chunk, cb) {
   return empty;
 };
 
+Zlib.prototype.reset = function reset() {
+  return this._binding.reset();
+};
+
 Zlib.prototype.flush = function flush(cb) {
   this._flush = binding.Z_SYNC_FLUSH;
   return this.write(cb);
diff --git a/src/node_zlib.cc b/src/node_zlib.cc
@@ -249,6 +249,17 @@ template <node_zlib_mode mode> class ZCtx : public ObjectWrap {
 
     Init(ctx, level, windowBits, memLevel, strategy,
          dictionary, dictionary_len);
+    SetDictionary(ctx);
+    return Undefined();
+  }
+
+  static Handle<Value> Reset(const Arguments &args) {
+    HandleScope scope;
+
+    ZCtx<mode> *ctx = ObjectWrap::Unwrap< ZCtx<mode> >(args.This());
+
+    Reset(ctx);
+    SetDictionary(ctx);
     return Undefined();
   }
 
@@ -304,23 +315,46 @@ template <node_zlib_mode mode> class ZCtx : public ObjectWrap {
     ctx->dictionary_ = reinterpret_cast<Bytef *>(dictionary);
     ctx->dictionary_len_ = dictionary_len;
 
-    if (dictionary != NULL) {
-      switch (mode) {
-        case DEFLATE:
-        case DEFLATERAW:
-          err = deflateSetDictionary(&ctx->strm_,
-                                     ctx->dictionary_,
-                                     dictionary_len);
-          break;
-        default:
-          break;
-      }
+    ctx->write_in_progress_ = false;
+    ctx->init_done_ = true;
+  }
+
+  static void SetDictionary(ZCtx* ctx) {
+    if (ctx->dictionary_ == NULL) return;
 
-      assert(err == Z_OK && "Failed to set dictionary");
+    int err;
+
+    switch (mode) {
+      case DEFLATE:
+      case DEFLATERAW:
+        err = deflateSetDictionary(&ctx->strm_,
+                                   ctx->dictionary_,
+                                   ctx->dictionary_len_);
+        break;
+      default:
+        break;
     }
 
-    ctx->write_in_progress_ = false;
-    ctx->init_done_ = true;
+    assert(err == Z_OK && "Failed to set dictionary");
+  }
+
+  static void Reset(ZCtx* ctx) {
+    int err;
+
+    switch (mode) {
+      case DEFLATE:
+      case DEFLATERAW:
+        err = deflateReset(&ctx->strm_);
+        break;
+      case INFLATE:
+      case INFLATERAW:
+        err = inflateReset(&ctx->strm_);
+        break;
+      default:
+        break;
+    }
+
+    assert(err == Z_OK && "Failed to reset stream");
   }
 
  private:
@@ -352,6 +386,7 @@ template <node_zlib_mode mode> class ZCtx : public ObjectWrap {
     z->InstanceTemplate()->SetInternalFieldCount(1); \
     NODE_SET_PROTOTYPE_METHOD(z, "write", ZCtx<mode>::Write); \
     NODE_SET_PROTOTYPE_METHOD(z, "init", ZCtx<mode>::Init); \
+    NODE_SET_PROTOTYPE_METHOD(z, "reset", ZCtx<mode>::Reset); \
     z->SetClassName(String::NewSymbol(name)); \
     target->Set(String::NewSymbol(name), z->GetFunction()); \
   }
diff --git a/test/simple/test-zlib-dictionary.js b/test/simple/test-zlib-dictionary.js
@@ -43,7 +43,6 @@ var spdyDict = new Buffer([
 ].join(''));
 
 var deflate = zlib.createDeflate({ dictionary: spdyDict });
-var inflate = zlib.createInflate({ dictionary: spdyDict });
 
 var input = [
   'HTTP/1.1 200 Ok',
@@ -52,22 +51,45 @@ var input = [
   ''
 ].join('\r\n');
 
-// Put data into deflate stream
-deflate.on('data', function(chunk) {
-  inflate.write(chunk);
-});
-deflate.on('end', function() {
-  inflate.end();
-});
+var called = 0;
 
-// Get data from inflate stream
-var output = [];
-inflate.on('data', function(chunk) {
-  output.push(chunk);
-});
-inflate.on('end', function() {
-  assert.equal(output.join(''), input);
-});
+//
+// We'll use clean-new inflate stream each time
+// and .reset() old dirty deflate one
+//
+function run(num) {
+  var inflate = zlib.createInflate({ dictionary: spdyDict });
+
+  if (num === 2) {
+    deflate.reset();
+    deflate.removeAllListeners('data');
+  }
 
-deflate.write(input);
-deflate.end();
+  // Put data into deflate stream
+  deflate.on('data', function(chunk) {
+    inflate.write(chunk);
+  });
+
+  // Get data from inflate stream
+  var output = [];
+  inflate.on('data', function(chunk) {
+    output.push(chunk);
+  });
+  inflate.on('end', function() {
+    called++;
+
+    assert.equal(output.join(''), input);
+
+    if (num < 2) run(num + 1);
+  });
+
+  deflate.write(input);
+  deflate.flush(function() {
+    inflate.end();
+  });
+}
+run(1);
+
+process.on('exit', function() {
+  assert.equal(called, 2);
+});
