Skip to content

Commit

Permalink
Merge 18bef28 into c82b087
Browse files Browse the repository at this point in the history
  • Loading branch information
lpinca committed Dec 20, 2021
2 parents c82b087 + 18bef28 commit 0106cb8
Show file tree
Hide file tree
Showing 5 changed files with 90 additions and 5 deletions.
3 changes: 3 additions & 0 deletions doc/ws.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
takes a `Buffer` that must be filled synchronously. 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
8 changes: 7 additions & 1 deletion lib/receiver.js
Original file line number Diff line number Diff line change
Expand Up @@ -417,7 +417,13 @@ 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);
Expand Down
42 changes: 39 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,14 +93,24 @@ 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];
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];
Expand Down Expand Up @@ -156,6 +178,8 @@ class Sender {
rsv1: false,
opcode: 0x08,
mask,
maskBuffer: this._maskBuffer,
generateMask: this._generateMask,
readOnly: false
}),
cb
Expand Down Expand Up @@ -200,6 +224,8 @@ class Sender {
rsv1: false,
opcode: 0x09,
mask,
maskBuffer: this._maskBuffer,
generateMask: this.generateMask,
readOnly
}),
cb
Expand Down Expand Up @@ -244,6 +270,8 @@ class Sender {
rsv1: false,
opcode: 0x0a,
mask,
maskBuffer: this._maskBuffer,
generateMask: this._generateMask,
readOnly
}),
cb
Expand Down Expand Up @@ -299,6 +327,8 @@ class Sender {
rsv1,
opcode,
mask: options.mask,
maskBuffer: this._maskBuffer,
generateMask: this._generateMask,
readOnly: toBuffer.readOnly
};

Expand All @@ -314,6 +344,8 @@ class Sender {
rsv1: false,
opcode,
mask: options.mask,
maskBuffer: this._maskBuffer,
generateMask: this._generateMask,
readOnly: toBuffer.readOnly
}),
cb
Expand All @@ -331,8 +363,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 0106cb8

Please sign in to comment.