From 89ee98864d7e4e9dbf1fe4ac5060fbe1867ef9b7 Mon Sep 17 00:00:00 2001
From: Fedor Indutnyy <indutny@signal.org>
Date: Wed, 10 Aug 2022 15:52:17 -0700
Subject: [PATCH] Process pending frames when handling socket end

When the socket emits `end` event we might still have to process the
frames from the bufferList. One of the frames could be CONNECTION_CLOSE
and we don't want to miss that and think that the server has terminated
connection incorrectly.
---
 lib/WebSocketConnection.js | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/lib/WebSocketConnection.js b/lib/WebSocketConnection.js
index 219de631..93d38000 100644
--- a/lib/WebSocketConnection.js
+++ b/lib/WebSocketConnection.js
@@ -271,7 +271,7 @@ WebSocketConnection.prototype.handleSocketData = function(data) {
     this.processReceivedData();
 };
 
-WebSocketConnection.prototype.processReceivedData = function() {
+WebSocketConnection.prototype.processReceivedData = function(isSync = false) {
     this._debug('processReceivedData');
     // If we're not connected, we should ignore any data remaining on the buffer.
     if (!this.connected) { return; }
@@ -320,7 +320,11 @@ WebSocketConnection.prototype.processReceivedData = function() {
         process.nextTick(function() { self.emit('frame', frame); });
     }
 
-    process.nextTick(function() { self.processFrame(frame); });
+    if (isSync) {
+        self.processFrame(frame);
+    } else {
+        process.nextTick(function() { self.processFrame(frame); });
+    }
     
     this.currentFrame = new WebSocketFrame(this.maskBytes, this.frameHeader, this.config);
 
@@ -329,7 +333,11 @@ 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);
+        if (isSync) {
+            this.receivedDataHandler();
+        } else {
+            setImmediateImpl(this.receivedDataHandler);
+        }
     }
 };
 
@@ -353,6 +361,11 @@ WebSocketConnection.prototype.handleSocketError = function(error) {
 };
 
 WebSocketConnection.prototype.handleSocketEnd = function() {
+    // We might have socket data scheduled for a next tick, process it now.
+    if (this.bufferList.length > 0) {
+        this.receivedDataHandler(true);
+    }
+
     this._debug('handleSocketEnd: received socket end.  state = %s', this.state);
     this.receivedEnd = true;
     if (this.state === STATE_CLOSED) {