[minor] Add ability to specify the masking key

Refs: #1986
Refs: #1988

Author: lpinca

doc/ws.md

@@ -462,8 +462,9 @@ is a noop if the ready state is `CONNECTING` or `CLOSED`.
 
 - `data` {Array|Number|Object|String|ArrayBuffer|Buffer|DataView|TypedArray} The
   data to send in the ping frame.
-- `mask` {Boolean} Specifies whether `data` should be masked or not. Defaults to
-  `true` when `websocket` is not a server client.
+- `mask` {Boolean|Buffer} Specifies whether `data` should be masked or not. If a
+  `Buffer` is provided, then its contents is used as the masking key. Defaults
+  to `true` when `websocket` is not a server client.
 - `callback` {Function} An optional callback which is invoked when the ping
   frame is written out. If an error occurs, the callback is called with the
   error as its first argument.
@@ -474,8 +475,9 @@ Send a ping. This method throws an error if the ready state is `CONNECTING`.
 
 - `data` {Array|Number|Object|String|ArrayBuffer|Buffer|DataView|TypedArray} The
   data to send in the pong frame.
-- `mask` {Boolean} Specifies whether `data` should be masked or not. Defaults to
-  `true` when `websocket` is not a server client.
+- `mask` {Boolean|Buffer} Specifies whether `data` should be masked or not. If a
+  `Buffer` is provided, then its contents is used as the masking key. Defaults
+  to `true` when `websocket` is not a server client.
 - `callback` {Function} An optional callback which is invoked when the pong
   frame is written out. If an error occurs, the callback is called with the
   error as its first argument.
@@ -519,8 +521,9 @@ only removes listeners added with
     Defaults to `true` when permessage-deflate is enabled.
   - `fin` {Boolean} Specifies whether `data` is the last fragment of a message
     or not. Defaults to `true`.
-  - `mask` {Boolean} Specifies whether `data` should be masked or not. Defaults
-    to `true` when `websocket` is not a server client.
+  - `mask` {Boolean|Buffer} Specifies whether `data` should be masked or not. If
+    a `Buffer` is provided, then its contents is used as the masking key.
+    Defaults to `true` when `websocket` is not a server client.
 - `callback` {Function} An optional callback which is invoked when `data` is
   written out. If an error occurs, the callback is called with the error as its
   first argument.

lib/receiver.js

@@ -417,7 +417,12 @@ class Receiver extends Writable {
       }
 
       data = this.consume(this._payloadLength);
-      if (this._masked) unmask(data, this._mask);
+      if (
+        this._masked &&
+        (this._mask[0] | this._mask[1] | this._mask[2] | this._mask[3]) !== 0
+      ) {
+        unmask(data, this._mask);
+      }
     }
 
     if (this._opcode > 0x07) return this.controlMessage(data);

lib/sender.js

@@ -11,7 +11,7 @@ const { EMPTY_BUFFER } = require('./constants');
 const { isValidStatusCode } = require('./validation');
 const { mask: applyMask, toBuffer } = require('./buffer-util');
 
-const mask = Buffer.alloc(4);
+const _mask = Buffer.alloc(4);
 
 /**
  * HyBi Sender implementation.
@@ -42,8 +42,9 @@ class Sender {
    * @param {Object} options Options object
    * @param {Boolean} [options.fin=false] Specifies whether or not to set the
    *     FIN bit
-   * @param {Boolean} [options.mask=false] Specifies whether or not to mask
-   *     `data`
+   * @param {(Boolean|Buffer)} [options.mask=false] Specifies whether or not to
+   *     mask `data`. If a `Buffer` is provided, then its contents is used as
+   *     the masking key
    * @param {Number} options.opcode The opcode
    * @param {Boolean} [options.readOnly=false] Specifies whether `data` can be
    *     modified
@@ -81,14 +82,24 @@ class Sender {
 
     if (!options.mask) return [target, data];
 
-    randomFillSync(mask, 0, 4);
+    let mask = _mask;
+
+    if (Buffer.isBuffer(options.mask)) {
+      mask = options.mask;
+    } else {
+      randomFillSync(_mask, 0, 4);
+    }
 
     target[1] |= 0x80;
     target[offset - 4] = mask[0];
     target[offset - 3] = mask[1];
     target[offset - 2] = mask[2];
     target[offset - 1] = mask[3];
 
+    if ((mask[0] | mask[1] | mask[2] | mask[3]) === 0) {
+      return [target, data];
+    }
+
     if (merge) {
       applyMask(data, mask, target, offset, data.length);
       return [target];
@@ -103,7 +114,9 @@ class Sender {
    *
    * @param {Number} [code] The status code component of the body
    * @param {(String|Buffer)} [data] The message component of the body
-   * @param {Boolean} [mask=false] Specifies whether or not to mask the message
+   * @param {(Boolean|Buffer)} [mask=false] Specifies whether or not to mask the
+   *     message. If a `Buffer` is provided, then its contents is used as the
+   *     masking key
    * @param {Function} [cb] Callback
    * @public
    */
@@ -145,7 +158,9 @@ class Sender {
    * Frames and sends a close message.
    *
    * @param {Buffer} data The message to send
-   * @param {Boolean} [mask=false] Specifies whether or not to mask `data`
+   * @param {(Boolean|Buffer)} [mask=false] Specifies whether or not to mask
+   *     `data`. If a `Buffer` is provided, then its contents is used as the
+   *     masking key
    * @param {Function} [cb] Callback
    * @private
    */
@@ -166,7 +181,9 @@ class Sender {
    * Sends a ping message to the other peer.
    *
    * @param {*} data The message to send
-   * @param {Boolean} [mask=false] Specifies whether or not to mask `data`
+   * @param {(Boolean|Buffer)} [mask=false] Specifies whether or not to mask
+   *     `data`. If a `Buffer` is provided, then its contents is used as the
+   *     masking key
    * @param {Function} [cb] Callback
    * @public
    */
@@ -188,7 +205,9 @@ class Sender {
    * Frames and sends a ping message.
    *
    * @param {Buffer} data The message to send
-   * @param {Boolean} [mask=false] Specifies whether or not to mask `data`
+   * @param {(Boolean|Buffer)} [mask=false] Specifies whether or not to mask
+   *     `data`. If a `Buffer` is provided, then its contents is used as the
+   *     masking key
    * @param {Boolean} [readOnly=false] Specifies whether `data` can be modified
    * @param {Function} [cb] Callback
    * @private
@@ -210,7 +229,9 @@ class Sender {
    * Sends a pong message to the other peer.
    *
    * @param {*} data The message to send
-   * @param {Boolean} [mask=false] Specifies whether or not to mask `data`
+   * @param {(Boolean|Buffer)} [mask=false] Specifies whether or not to mask
+   *     `data`. If a `Buffer` is provided, then its contents is used as the
+   *     masking key
    * @param {Function} [cb] Callback
    * @public
    */
@@ -232,7 +253,9 @@ class Sender {
    * Frames and sends a pong message.
    *
    * @param {Buffer} data The message to send
-   * @param {Boolean} [mask=false] Specifies whether or not to mask `data`
+   * @param {(Boolean|Buffer)} [mask=false] Specifies whether or not to mask
+   *     `data`. If a `Buffer` is provided, then its contents is used as the
+   *     masking key
    * @param {Boolean} [readOnly=false] Specifies whether `data` can be modified
    * @param {Function} [cb] Callback
    * @private
@@ -261,8 +284,9 @@ class Sender {
    *     compress `data`
    * @param {Boolean} [options.fin=false] Specifies whether the fragment is the
    *     last one
-   * @param {Boolean} [options.mask=false] Specifies whether or not to mask
-   *     `data`
+   * @param {(Boolean|Buffer)} [options.mask=false] Specifies whether or not to
+   *     mask `data`. If a `Buffer` is provided, then its contents is used as
+   *     the masking key
    * @param {Function} [cb] Callback
    * @public
    */
@@ -331,8 +355,9 @@ class Sender {
    * @param {Number} options.opcode The opcode
    * @param {Boolean} [options.fin=false] Specifies whether or not to set the
    *     FIN bit
-   * @param {Boolean} [options.mask=false] Specifies whether or not to mask
-   *     `data`
+   * @param {(Boolean|Buffer)} [options.mask=false] Specifies whether or not to
+   *     mask `data`. If a `Buffer` is provided, then its contents is used as
+   *     the masking key
    * @param {Boolean} [options.readOnly=false] Specifies whether `data` can be
    *     modified
    * @param {Boolean} [options.rsv1=false] Specifies whether or not to set the
@@ -348,6 +373,7 @@ class Sender {
 
     const perMessageDeflate = this._extensions[PerMessageDeflate.extensionName];
 
+    if (options.mask && options.mask.length) this._bufferedBytes += 4;
     this._bufferedBytes += data.length;
     this._deflating = true;
     perMessageDeflate.compress(data, options.fin, (_, buf) => {
@@ -367,6 +393,7 @@ class Sender {
         return;
       }
 
+      if (options.mask && options.mask.length) this._bufferedBytes -= 4;
       this._bufferedBytes -= data.length;
       this._deflating = false;
       options.readOnly = false;
@@ -383,7 +410,9 @@ class Sender {
   dequeue() {
     while (!this._deflating && this._queue.length) {
       const params = this._queue.shift();
+      const mask = params[0] === this.dispatch ? params[3].mask : params[2];
 
+      if (mask && mask.length) this._bufferedBytes -= 4;
       this._bufferedBytes -= params[1].length;
       Reflect.apply(params[0], this, params.slice(1));
     }
@@ -396,7 +425,11 @@ class Sender {
    * @private
    */
   enqueue(params) {
+    const mask = params[0] === this.dispatch ? params[3].mask : params[2];
+
+    if (mask && mask.length) this._bufferedBytes += 4;
     this._bufferedBytes += params[1].length;
+
     this._queue.push(params);
   }
 

lib/websocket.js

@@ -341,7 +341,8 @@ class WebSocket extends EventEmitter {
    * Send a ping.
    *
    * @param {*} [data] The data to send
-   * @param {Boolean} [mask] Indicates whether or not to mask `data`
+   * @param {(Boolean|Buffer)} [mask] Indicates whether or not to mask `data`.
+   *     If a `Buffer` is provided, then its contents is used as the masking key
    * @param {Function} [cb] Callback which is executed when the ping is sent
    * @public
    */
@@ -373,7 +374,8 @@ class WebSocket extends EventEmitter {
    * Send a pong.
    *
    * @param {*} [data] The data to send
-   * @param {Boolean} [mask] Indicates whether or not to mask `data`
+   * @param {(Boolean|Buffer)} [mask] Indicates whether or not to mask `data`.
+   *     If a `Buffer` is provided, then its contents is used as the masking key
    * @param {Function} [cb] Callback which is executed when the pong is sent
    * @public
    */
@@ -429,7 +431,9 @@ class WebSocket extends EventEmitter {
    *     `data`
    * @param {Boolean} [options.fin=true] Specifies whether the fragment is the
    *     last one
-   * @param {Boolean} [options.mask] Specifies whether or not to mask `data`
+   * @param {(Boolean|Buffer)} [options.mask] Specifies whether or not to mask
+   *     `data`. If a `Buffer` is provided, then its contents is used as the
+   *     masking key
    * @param {Function} [cb] Callback which is executed when data is written out
    * @public
    */

test/websocket.test.js

@@ -1299,10 +1299,11 @@ describe('WebSocket', () => {
 
     it('can send a ping with data', (done) => {
       const wss = new WebSocket.Server({ port: 0 }, () => {
+        const mask = Buffer.alloc(4);
         const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
 
         ws.on('open', () => {
-          ws.ping('hi', () => {
+          ws.ping('hi', mask, () => {
             ws.ping('hi', true);
             ws.close();
           });
@@ -1468,10 +1469,11 @@ describe('WebSocket', () => {
 
     it('can send a pong with data', (done) => {
       const wss = new WebSocket.Server({ port: 0 }, () => {
+        const mask = Buffer.alloc(4);
         const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
 
         ws.on('open', () => {
-          ws.pong('hi', () => {
+          ws.pong('hi', mask, () => {
             ws.pong('hi', true);
             ws.close();
           });
@@ -1877,7 +1879,7 @@ describe('WebSocket', () => {
       });
     });
 
-    it('honors the `mask` option', (done) => {
+    it('honors the `mask` option (1/2)', (done) => {
       let clientCloseEventEmitted = false;
       let serverClientCloseEventEmitted = false;
 
@@ -1921,6 +1923,51 @@ describe('WebSocket', () => {
         });
       });
     });
+
+    it('honors the `mask` option (2/2)', (done) => {
+      const mask1 = Buffer.from('00000000', 'hex');
+      const mask2 = Buffer.from('00000001', 'hex');
+      const wss = new WebSocket.Server(
+        { perMessageDeflate: true, port: 0 },
+        () => {
+          const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
+
+          ws.on('open', () => {
+            ws.send('foo', { mask: mask1 });
+            ws.send('bar', { mask: mask2 });
+
+            assert.strictEqual(ws.bufferedAmount, 14);
+
+            ws.on('close', (code, reason) => {
+              assert.strictEqual(code, 1005);
+              assert.deepStrictEqual(reason, EMPTY_BUFFER);
+
+              wss.close(done);
+            });
+          });
+        }
+      );
+
+      wss.on('connection', (ws) => {
+        const chunks = [];
+
+        ws._socket.prependListener('data', (chunk) => {
+          chunks.push(chunk);
+        });
+
+        ws.on('message', (message) => {
+          if (message.toString() === 'foo') return;
+
+          const data = Buffer.concat(chunks);
+          const length = data[1] & 0x7f;
+
+          assert.ok(data.slice(2, 6).equals(mask1));
+          assert.ok(data.slice(length + 8, length + 12).equals(mask2));
+
+          ws.close();
+        });
+      });
+    });
   });
 
   describe('#close', () => {