[minor] Refactor server client initialization

websockets · Jan 5, 2018 · 46461a9 · 46461a9
1 parent ee9b5f3
commit 46461a9
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 60 deletions.
diff --git a/lib/websocket-server.js b/lib/websocket-server.js
@@ -249,13 +249,19 @@ class WebSocketServer extends EventEmitter {
       `Sec-WebSocket-Accept: ${key}`
     ];
 
-    if (protocol) headers.push(`Sec-WebSocket-Protocol: ${protocol}`);
+    const ws = new WebSocket(null);
+
+    if (protocol) {
+      headers.push(`Sec-WebSocket-Protocol: ${protocol}`);
+      ws.protocol = protocol;
+    }
     if (extensions[PerMessageDeflate.extensionName]) {
       const params = extensions[PerMessageDeflate.extensionName].params;
       const value = extension.format({
         [PerMessageDeflate.extensionName]: [params]
       });
       headers.push(`Sec-WebSocket-Extensions: ${value}`);
+      ws.extensions = extensions;
     }
 
     //
@@ -264,20 +270,16 @@ class WebSocketServer extends EventEmitter {
     this.emit('headers', headers, req);
 
     socket.write(headers.concat('\r\n').join('\r\n'));
+    socket.removeListener('error', socketError);
 
-    const client = new WebSocket([socket, head], null, {
-      maxPayload: this.options.maxPayload,
-      extensions,
-      protocol
-    });
+    ws.setSocket(socket, head, this.options.maxPayload);
 
     if (this.clients) {
-      this.clients.add(client);
-      client.on('close', () => this.clients.delete(client));
+      this.clients.add(ws);
+      ws.on('close', () => this.clients.delete(ws));
     }
 
-    socket.removeListener('error', socketError);
-    cb(client);
+    cb(ws);
   }
 }
 

diff --git a/lib/websocket.js b/lib/websocket.js
@@ -40,15 +40,6 @@ class WebSocket extends EventEmitter {
   constructor (address, protocols, options) {
     super();
 
-    if (!protocols) {
-      protocols = [];
-    } else if (typeof protocols === 'string') {
-      protocols = [protocols];
-    } else if (!Array.isArray(protocols)) {
-      options = protocols;
-      protocols = [];
-    }
-
     this.readyState = WebSocket.CONNECTING;
     this.extensions = {};
     this.protocol = '';
@@ -61,14 +52,22 @@ class WebSocket extends EventEmitter {
     this._closeTimer = null;
     this._finalized = false;
     this._closeCode = 1006;
+    this._isServer = true;
     this._receiver = null;
     this._sender = null;
     this._socket = null;
     this._ultron = null;
 
-    if (Array.isArray(address)) {
-      initAsServerClient.call(this, address[0], address[1], options);
-    } else {
+    if (address !== null) {
+      if (!protocols) {
+        protocols = [];
+      } else if (typeof protocols === 'string') {
+        protocols = [protocols];
+      } else if (!Array.isArray(protocols)) {
+        options = protocols;
+        protocols = [];
+      }
+
       initAsClient.call(this, address, protocols, options);
     }
   }
@@ -116,13 +115,14 @@ class WebSocket extends EventEmitter {
    *
    * @param {net.Socket} socket The network socket between the server and client
    * @param {Buffer} head The first packet of the upgraded stream
+   * @param {Number} maxPayload The maximum allowed message size
    * @private
    */
-  setSocket (socket, head) {
+  setSocket (socket, head, maxPayload) {
     socket.setTimeout(0);
     socket.setNoDelay();
 
-    this._receiver = new Receiver(this.extensions, this._maxPayload, this.binaryType);
+    this._receiver = new Receiver(this.extensions, maxPayload, this.binaryType);
     this._sender = new Sender(socket, this.extensions);
     this._ultron = new Ultron(socket);
     this._socket = socket;
@@ -485,28 +485,6 @@ WebSocket.prototype.removeEventListener = EventTarget.removeEventListener;
 
 module.exports = WebSocket;
 
-/**
- * Initialize a WebSocket server client.
- *
- * @param {http.IncomingMessage} req The request object
- * @param {net.Socket} socket The network socket between the server and client
- * @param {Buffer} head The first packet of the upgraded stream
- * @param {Object} options WebSocket attributes
- * @param {Object} options.extensions The negotiated extensions
- * @param {Number} options.maxPayload The maximum allowed message size
- * @param {String} options.protocol The chosen subprotocol
- * @private
- */
-function initAsServerClient (socket, head, options) {
-  this._maxPayload = options.maxPayload;
-  this.extensions = options.extensions;
-  this.protocol = options.protocol;
-
-  this._isServer = true;
-
-  this.setSocket(socket, head);
-}
-
 /**
  * Initialize a WebSocket client.
  *
@@ -753,6 +731,6 @@ function initAsClient (address, protocols, options) {
       }
     }
 
-    this.setSocket(socket, head);
+    this.setSocket(socket, head, 0);
   });
 }
diff --git a/test/websocket-server.test.js b/test/websocket-server.test.js
@@ -250,19 +250,6 @@ describe('WebSocketServer', function () {
   });
 
   describe('#maxpayload', function () {
-    it('maxpayload is passed on to clients', function (done) {
-      const maxPayload = 20480;
-      const wss = new WebSocket.Server({ port: 0, maxPayload }, () => {
-        const port = wss._server.address().port;
-        const ws = new WebSocket(`ws://localhost:${port}`);
-      });
-
-      wss.on('connection', (client) => {
-        assert.strictEqual(client._maxPayload, maxPayload);
-        wss.close(done);
-      });
-    });
-
     it('maxpayload is passed on to hybi receivers', function (done) {
       const maxPayload = 20480;
       const wss = new WebSocket.Server({ port: 0, maxPayload }, () => {