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>
(cherry picked from commit 32dbf0b)

Author: lpinca

docs/docs/api/Client.md

@@ -24,6 +24,8 @@ Returns: `Client`
 * **keepAliveTimeoutThreshold** `number | null` (optional) - Default: `2e3` - A number of milliseconds subtracted from server *keep-alive* hints when overriding `keepAliveTimeout` to account for timing inaccuracies caused by e.g. transport latency. Defaults to 2 seconds.
 * **maxHeaderSize** `number | null` (optional) - Default: `--max-http-header-size` or `16384` - The maximum length of request headers in bytes. Defaults to Node.js' --max-http-header-size or 16KiB.
 * **maxResponseSize** `number | null` (optional) - Default: `-1` - The maximum length of response body in bytes. Set to `-1` to disable.
+* **webSocket** `WebSocketOptions` (optional) - WebSocket-specific configuration options.
+  * **maxFragments** `number` (optional) - Defailt: `131072` - Maximum number of fragments in a message. Set to 0 to disable the limit.
 * **pipelining** `number | null` (optional) - Default: `1` - The amount of concurrent requests to be sent over the single TCP/TLS connection according to [RFC7230](https://tools.ietf.org/html/rfc7230#section-6.3.2). Carefully consider your workload and environment before enabling concurrent requests as pipelining may reduce performance if used incorrectly. Pipelining is sensitive to network stack settings as well as head of line blocking caused by e.g. long running requests. Set to `0` to disable keep-alive connections.
 * **connect** `ConnectOptions | Function | null` (optional) - Default: `null`.
 * **strictContentLength** `Boolean` (optional) - Default: `true` - Whether to treat request content length mismatches as errors. If true, an error is thrown when the request content-length header doesn't match the length of the request body. **Security Warning:** Disabling this option can expose your application to HTTP Request Smuggling attacks, where mismatched content-length headers cause servers and proxies to interpret request boundaries differently. This can lead to cache poisoning, credential hijacking, and bypassing security controls. Only disable this in controlled environments where you fully trust the request source.

lib/dispatcher/agent.js

@@ -35,7 +35,7 @@ class Agent extends DispatcherBase {
       throw new InvalidArgumentError('maxOrigins must be a number greater than 0')
     }
 
-    super()
+    super(options)
 
     if (connect && typeof connect !== 'function') {
       connect = { ...connect }

lib/dispatcher/balanced-pool.js

@@ -54,7 +54,7 @@ class BalancedPool extends PoolBase {
       throw new InvalidArgumentError('factory must be a function.')
     }
 
-    super()
+    super(opts)
 
     this[kOptions] = { ...util.deepClone(opts) }
     this[kOptions].interceptors = opts.interceptors

lib/dispatcher/client.js

@@ -114,7 +114,8 @@ class Client extends DispatcherBase {
     useH2c,
     initialWindowSize,
     connectionWindowSize,
-    pingInterval
+    pingInterval,
+    webSocket
   } = {}) {
     if (keepAlive !== undefined) {
       throw new InvalidArgumentError('unsupported keepAlive, use pipelining=0 instead')
@@ -222,7 +223,7 @@ class Client extends DispatcherBase {
       throw new InvalidArgumentError('pingInterval must be a positive integer, greater or equal to 0')
     }
 
-    super()
+    super({ webSocket })
 
     if (typeof connect !== 'function') {
       connect = buildConnector({

lib/dispatcher/dispatcher-base.js

@@ -11,6 +11,7 @@ const { kDestroy, kClose, kClosed, kDestroyed, kDispatch } = require('../core/sy
 
 const kOnDestroyed = Symbol('onDestroyed')
 const kOnClosed = Symbol('onClosed')
+const kWebSocketOptions = Symbol('web socket options')
 
 class DispatcherBase extends Dispatcher {
   /** @type {boolean} */
@@ -25,6 +26,20 @@ class DispatcherBase extends Dispatcher {
   /** @type {Array<Function>|null} */
   [kOnClosed] = null
 
+  /**
+   * @param {{ webSocket?: { maxFragments?: number } }} [opts]
+   */
+  constructor (opts) {
+    super()
+    this[kWebSocketOptions] = opts?.webSocket ?? {}
+  }
+
+  get webSocketOptions () {
+    return {
+      maxFragments: this[kWebSocketOptions].maxFragments ?? 131072
+    }
+  }
+
   /** @returns {boolean} */
   get destroyed () {
     return this[kDestroyed]

lib/dispatcher/pool.js

@@ -63,7 +63,7 @@ class Pool extends PoolBase {
       })
     }
 
-    super()
+    super(options)
 
     this[kConnections] = connections || null
     this[kUrl] = util.parseOrigin(origin)

lib/web/websocket/receiver.js

@@ -39,15 +39,20 @@ class ByteParser extends Writable {
   /** @type {import('./websocket').Handler} */
   #handler
 
+  /** @type {number} */
+  #maxFragments
+
   /**
    * @param {import('./websocket').Handler} handler
    * @param {Map<string, string>|null} extensions
+   * @param {{ maxFragments?: number }} [options]
    */
-  constructor (handler, extensions) {
+  constructor (handler, extensions, options = {}) {
     super()
 
     this.#handler = handler
     this.#extensions = extensions == null ? new Map() : extensions
+    this.#maxFragments = options.maxFragments ?? 0
 
     if (this.#extensions.has('permessage-deflate')) {
       this.#extensions.set('permessage-deflate', new PerMessageDeflate(extensions))
@@ -212,7 +217,9 @@ class ByteParser extends Writable {
           this.#state = parserStates.INFO
         } else {
           if (!this.#info.compressed) {
-            this.writeFragments(body)
+            if (body.length && !this.writeFragments(body)) {
+              return
+            }
 
             // If the frame is not fragmented, a message has been received.
             // If the frame is fragmented, it will terminate with a fin bit set
@@ -232,7 +239,9 @@ class ByteParser extends Writable {
                 return
               }
 
-              this.writeFragments(data)
+              if (data.length && !this.writeFragments(data)) {
+                return
+              }
 
               if (!this.#info.fin) {
                 this.#state = parserStates.INFO
@@ -305,8 +314,17 @@ class ByteParser extends Writable {
   }
 
   writeFragments (fragment) {
+    if (
+      this.#maxFragments > 0 &&
+      this.#fragments.length === this.#maxFragments
+    ) {
+      failWebsocketConnection(this.#handler, 1008, 'Too many message fragments')
+      return false
+    }
+
     this.#fragmentsBytes += fragment.length
     this.#fragments.push(fragment)
+    return true
   }
 
   consumeFragments () {

lib/web/websocket/websocket.js

@@ -468,7 +468,11 @@ class WebSocket extends EventTarget {
     // once this happens, the connection is open
     this.#handler.socket = response.socket
 
-    const parser = new ByteParser(this.#handler, parsedExtensions)
+    const maxFragments = this.#handler.controller.dispatcher?.webSocketOptions?.maxFragments
+
+    const parser = new ByteParser(this.#handler, parsedExtensions, {
+      maxFragments
+    })
     parser.on('drain', () => this.#handler.onParserDrain())
     parser.on('error', (err) => this.#handler.onParserError(err))
 

test/websocket/fragments.js

@@ -2,7 +2,7 @@
 
 const { test, after } = require('node:test')
 const { WebSocketServer } = require('ws')
-const { WebSocket } = require('../..')
+const { Agent, WebSocket } = require('../..')
 const diagnosticsChannel = require('node:diagnostics_channel')
 
 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) => {
     })
   })
 })
+
+test('Too many fragments (uncompressed)', (t, done) => {
+  t.plan(4)
+
+  const agent = new Agent({
+    webSocket: {
+      maxFragments: 3
+    }
+  })
+
+  const server = new WebSocketServer({ port: 0 }, () => {
+    const { port } = server.address()
+    const client = new WebSocket(`ws://127.0.0.1:${port}`, {
+      dispatcher: agent
+    })
+
+    client.addEventListener('error', (event) => {
+      t.assert.ok(true)
+    })
+
+    client.addEventListener('close', (event) => {
+      t.assert.deepStrictEqual(event.code, 1006)
+    })
+  })
+
+  server.on('connection', (ws) => {
+    ws.on('close', (code, reason) => {
+      t.assert.deepStrictEqual(code, 1008)
+      t.assert.deepStrictEqual(reason.toString(), 'Too many message fragments')
+      agent.close()
+      server.close(done)
+    })
+
+    const fragment = Buffer.from('a')
+    const options = { fin: false }
+
+    ws.send(fragment, options)
+    ws.send(fragment, options)
+    ws.send(fragment, options)
+    ws.send(fragment, options)
+  })
+})
+
+test('Too many fragments (compressed)', (t, done) => {
+  t.plan(4)
+
+  const agent = new Agent({
+    webSocket: {
+      maxFragments: 3
+    }
+  })
+
+  const server = new WebSocketServer({
+    perMessageDeflate: { threshold: 0 },
+    port: 0
+  }, () => {
+    const { port } = server.address()
+    const client = new WebSocket(`ws://127.0.0.1:${port}`, {
+      dispatcher: agent
+    })
+
+    client.addEventListener('error', (event) => {
+      t.assert.ok(true)
+    })
+
+    client.addEventListener('close', (event) => {
+      t.assert.deepStrictEqual(event.code, 1006)
+    })
+  })
+
+  server.on('connection', (ws) => {
+    ws.on('close', (code, reason) => {
+      t.assert.deepStrictEqual(code, 1008)
+      t.assert.deepStrictEqual(reason.toString(), 'Too many message fragments')
+      agent.close()
+      server.close(done)
+    })
+
+    const fragment = Buffer.from('a')
+    const options = { fin: false }
+
+    ws.send(fragment, options)
+    ws.send(fragment, options)
+    ws.send(fragment, options)
+    ws.send(fragment, options)
+  })
+})

types/client.d.ts

@@ -107,6 +107,8 @@ export declare namespace Client {
      * @default 60000
      */
     pingInterval?: number;
+    /** WebSocket-specific configuration options. */
+    webSocket?: WebSocketOptions;
   }
   export interface SocketInfo {
     localAddress?: string
@@ -118,6 +120,13 @@ export declare namespace Client {
     bytesWritten?: number
     bytesRead?: number
   }
+  export interface WebSocketOptions {
+    /**
+     * Maximum number of fragments in a message. Set to 0 to disable the limit.
+     * @default 131072
+     */
+    maxFragments?: number;
+  }
 }
 
 export default Client