fix(discord): bound gateway websocket payloads (#99998)

* fix(discord): bound gateway websocket payloads

* fix(discord): raise gateway payload cap for large events

* refactor(discord): centralize gateway websocket policy

---------

Co-authored-by: Peter Steinberger <steipete@gmail.com>

extensions/discord/src/internal/gateway-websocket-options.test.ts

@@ -0,0 +1,44 @@
+// Discord tests cover gateway websocket transport options.
+import { EventEmitter } from "node:events";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+const { webSocketCtorCalls } = vi.hoisted(() => ({
+  webSocketCtorCalls: [] as Array<{ url: string; options: unknown }>,
+}));
+
+vi.mock("ws", () => ({
+  WebSocket: class MockWebSocket extends EventEmitter {
+    readyState = 1;
+    send = vi.fn();
+    close = vi.fn();
+
+    constructor(url: string, options?: unknown) {
+      super();
+      webSocketCtorCalls.push({ url, options });
+    }
+  },
+}));
+
+describe("GatewayPlugin websocket options", () => {
+  let GatewayPlugin: typeof import("./gateway.js").GatewayPlugin;
+
+  beforeEach(async () => {
+    webSocketCtorCalls.length = 0;
+    ({ GatewayPlugin } = await import("./gateway.js"));
+  });
+
+  it("bounds inbound gateway websocket payloads", () => {
+    const gateway = new GatewayPlugin({
+      autoInteractions: false,
+      url: "wss://gateway.example.test",
+    });
+
+    gateway.connect(false);
+
+    expect(webSocketCtorCalls).toHaveLength(1);
+    expect(webSocketCtorCalls[0]).toEqual({
+      url: "wss://gateway.example.test/?v=10&encoding=json",
+      options: { maxPayload: 16 * 1024 * 1024 },
+    });
+  });
+});

extensions/discord/src/internal/gateway.ts

@@ -47,6 +47,11 @@ type GatewayPluginOptions = {
 const READY_STATE_OPEN = 1;
 const DEFAULT_GATEWAY_URL = "wss://gateway.discord.gg/";
 const DISCORD_GATEWAY_PAYLOAD_LIMIT_BYTES = 4096;
+// Discord can send multi-megabyte member chunks. Keep generous headroom while
+// bounding ws's 100 MiB default before an inbound payload reaches JSON parsing.
+export const DISCORD_GATEWAY_WS_CLIENT_OPTIONS = Object.freeze({
+  maxPayload: 16 * 1024 * 1024,
+}) satisfies ws.ClientOptions;
 const INVALID_SESSION_MIN_DELAY_MS = 1_000;
 const INVALID_SESSION_JITTER_MS = 4_000;
 
@@ -171,7 +176,7 @@ export class GatewayPlugin extends Plugin {
   }
 
   protected createWebSocket(url: string): ws.WebSocket {
-    return new ws.WebSocket(url);
+    return new ws.WebSocket(url, DISCORD_GATEWAY_WS_CLIENT_OPTIONS);
   }
 
   private setupWebSocket(resume: boolean): void {

extensions/discord/src/monitor/gateway-plugin.test.ts

@@ -54,6 +54,7 @@ const { GatewayIntents, GatewayPlugin } = vi.hoisted(() => {
 });
 
 vi.mock("../internal/gateway.js", () => ({
+  DISCORD_GATEWAY_WS_CLIENT_OPTIONS: { maxPayload: 16 * 1024 * 1024 },
   GatewayIntents,
   GatewayPlugin,
 }));

extensions/discord/src/monitor/gateway-plugin.ts

@@ -37,10 +37,7 @@ const DISCORD_GATEWAY_WS_RECEIVER_LIMIT_CODE = "WS_ERR_TOO_MANY_BUFFERED_PARTS";
 const DISCORD_GATEWAY_CLOSE_REASON_LOG_MAX_CHARS = 240;
 const discordDnsLookup = createDiscordDnsLookup();
 
-type DiscordGatewayWebSocketCtor = new (
-  url: string,
-  options?: { agent?: unknown; handshakeTimeout?: number },
-) => ws.WebSocket;
+type DiscordGatewayWebSocketCtor = typeof ws.WebSocket;
 type DiscordGatewayWebSocketAgent = InstanceType<typeof HttpsAgent> | HttpAgent;
 const registrationPromises = new WeakMap<discordGateway.GatewayPlugin, Promise<void>>();
 type DiscordGatewayClient = Parameters<discordGateway.GatewayPlugin["registerClient"]>[0];
@@ -268,6 +265,7 @@ function createGatewayPlugin(params: {
       // already our proxy path and behaves predictably for lifecycle cleanup.
       const WebSocketCtor = params.testing?.webSocketCtor ?? ws.default;
       const socket = new WebSocketCtor(url, {
+        ...discordGateway.DISCORD_GATEWAY_WS_CLIENT_OPTIONS,
         handshakeTimeout: DISCORD_GATEWAY_HANDSHAKE_TIMEOUT_MS,
         ...(params.wsAgent ? { agent: params.wsAgent } : {}),
       });

extensions/discord/src/monitor/provider.proxy.test.ts

@@ -146,11 +146,7 @@ const {
 
 // Unit test: don't import the real gateway just to check the prototype chain.
 vi.mock("../internal/gateway.js", () => ({
-  GatewayIntents,
-  GatewayPlugin,
-}));
-
-vi.mock("../internal/gateway.js", () => ({
+  DISCORD_GATEWAY_WS_CLIENT_OPTIONS: { maxPayload: 16 * 1024 * 1024 },
   GatewayIntents,
   GatewayPlugin,
 }));
@@ -162,7 +158,7 @@ vi.mock("node:https", () => ({
 vi.mock("ws", () => ({
   default: function MockWebSocket(
     url: string,
-    options?: { agent?: unknown; handshakeTimeout?: number },
+    options?: { agent?: unknown; handshakeTimeout?: number; maxPayload?: number },
   ) {
     webSocketSpy(url, options);
   },
@@ -398,6 +394,7 @@ describe("createDiscordGatewayPlugin", () => {
     expect(webSocketSpy).toHaveBeenCalledWith("wss://gateway.discord.gg", {
       agent: getLastAgent(),
       handshakeTimeout: 30_000,
+      maxPayload: 16 * 1024 * 1024,
     });
     expect(wsProxyAgentSpy).not.toHaveBeenCalled();
   });
@@ -515,6 +512,7 @@ describe("createDiscordGatewayPlugin", () => {
     expect(webSocketSpy).toHaveBeenCalledWith("wss://gateway.discord.gg", {
       agent: getLastProxyAgent(),
       handshakeTimeout: 30_000,
+      maxPayload: 16 * 1024 * 1024,
     });
     expect(runtime.log).toHaveBeenCalledWith("discord: gateway proxy enabled");
     expect(runtime.error).not.toHaveBeenCalled();
@@ -537,6 +535,7 @@ describe("createDiscordGatewayPlugin", () => {
     expect(webSocketSpy).toHaveBeenCalledWith("wss://gateway.discord.gg", {
       agent: getLastProxyAgent(),
       handshakeTimeout: 30_000,
+      maxPayload: 16 * 1024 * 1024,
     });
     expect(runtime.log).toHaveBeenCalledWith("discord: gateway proxy enabled");
     expect(runtime.error).not.toHaveBeenCalled();
@@ -559,6 +558,7 @@ describe("createDiscordGatewayPlugin", () => {
     expect(webSocketSpy).toHaveBeenCalledWith("wss://gateway.discord.gg", {
       agent: getLastProxyAgent(),
       handshakeTimeout: 30_000,
+      maxPayload: 16 * 1024 * 1024,
     });
     expect(runtime.error).not.toHaveBeenCalled();
     expect(runtime.log).toHaveBeenCalledWith("discord: gateway proxy enabled");
@@ -580,6 +580,7 @@ describe("createDiscordGatewayPlugin", () => {
     expect(webSocketSpy).toHaveBeenCalledWith("wss://gateway.discord.gg", {
       agent: getLastAgent(),
       handshakeTimeout: 30_000,
+      maxPayload: 16 * 1024 * 1024,
     });
     expect(runtime.log).not.toHaveBeenCalled();
   });
@@ -600,6 +601,7 @@ describe("createDiscordGatewayPlugin", () => {
     expect(webSocketSpy).toHaveBeenCalledWith("wss://gateway.discord.gg", {
       agent: getLastProxyAgent(),
       handshakeTimeout: 30_000,
+      maxPayload: 16 * 1024 * 1024,
     });
     expect(runtime.log).toHaveBeenCalledTimes(1);
     expect(runtime.log).toHaveBeenCalledWith("discord: gateway proxy enabled");
@@ -691,6 +693,7 @@ describe("createDiscordGatewayPlugin", () => {
     expect(webSocketSpy).toHaveBeenCalledWith("wss://gateway.discord.gg", {
       agent: getLastProxyAgent(),
       handshakeTimeout: 30_000,
+      maxPayload: 16 * 1024 * 1024,
     });
     expect(runtime.error).not.toHaveBeenCalled();
     expect(runtime.log).toHaveBeenCalledWith("discord: gateway proxy enabled");