PerMessageDeflate#cleanup - properly cleanup on close

Author: yosiat

lib/permessage-deflate.js

@@ -133,6 +133,10 @@ class PerMessageDeflate {
     }
 
     if (this._deflate) {
+      if (this._deflate[kCallback]) {
+        this._deflate[kCallback]();
+      }
+
       this._deflate.close();
       this._deflate = null;
     }
@@ -312,7 +316,9 @@ class PerMessageDeflate {
     zlibLimiter.push((done) => {
       this._compress(data, fin, (err, result) => {
         done();
-        callback(err, result);
+        if (err || result) {
+          callback(err, result);
+        }
       });
     });
   }
@@ -419,6 +425,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 +446,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;