From 8af1d9fd837f33f54999d2ab30180bd7cb56bc88 Mon Sep 17 00:00:00 2001
From: Mathieu Bruyen <code@mais-h.eu>
Date: Fri, 9 Sep 2016 22:28:24 +0200
Subject: [PATCH] Remove race condition when server closes connection

When server closes connection right after having sent some
frames, those are ignored by the client.

Frames are lost because they are processed in a
nextTick/setImmediate callback, thus socket close event
may arrive before they are actually processed.
---
 lib/WebSocketConnection.js |  3 ++-
 test/unit/request.js       | 54 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/lib/WebSocketConnection.js b/lib/WebSocketConnection.js
index fd87264e..aff6f9a4 100644
--- a/lib/WebSocketConnection.js
+++ b/lib/WebSocketConnection.js
@@ -327,7 +327,8 @@ WebSocketConnection.prototype.processReceivedData = function() {
     // processed.  We use setImmediate here instead of process.nextTick to
     // explicitly indicate that we wish for other I/O to be handled first.
     if (this.bufferList.length > 0) {
-        setImmediateImpl(this.receivedDataHandler);
+        this.receivedDataHandler();
+        //setImmediateImpl(this.receivedDataHandler);
     }
 };
 
diff --git a/test/unit/request.js b/test/unit/request.js
index f5cc69a4..a34f386f 100644
--- a/test/unit/request.js
+++ b/test/unit/request.js
@@ -103,3 +103,57 @@ test('Protocol mismatch should be handled gracefully', function(t) {
     t.end();
   });
 });
+
+test('Server pushing and closing should not loose data', function(t) {
+  var wsServer;
+  
+  t.test('setup', function(t) {
+    server.prepare(function(err, result) {
+      if (err) {
+        t.fail('Unable to start test server');
+        return t.end();
+      }
+      
+      wsServer = result;
+      t.end();
+    });
+  });
+  
+  t.test('server push and close', function(t) {
+    t.plan(2);
+    var count = 100;
+    var received = 0;
+    wsServer.on('request', handleRequest);
+    
+    var client = new WebSocketClient();
+    
+    var timer = setTimeout(function() {
+      t.fail('Timeout waiting for client event - received ' + received + ' frames out of ' + count);
+    }, 2000);
+    
+    client.connect('ws://localhost:64321/', 'echo');
+    client.on('connect', function(connection) {
+      connection.on('message', function (msg) {
+        received++;
+        if (count === received) {
+          clearTimeout(timer);
+          t.pass('received enough frames');
+        }
+      });
+    });
+    
+    function handleRequest(request) {
+      t.pass('connected');
+      var connection = request.accept('echo', request.origin);
+      for (var i = 0; i < count; i++) {
+        connection.sendUTF('some content');
+      }
+      connection.close();
+    }
+  });
+  
+  t.test('teardown', function(t) {
+    stopServer();
+    t.end();
+  });
+});