fix: only unassign currentMessage after fully processing it

Author: arthurschreiber

src/incoming-message-stream.js

@@ -64,23 +64,21 @@ class IncomingMessageStream extends Transform {
         }
 
         if (packet.isLast()) {
-          this.currentMessage = undefined;
           // Wait until the current message was fully processed before we
           // continue processing any remaining messages.
           message.once('end', () => {
+            this.currentMessage = undefined;
             this.processBufferedData(callback);
           });
           message.end(packet.data());
           return;
-        } else {
+        } else if (!message.write(packet.data())) {
           // If too much data is buffering up in the
           // current message, wait for it to drain.
-          if (!message.write(packet.data())) {
-            message.once('drain', () => {
-              this.processBufferedData(callback);
-            });
-            return;
-          }
+          message.once('drain', () => {
+            this.processBufferedData(callback);
+          });
+          return;
         }
       } else {
         break;

test/unit/incoming-message-stream-test.js

@@ -106,4 +106,73 @@ describe('IncomingMessageStream', function() {
       });
     });
   });
+
+  it('correctly handles the last package coming in after the stream was paused', function(done) {
+    const packetData = Buffer.from('test1234');
+    const packetHeader = Buffer.alloc(8);
+
+    let offset = 0;
+    offset = packetHeader.writeUInt8(0x11, offset);
+    offset = packetHeader.writeUInt8(0x00, offset);
+    offset = packetHeader.writeUInt16BE(8 + packetData.length, offset);
+    offset = packetHeader.writeUInt16BE(0x0000, offset);
+    offset = packetHeader.writeUInt8(1, offset);
+    packetHeader.writeUInt8(0x00, offset);
+
+    const firstPacket = Buffer.concat([packetHeader, packetData]);
+
+    offset = 0;
+    offset = packetHeader.writeUInt8(0x11, offset);
+    offset = packetHeader.writeUInt8(0x01, offset);
+    offset = packetHeader.writeUInt16BE(8 + packetData.length, offset);
+    offset = packetHeader.writeUInt16BE(0x0000, offset);
+    offset = packetHeader.writeUInt8(1, offset);
+    packetHeader.writeUInt8(0x00, offset);
+
+    const secondPacket = Buffer.concat([packetHeader, packetData]);
+
+    const incoming = new IncomingMessageStream(new Debug());
+
+    const result = new BufferList(function(err, res) {
+      if (err) {
+        return done(err);
+      }
+
+      assert.deepEqual(res, Buffer.concat([ packetData, packetData ]));
+
+      done();
+    });
+
+    let messageEnded = false;
+    incoming.on('data', function(message) {
+      assert.instanceOf(message, Message);
+
+      message.on('end', function() {
+        messageEnded = true;
+      });
+
+      message.pipe(result);
+    });
+
+    incoming.write(firstPacket, function() {
+      const writtenData = result.slice();
+
+      assert.strictEqual(writtenData.length, 8);
+      assert.deepEqual(writtenData, packetData);
+
+      incoming.pause();
+
+      incoming.write(secondPacket, function() {
+        const writtenData = result.slice();
+
+        assert.strictEqual(writtenData.length, 16);
+        assert.deepEqual(writtenData, Buffer.concat([ packetData, packetData ]));
+
+        assert.strictEqual(messageEnded, true);
+      });
+
+      assert.isFalse(messageEnded);
+      incoming.resume();
+    });
+  });
 });