From 6a6ae04ef3514e2f8b1b0890649523d59dcf27f4 Mon Sep 17 00:00:00 2001
From: Luigi Pinca <luigipinca@gmail.com>
Date: Sat, 16 Dec 2017 18:54:35 +0100
Subject: [PATCH] [minor] Send the close status code only when necessary

---
 lib/Sender.js          | 23 ++++++++++++++++++-----
 test/Sender.test.js    |  2 +-
 test/WebSocket.test.js | 28 ++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/lib/Sender.js b/lib/Sender.js
index 53953d8d2..046a0e1d4 100644
--- 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
index 385011476..ac19d1ecd 100644
--- 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
index 32a4245b8..f293491b9 100644
--- 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;