websocket: limit the number of fragments in a message · nodejs/undici@32dbf0b · GitHub
Skip to content

Commit 32dbf0b

Browse files
lpincamcollina
lpinca
authored and
mcollina
committed
websocket: limit the number of fragments in a message
Introduce the `maxFragments` option to limit the number of fragments in a message. Without this limit, a remote peer could exploit structural wrapper overhead by sending many tiny fragments, consuming far more memory than the payload itself and crashing the process due to OOM. When the limit is hit, the connection is closed with close code 1008. Signed-off-by: Luigi Pinca <luigipinca@gmail.com>
1 parent 0d6ecc5 commit 32dbf0b

6 files changed

Lines changed: 118 additions & 4 deletions

File tree

‎docs/docs/api/Client.md‎

Lines changed: 1 addition & 0 deletions

‎lib/dispatcher/dispatcher-base.js‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ class DispatcherBase extends Dispatcher {
3838
*/
3939
get webSocketOptions () {
4040
return {
41+
maxFragments: this[kWebSocketOptions].maxFragments ?? 131072,
4142
maxPayloadSize: this[kWebSocketOptions].maxPayloadSize ?? 128 * 1024 * 1024 // 128 MB default
4243
}
4344
}

‎lib/web/websocket/receiver.js‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@ class ByteParser extends Writable {
3939
/** @type {import('./websocket').Handler} */
4040
#handler
4141

42+
43+
/** @type {number} */
44+
#maxFragments
45+
4246
/** @type {number} */
4347
#maxPayloadSize
4448

@@ -52,6 +56,7 @@ class ByteParser extends Writable {
5256

5357
this.#handler = handler
5458
this.#extensions = extensions == null ? new Map() : extensions
59+
this.#maxFragments = options.maxFragments ?? 0
5560
this.#maxPayloadSize = options.maxPayloadSize ?? 0
5661

5762
if (this.#extensions.has('permessage-deflate')) {
@@ -242,7 +247,9 @@ class ByteParser extends Writable {
242247
this.#state = parserStates.INFO
243248
} else {
244249
if (!this.#info.compressed) {
245-
this.writeFragments(body)
250+
if (body.length && !this.writeFragments(body)) {
251+
return
252+
}
246253

247254
// If the frame is not fragmented, a message has been received.
248255
// If the frame is fragmented, it will terminate with a fin bit set
@@ -264,7 +271,9 @@ class ByteParser extends Writable {
264271
return
265272
}
266273

267-
this.writeFragments(data)
274+
if (data.length && !this.writeFragments(data)) {
275+
return
276+
}
268277

269278
// Check cumulative fragment size
270279
if (this.#maxPayloadSize > 0 && this.#fragmentsBytes > this.#maxPayloadSize) {
@@ -345,8 +354,17 @@ class ByteParser extends Writable {
345354
}
346355

347356
writeFragments (fragment) {
357+
if (
358+
this.#maxFragments > 0 &&
359+
this.#fragments.length === this.#maxFragments
360+
) {
361+
failWebsocketConnection(this.#handler, 1008, 'Too many message fragments')
362+
return false
363+
}
364+
348365
this.#fragmentsBytes += fragment.length
349366
this.#fragments.push(fragment)
367+
return true
350368
}
351369

352370
consumeFragments () {

‎lib/web/websocket/websocket.js‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -468,10 +468,12 @@ class WebSocket extends EventTarget {
468468
// once this happens, the connection is open
469469
this.#handler.socket = response.socket
470470

471-
// Get maxPayloadSize from dispatcher options
471+
// Get options from dispatcher options
472+
const maxFragments = this.#handler.controller.dispatcher?.webSocketOptions?.maxFragments
472473
const maxPayloadSize = this.#handler.controller.dispatcher?.webSocketOptions?.maxPayloadSize
473474

474475
const parser = new ByteParser(this.#handler, parsedExtensions, {
476+
maxFragments,
475477
maxPayloadSize
476478
})
477479
parser.on('drain', () => this.#handler.onParserDrain())

‎test/websocket/fragments.js‎

Lines changed: 88 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
const { test, after } = require('node:test')
44
const { WebSocketServer } = require('ws')
5-
const { WebSocket } = require('../..')
5+
const { Agent, WebSocket } = require('../..')
66
const diagnosticsChannel = require('node:diagnostics_channel')
77

88
test('Fragmented frame with a ping frame in the middle of it', (t) => {
@@ -39,3 +39,90 @@ test('Fragmented frame with a ping frame in the middle of it', (t) => {
3939
})
4040
})
4141
})
42+
43+
test('Too many fragments (uncompressed)', (t, done) => {
44+
t.plan(4)
45+
46+
const agent = new Agent({
47+
webSocket: {
48+
maxFragments: 3
49+
}
50+
})
51+
52+
const server = new WebSocketServer({ port: 0 }, () => {
53+
const { port } = server.address()
54+
const client = new WebSocket(`ws://127.0.0.1:${port}`, {
55+
dispatcher: agent
56+
})
57+
58+
client.addEventListener('error', (event) => {
59+
t.assert.ok(true)
60+
})
61+
62+
client.addEventListener('close', (event) => {
63+
t.assert.deepStrictEqual(event.code, 1006)
64+
})
65+
})
66+
67+
server.on('connection', (ws) => {
68+
ws.on('close', (code, reason) => {
69+
t.assert.deepStrictEqual(code, 1008)
70+
t.assert.deepStrictEqual(reason.toString(), 'Too many message fragments')
71+
agent.close()
72+
server.close(done)
73+
})
74+
75+
const fragment = Buffer.from('a')
76+
const options = { fin: false }
77+
78+
ws.send(fragment, options)
79+
ws.send(fragment, options)
80+
ws.send(fragment, options)
81+
ws.send(fragment, options)
82+
})
83+
})
84+
85+
test('Too many fragments (compressed)', (t, done) => {
86+
t.plan(4)
87+
88+
const agent = new Agent({
89+
webSocket: {
90+
maxFragments: 3
91+
}
92+
})
93+
94+
const server = new WebSocketServer({
95+
perMessageDeflate: { threshold: 0 },
96+
port: 0
97+
}, () => {
98+
const { port } = server.address()
99+
const client = new WebSocket(`ws://127.0.0.1:${port}`, {
100+
dispatcher: agent
101+
})
102+
103+
client.addEventListener('error', (event) => {
104+
t.assert.ok(true)
105+
})
106+
107+
client.addEventListener('close', (event) => {
108+
t.assert.deepStrictEqual(event.code, 1006)
109+
})
110+
})
111+
112+
server.on('connection', (ws) => {
113+
ws.on('close', (code, reason) => {
114+
t.assert.deepStrictEqual(code, 1008)
115+
t.assert.deepStrictEqual(reason.toString(), 'Too many message fragments')
116+
agent.close()
117+
server.close(done)
118+
})
119+
120+
const fragment = Buffer.from('a')
121+
const options = { fin: false }
122+
123+
ws.send(fragment, options)
124+
ws.send(fragment, options)
125+
ws.send(fragment, options)
126+
ws.send(fragment, options)
127+
})
128+
})

‎types/client.d.ts‎

Lines changed: 5 additions & 0 deletions

0 commit comments

Comments
 (0)