PerMessageDeflate#cleanup - properly cleanup on close

Author: yosiat

lib/permessage-deflate.js

@@ -61,6 +61,8 @@ class PerMessageDeflate {
 
     this.params = null;
 
+    this._stopWriting = false;
+
     if (!zlibLimiter) {
       const concurrency =
         this._options.concurrencyLimit !== undefined
@@ -127,12 +129,18 @@ class PerMessageDeflate {
    * @public
    */
   cleanup() {
+    this._stopWriting = true;
+
     if (this._inflate) {
       this._inflate.close();
       this._inflate = null;
     }
 
     if (this._deflate) {
+      if (this._deflate[kCallback]) {
+        this._deflate[kCallback]();
+      }
+
       this._deflate.close();
       this._deflate = null;
     }
@@ -312,7 +320,9 @@ class PerMessageDeflate {
     zlibLimiter.push((done) => {
       this._compress(data, fin, (err, result) => {
         done();
-        callback(err, result);
+        if (!this._stopWriting) {
+          callback(err, result);
+        }
       });
     });
   }
@@ -419,6 +429,8 @@ class PerMessageDeflate {
       this._deflate.on('data', deflateOnData);
     }
 
+    this._deflate[kCallback] = callback;
+
     this._deflate.write(data);
     this._deflate.flush(zlib.Z_SYNC_FLUSH, () => {
       if (!this._deflate) {
@@ -438,6 +450,11 @@ class PerMessageDeflate {
 
       if (fin) data = data.slice(0, data.length - 4);
 
+      // Set callback to null so we won't call it twice if socket is closed
+      // (see PerMessageDeflate#cleanup)
+      // ideally, we would set it to `null` before calling, but we might set it to null
+      this._deflate[kCallback] = null;
+
       if (fin && this.params[`${endpoint}_no_context_takeover`]) {
         this._deflate.close();
         this._deflate = null;