[minor] Send the close status code only when necessary

websockets · Dec 16, 2017 · 6a6ae04 · 6a6ae04
1 parent 85919f2
commit 6a6ae04
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 6 deletions.
diff --git a/lib/Sender.js b/lib/Sender.js
@@ -12,6 +12,7 @@ const crypto = require('crypto');
 const PerMessageDeflate = require('./PerMessageDeflate');
 const bufferUtil = require('./BufferUtil');
 const ErrorCodes = require('./ErrorCodes');
+const constants = require('./Constants');
 
 const Buffer = safeBuffer.Buffer;
 
@@ -112,14 +113,26 @@ class Sender {
    * @public
    */
   close (code, data, mask, cb) {
-    if (code !== undefined && (typeof code !== 'number' || !ErrorCodes.isValidErrorCode(code))) {
+    var buf;
+
+    if (code === undefined) {
+      code = 1000;
+    } else if (typeof code !== 'number' || !ErrorCodes.isValidErrorCode(code)) {
       throw new Error('first argument must be a valid error code number');
     }
 
-    const buf = Buffer.allocUnsafe(2 + (data ? Buffer.byteLength(data) : 0));
-
-    buf.writeUInt16BE(code || 1000, 0, true);
-    if (buf.length > 2) buf.write(data, 2);
+    if (data === undefined || data === '') {
+      if (code === 1000) {
+        buf = constants.EMPTY_BUFFER;
+      } else {
+        buf = Buffer.allocUnsafe(2);
+        buf.writeUInt16BE(code, 0, true);
+      }
+    } else {
+      buf = Buffer.allocUnsafe(2 + Buffer.byteLength(data));
+      buf.writeUInt16BE(code, 0, true);
+      buf.write(data, 2);
+    }
 
     if (this._deflating) {
       this.enqueue([this.doClose, buf, mask, cb]);

diff --git a/test/Sender.test.js b/test/Sender.test.js
@@ -267,7 +267,7 @@ describe('Sender', function () {
       sender.send('bar', { compress: true, fin: true });
       sender.send('baz', { compress: true, fin: true });
 
-      sender.close(1000, null, false, () => {
+      sender.close(1000, undefined, false, () => {
         assert.strictEqual(count, 4);
         done();
       });

diff --git a/test/WebSocket.test.js b/test/WebSocket.test.js
@@ -1203,6 +1203,34 @@ describe('WebSocket', function () {
       });
     });
 
+    it('sends the close status code only when necessary', function (done) {
+      let sent;
+      const wss = new WebSocket.Server({ port: 0 }, () => {
+        const port = wss._server.address().port;
+        const ws = new WebSocket(`ws://localhost:${port}`);
+
+        ws.on('open', () => {
+          ws._socket.once('data', (data) => {
+            sent = data;
+          });
+        });
+      });
+
+      wss.on('connection', (ws) => {
+        ws._socket.once('data', (received) => {
+          assert.ok(received.slice(0, 2).equals(Buffer.from([0x88, 0x80])));
+          assert.ok(sent.equals(Buffer.from([0x88, 0x00])));
+
+          ws.on('close', (code, reason) => {
+            assert.strictEqual(code, 1000);
+            assert.strictEqual(reason, '');
+            wss.close(done);
+          });
+        });
+        ws.close();
+      });
+    });
+
     it('works when close reason is not specified', function (done) {
       const wss = new WebSocket.Server({ port: 0 }, () => {
         const port = wss._server.address().port;