[fix] Use status code from close frame if received

Author: lpinca

lib/EventTarget.js

@@ -55,7 +55,7 @@ class CloseEvent extends Event {
   constructor (code, reason, target) {
     super('close', target);
 
-    this.wasClean = code === undefined || code === 1000 || (code >= 3000 && code <= 4999);
+    this.wasClean = target._closeFrameReceived && target._closeFrameSent;
     this.reason = reason;
     this.code = code;
   }

lib/WebSocket.js

@@ -55,11 +55,12 @@ class WebSocket extends EventEmitter {
 
     this._binaryType = constants.BINARY_TYPES[0];
     this._finalize = this.finalize.bind(this);
+    this._closeFrameReceived = false;
     this._closeFrameSent = false;
-    this._closeMessage = null;
+    this._closeMessage = '';
     this._closeTimer = null;
     this._finalized = false;
-    this._closeCode = null;
+    this._closeCode = 1006;
     this._receiver = null;
     this._sender = null;
     this._socket = null;
@@ -126,36 +127,40 @@ class WebSocket extends EventEmitter {
     this._ultron = new Ultron(socket);
     this._socket = socket;
 
-    // socket cleanup handlers
     this._ultron.on('close', this._finalize);
     this._ultron.on('error', this._finalize);
     this._ultron.on('end', this._finalize);
 
-    // ensure that the head is added to the receiver
     if (head.length > 0) socket.unshift(head);
 
-    // subsequent packets are pushed to the receiver
     this._ultron.on('data', (data) => {
       this.bytesReceived += data.length;
       this._receiver.add(data);
     });
 
-    // receiver event handlers
     this._receiver.onmessage = (data) => this.emit('message', data);
     this._receiver.onping = (data) => {
       this.pong(data, !this._isServer, true);
       this.emit('ping', data);
     };
     this._receiver.onpong = (data) => this.emit('pong', data);
     this._receiver.onclose = (code, reason) => {
+      this._closeFrameReceived = true;
       this._closeMessage = reason;
       this._closeCode = code;
       if (!this._finalized) this.close(code, reason);
     };
     this._receiver.onerror = (error, code) => {
-      // close the connection when the receiver reports a HyBi error code
-      this.close(code, '');
+      this._closeMessage = '';
+      this._closeCode = code;
+
+      //
+      // Ensure that the error is emitted even if `WebSocket#finalize()` has
+      // already been called.
+      //
+      this.readyState = WebSocket.CLOSING;
       this.emit('error', error);
+      this.finalize(true);
     };
 
     this.readyState = WebSocket.OPEN;
@@ -174,7 +179,8 @@ class WebSocket extends EventEmitter {
     this.readyState = WebSocket.CLOSING;
     this._finalized = true;
 
-    if (!this._socket) return this.emitClose(error);
+    if (typeof error === 'object') this.emit('error', error);
+    if (!this._socket) return this.emitClose();
 
     clearTimeout(this._closeTimer);
     this._closeTimer = null;
@@ -190,29 +196,19 @@ class WebSocket extends EventEmitter {
     this._socket = null;
     this._sender = null;
 
-    this._receiver.cleanup(() => this.emitClose(error));
+    this._receiver.cleanup(() => this.emitClose());
     this._receiver = null;
   }
 
   /**
    * Emit the `close` event.
    *
-   * @param {(Boolean|Error)} error Indicates whether or not an error occurred
    * @private
    */
-  emitClose (error) {
+  emitClose () {
     this.readyState = WebSocket.CLOSED;
 
-    //
-    // If the connection was closed abnormally (with an error), or if the close
-    // control frame was not sent or received then the close code must be 1006.
-    //
-    if (error || !this._closeFrameSent) {
-      this._closeMessage = '';
-      this._closeCode = 1006;
-    }
-
-    this.emit('close', this._closeCode || 1006, this._closeMessage || '');
+    this.emit('close', this._closeCode, this._closeMessage);
 
     if (this.extensions[PerMessageDeflate.extensionName]) {
       this.extensions[PerMessageDeflate.extensionName].cleanup();
@@ -271,37 +267,35 @@ class WebSocket extends EventEmitter {
   close (code, data) {
     if (this.readyState === WebSocket.CLOSED) return;
     if (this.readyState === WebSocket.CONNECTING) {
-      if (this._req && !this._req.aborted) {
-        this._req.abort();
-        this.emit('error', new Error('closed before the connection is established'));
-        this.finalize(true);
-      }
+      this._req.abort();
+      this.finalize(new Error('closed before the connection is established'));
       return;
     }
 
     if (this.readyState === WebSocket.CLOSING) {
-      if (this._closeFrameSent && this._closeCode) this._socket.end();
+      if (this._closeFrameSent && this._closeFrameReceived) this._socket.end();
       return;
     }
 
     this.readyState = WebSocket.CLOSING;
     this._sender.close(code, data, !this._isServer, (err) => {
-      if (this._finalized) return;
-
-      if (err) {
-        this.emit('error', err);
-        this.finalize(true);
-        return;
-      }
+      //
+      // This error is handled by the `'error'` listener on the socket. We only
+      // want to know if the close frame has been sent here.
+      //
+      if (err) return;
 
-      if (this._closeCode) this._socket.end();
       this._closeFrameSent = true;
 
-      //
-      // Ensure that the connection is cleaned up even when the closing
-      // handshake fails.
-      //
-      this._closeTimer = setTimeout(this._finalize, closeTimeout, true);
+      if (!this._finalized) {
+        if (this._closeFrameReceived) this._socket.end();
+
+        //
+        // Ensure that the connection is cleaned up even when the closing
+        // handshake fails.
+        //
+        this._closeTimer = setTimeout(this._finalize, closeTimeout, true);
+      }
     });
   }
 
@@ -391,11 +385,8 @@ class WebSocket extends EventEmitter {
   terminate () {
     if (this.readyState === WebSocket.CLOSED) return;
     if (this.readyState === WebSocket.CONNECTING) {
-      if (this._req && !this._req.aborted) {
-        this._req.abort();
-        this.emit('error', new Error('closed before the connection is established'));
-        this.finalize(true);
-      }
+      this._req.abort();
+      this.finalize(new Error('closed before the connection is established'));
       return;
     }
 
@@ -645,24 +636,21 @@ function initAsClient (address, protocols, options) {
   if (options.handshakeTimeout) {
     this._req.setTimeout(options.handshakeTimeout, () => {
       this._req.abort();
-      this.emit('error', new Error('opening handshake has timed out'));
-      this.finalize(true);
+      this.finalize(new Error('opening handshake has timed out'));
     });
   }
 
   this._req.on('error', (error) => {
     if (this._req.aborted) return;
 
     this._req = null;
-    this.emit('error', error);
-    this.finalize(true);
+    this.finalize(error);
   });
 
   this._req.on('response', (res) => {
     if (!this.emit('unexpected-response', this._req, res)) {
       this._req.abort();
-      this.emit('error', new Error(`unexpected server response (${res.statusCode})`));
-      this.finalize(true);
+      this.finalize(new Error(`unexpected server response (${res.statusCode})`));
     }
   });
 
@@ -683,8 +671,7 @@ function initAsClient (address, protocols, options) {
 
     if (res.headers['sec-websocket-accept'] !== digest) {
       socket.destroy();
-      this.emit('error', new Error('invalid server key'));
-      return this.finalize(true);
+      return this.finalize(new Error('invalid server key'));
     }
 
     const serverProt = res.headers['sec-websocket-protocol'];
@@ -701,8 +688,7 @@ function initAsClient (address, protocols, options) {
 
     if (protError) {
       socket.destroy();
-      this.emit('error', new Error(protError));
-      return this.finalize(true);
+      return this.finalize(new Error(protError));
     }
 
     if (serverProt) this.protocol = serverProt;
@@ -721,8 +707,8 @@ function initAsClient (address, protocols, options) {
         }
       } catch (err) {
         socket.destroy();
-        this.emit('error', new Error('invalid Sec-WebSocket-Extensions header'));
-        return this.finalize(true);
+        this.finalize(new Error('invalid Sec-WebSocket-Extensions header'));
+        return;
       }
     }