[feature] Introduce the generateMask option

The `generateMask` option specifies a function that can be used to generate custom masking keys. Refs: #1986 Refs: #1988 Refs: #1989
websockets · Dec 19, 2021 · 7e87037 · 7e87037
1 parent c82b087
commit 7e87037
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 5 deletions.
diff --git a/doc/ws.md b/doc/ws.md
@@ -270,6 +270,9 @@ This class represents a WebSocket. It extends the `EventEmitter`.
 - `options` {Object}
   - `followRedirects` {Boolean} Whether or not to follow redirects. Defaults to
     `false`.
+  - `generateMask` {Function} The function used to generate the making key. It
+    must return a `Buffer` of length 4 synchronously. By default the masking key
+    is generated with cryptographically strong random bytes.
   - `handshakeTimeout` {Number} Timeout in milliseconds for the handshake
     request. This is reset after every redirection.
   - `maxPayload` {Number} The maximum allowed message size in bytes.

diff --git a/lib/sender.js b/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.
@@ -22,9 +22,12 @@ class Sender {
    *
    * @param {(net.Socket|tls.Socket)} socket The connection socket
    * @param {Object} [extensions] An object containing the negotiated extensions
+   * @param {Function} [generateMask] The function used to generate the masking
+   *     key
    */
-  constructor(socket, extensions) {
+  constructor(socket, extensions, generateMask) {
     this._extensions = extensions || {};
+    this._generateMask = generateMask;
     this._socket = socket;
 
     this._firstFragment = true;
@@ -42,6 +45,8 @@ class Sender {
    * @param {Object} options Options object
    * @param {Boolean} [options.fin=false] Specifies whether or not to set the
    *     FIN bit
+   * @param {Function} [options.generateMask] The function used to generate the
+   *     masking key
    * @param {Boolean} [options.mask=false] Specifies whether or not to mask
    *     `data`
    * @param {Number} options.opcode The opcode
@@ -81,7 +86,13 @@ class Sender {
 
     if (!options.mask) return [target, data];
 
-    randomFillSync(mask, 0, 4);
+    let mask = _mask;
+
+    if (options.generateMask) {
+      mask = options.generateMask();
+    } else {
+      randomFillSync(_mask, 0, 4);
+    }
 
     target[1] |= 0x80;
     target[offset - 4] = mask[0];
@@ -156,6 +167,7 @@ class Sender {
         rsv1: false,
         opcode: 0x08,
         mask,
+        generateMask: this._generateMask,
         readOnly: false
       }),
       cb
@@ -200,6 +212,7 @@ class Sender {
         rsv1: false,
         opcode: 0x09,
         mask,
+        generateMask: this.generateMask,
         readOnly
       }),
       cb
@@ -244,6 +257,7 @@ class Sender {
         rsv1: false,
         opcode: 0x0a,
         mask,
+        generateMask: this._generateMask,
         readOnly
       }),
       cb
@@ -299,6 +313,7 @@ class Sender {
         rsv1,
         opcode,
         mask: options.mask,
+        generateMask: this._generateMask,
         readOnly: toBuffer.readOnly
       };
 
@@ -314,6 +329,7 @@ class Sender {
           rsv1: false,
           opcode,
           mask: options.mask,
+          generateMask: this._generateMask,
           readOnly: toBuffer.readOnly
         }),
         cb
@@ -331,6 +347,8 @@ class Sender {
    * @param {Number} options.opcode The opcode
    * @param {Boolean} [options.fin=false] Specifies whether or not to set the
    *     FIN bit
+   * @param {Function} [options.generateMask] The function used to generate the
+   *     masking key
    * @param {Boolean} [options.mask=false] Specifies whether or not to mask
    *     `data`
    * @param {Boolean} [options.readOnly=false] Specifies whether `data` can be

diff --git a/lib/websocket.js b/lib/websocket.js
@@ -7,7 +7,7 @@ const https = require('https');
 const http = require('http');
 const net = require('net');
 const tls = require('tls');
-const { randomBytes, createHash } = require('crypto');
+const { createHash, randomBytes, randomFillSync } = require('crypto');
 const { Readable } = require('stream');
 const { URL } = require('url');
 
@@ -192,6 +192,8 @@ class WebSocket extends EventEmitter {
    *     server and client
    * @param {Buffer} head The first packet of the upgraded stream
    * @param {Object} options Options object
+   * @param {Function} [options.generateMask] The function used to generate the
+   *     masking key
    * @param {Number} [options.maxPayload=0] The maximum allowed message size
    * @param {Boolean} [options.skipUTF8Validation=false] Specifies whether or
    *     not to skip UTF-8 validation for text and close messages
@@ -206,7 +208,7 @@ class WebSocket extends EventEmitter {
       skipUTF8Validation: options.skipUTF8Validation
     });
 
-    this._sender = new Sender(socket, this._extensions);
+    this._sender = new Sender(socket, this._extensions, options.generateMask);
     this._receiver = receiver;
     this._socket = socket;
 
@@ -613,6 +615,8 @@ module.exports = WebSocket;
  * @param {Object} [options] Connection options
  * @param {Boolean} [options.followRedirects=false] Whether or not to follow
  *     redirects
+ * @param {Function} [options.generateMask] The function used to generate the
+ *     masking key
  * @param {Number} [options.handshakeTimeout] Timeout in milliseconds for the
  *     handshake request
  * @param {Number} [options.maxPayload=104857600] The maximum allowed message
@@ -899,6 +903,7 @@ function initAsClient(websocket, address, protocols, options) {
     }
 
     websocket.setSocket(socket, head, {
+      generateMask: opts.generateMask,
       maxPayload: opts.maxPayload,
       skipUTF8Validation: opts.skipUTF8Validation
     });

diff --git a/test/websocket.test.js b/test/websocket.test.js
@@ -126,6 +126,42 @@ describe('WebSocket', () => {
           /^RangeError: Unsupported protocol version: 1000 \(supported versions: 8, 13\)$/
         );
       });
+
+      it('honors the `generateMask` option', (done) => {
+        const mask = Buffer.alloc(4);
+        const wss = new WebSocket.Server({ port: 0 }, () => {
+          const ws = new WebSocket(`ws://localhost:${wss.address().port}`, {
+            generateMask() {
+              return mask;
+            }
+          });
+
+          ws.on('open', () => {
+            ws.send('foo');
+          });
+
+          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', () => {
+            assert.ok(Buffer.concat(chunks).slice(2, 6).equals(mask));
+
+            ws.close();
+          });
+        });
+      });
     });
   });