Skip to content

Commit

Permalink
[feature] Introduce the generateMask option
Browse files Browse the repository at this point in the history
The `generateMask` option specifies a function that can be used to
generate custom masking keys.

Refs: #1986
Refs: #1988
Refs: #1989
  • Loading branch information
lpinca committed Dec 20, 2021
1 parent c82b087 commit eb2e3a8
Show file tree
Hide file tree
Showing 4 changed files with 80 additions and 4 deletions.
4 changes: 4 additions & 0 deletions doc/ws.md
Original file line number Diff line number Diff line change
Expand Up @@ -270,6 +270,10 @@ 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 masking key. It
takes a `Buffer` that must be filled synchronously and is called before a
message is sent, for each message. By default the buffer is filled 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.
Expand Down
38 changes: 35 additions & 3 deletions lib/sender.js
Original file line number Diff line number Diff line change
Expand Up @@ -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 maskBuffer = Buffer.alloc(4);

/**
* HyBi Sender implementation.
Expand All @@ -22,9 +22,17 @@ 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 || {};

if (generateMask) {
this._generateMask = generateMask;
this._maskBuffer = Buffer.alloc(4);
}

this._socket = socket;

this._firstFragment = true;
Expand All @@ -42,8 +50,12 @@ 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 {Buffer} [options.maskBuffer] The buffer used to store the masking
* key
* @param {Number} options.opcode The opcode
* @param {Boolean} [options.readOnly=false] Specifies whether `data` can be
* modified
Expand Down Expand Up @@ -81,7 +93,13 @@ class Sender {

if (!options.mask) return [target, data];

randomFillSync(mask, 0, 4);
const mask = options.maskBuffer ? options.maskBuffer : maskBuffer;

if (options.generateMask) {
options.generateMask(mask);
} else {
randomFillSync(mask, 0, 4);
}

target[1] |= 0x80;
target[offset - 4] = mask[0];
Expand Down Expand Up @@ -156,6 +174,8 @@ class Sender {
rsv1: false,
opcode: 0x08,
mask,
maskBuffer: this._maskBuffer,
generateMask: this._generateMask,
readOnly: false
}),
cb
Expand Down Expand Up @@ -200,6 +220,8 @@ class Sender {
rsv1: false,
opcode: 0x09,
mask,
maskBuffer: this._maskBuffer,
generateMask: this._generateMask,
readOnly
}),
cb
Expand Down Expand Up @@ -244,6 +266,8 @@ class Sender {
rsv1: false,
opcode: 0x0a,
mask,
maskBuffer: this._maskBuffer,
generateMask: this._generateMask,
readOnly
}),
cb
Expand Down Expand Up @@ -299,6 +323,8 @@ class Sender {
rsv1,
opcode,
mask: options.mask,
maskBuffer: this._maskBuffer,
generateMask: this._generateMask,
readOnly: toBuffer.readOnly
};

Expand All @@ -314,6 +340,8 @@ class Sender {
rsv1: false,
opcode,
mask: options.mask,
maskBuffer: this._maskBuffer,
generateMask: this._generateMask,
readOnly: toBuffer.readOnly
}),
cb
Expand All @@ -331,8 +359,12 @@ 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 {Buffer} [options.maskBuffer] The buffer used to store 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
Expand Down
7 changes: 6 additions & 1 deletion lib/websocket.js
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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;

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -899,6 +903,7 @@ function initAsClient(websocket, address, protocols, options) {
}

websocket.setSocket(socket, head, {
generateMask: opts.generateMask,
maxPayload: opts.maxPayload,
skipUTF8Validation: opts.skipUTF8Validation
});
Expand Down
35 changes: 35 additions & 0 deletions test/websocket.test.js
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,41 @@ describe('WebSocket', () => {
/^RangeError: Unsupported protocol version: 1000 \(supported versions: 8, 13\)$/
);
});

it('honors the `generateMask` option', (done) => {
const wss = new WebSocket.Server({ port: 0 }, () => {
const ws = new WebSocket(`ws://localhost:${wss.address().port}`, {
generateMask() {}
});

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(Buffer.alloc(4))
);

ws.close();
});
});
});
});
});

Expand Down

0 comments on commit eb2e3a8

Please sign in to comment.