From ffede67aedcca9221a0079b034b294bcf8aa4907 Mon Sep 17 00:00:00 2001
From: DrSmoothl <1787882683@qq.com>
Date: Fri, 10 Jul 2026 13:41:36 +0800
Subject: [PATCH] fix(websocket): cap decompressed message payloads

Apply the configured payload limit while per-message deflate is producing output, including fragmented frames. Abort oversized expansion before it can allocate unbounded memory.
---
 packages/websocket/src/extensions.ts        | 32 ++++++++++++++++++---
 packages/websocket/tests/extensions.test.ts | 31 ++++++++++++++++++++
 2 files changed, 59 insertions(+), 4 deletions(-)
 create mode 100644 packages/websocket/tests/extensions.test.ts

diff --git a/packages/websocket/src/extensions.ts b/packages/websocket/src/extensions.ts
index 7da9467d..f349a359 100644
--- a/packages/websocket/src/extensions.ts
+++ b/packages/websocket/src/extensions.ts
@@ -135,11 +135,35 @@ export function compressRaw(data: Buffer): Buffer {
 }
 
 export function decompressRaw(data: Buffer, maxPayload?: number): Buffer {
-  const inflated = zlib.inflateRawSync(Buffer.concat([data, TRAILER]), {
-    finishFlush: zlib.constants.Z_SYNC_FLUSH,
-  });
+  let inflated: Buffer;
+  try {
+    inflated = zlib.inflateRawSync(Buffer.concat([data, TRAILER]), {
+      finishFlush: zlib.constants.Z_SYNC_FLUSH,
+      // A post-inflate length check is too late for a compression bomb: zlib
+      // has already allocated the expanded buffer by then. Node enforces this
+      // limit while producing output. maxOutputLength cannot be zero, so keep
+      // the allocation ceiling at one byte for the maxPayload=0 edge case and
+      // let the exact check below reject any non-empty message.
+      ...(maxPayload === undefined
+        ? {}
+        : { maxOutputLength: Math.max(1, maxPayload) }),
+    });
+  } catch (cause) {
+    const zlibError = cause as Error & { code?: string };
+    if (maxPayload !== undefined && zlibError.code === 'ERR_BUFFER_TOO_LARGE') {
+      const err = new Error(
+        `Message too large after inflate (maxPayload=${maxPayload})`,
+        { cause },
+      ) as Error & { code?: number };
+      err.code = 1009;
+      throw err;
+    }
+    throw cause;
+  }
   if (maxPayload !== undefined && inflated.length > maxPayload) {
-    const err = new Error('Message too large after inflate') as Error & { code?: number };
+    const err = new Error(
+      `Message too large after inflate (maxPayload=${maxPayload})`,
+    ) as Error & { code?: number };
     err.code = 1009;
     throw err;
   }
diff --git a/packages/websocket/tests/extensions.test.ts b/packages/websocket/tests/extensions.test.ts
new file mode 100644
index 00000000..f294e236
--- /dev/null
+++ b/packages/websocket/tests/extensions.test.ts
@@ -0,0 +1,31 @@
+import { describe, expect, it } from 'vitest';
+
+import { compressRaw, decompressRaw } from '../src/extensions';
+
+describe('permessage-deflate decompression limits', () => {
+  it('accepts output exactly at maxPayload', () => {
+    const source = Buffer.alloc(4096, 0x61);
+    const compressed = compressRaw(source);
+
+    expect(decompressRaw(compressed, source.byteLength)).toEqual(source);
+  });
+
+  it('stops a compression bomb during inflate and maps it to WebSocket 1009', () => {
+    const maxPayload = 4096;
+    const compressed = compressRaw(Buffer.alloc(1024 * 1024, 0x61));
+    expect(compressed.byteLength).toBeLessThan(maxPayload);
+
+    let thrown: unknown;
+    try {
+      decompressRaw(compressed, maxPayload);
+    } catch (error) {
+      thrown = error;
+    }
+
+    expect(thrown).toMatchObject({
+      code: 1009,
+      message: `Message too large after inflate (maxPayload=${maxPayload})`,
+      cause: { code: 'ERR_BUFFER_TOO_LARGE' },
+    });
+  });
+});