Skip to content

Commit de1bac6

Browse files
sunlit-dengsteipete
sunlit-deng
and
steipete
authored
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 <[email protected]>
1 parent 1021784 commit de1bac6

5 files changed

Lines changed: 62 additions & 11 deletions

File tree

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

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
// Discord tests cover gateway websocket transport options.
2+
import { EventEmitter } from "node:events";
3+
import { beforeEach, describe, expect, it, vi } from "vitest";
4+
5+
const { webSocketCtorCalls } = vi.hoisted(() => ({
6+
webSocketCtorCalls: [] as Array<{ url: string; options: unknown }>,
7+
}));
8+
9+
vi.mock("ws", () => ({
10+
WebSocket: class MockWebSocket extends EventEmitter {
11+
readyState = 1;
12+
send = vi.fn();
13+
close = vi.fn();
14+
15+
constructor(url: string, options?: unknown) {
16+
super();
17+
webSocketCtorCalls.push({ url, options });
18+
}
19+
},
20+
}));
21+
22+
describe("GatewayPlugin websocket options", () => {
23+
let GatewayPlugin: typeof import("./gateway.js").GatewayPlugin;
24+
25+
beforeEach(async () => {
26+
webSocketCtorCalls.length = 0;
27+
({ GatewayPlugin } = await import("./gateway.js"));
28+
});
29+
30+
it("bounds inbound gateway websocket payloads", () => {
31+
const gateway = new GatewayPlugin({
32+
autoInteractions: false,
33+
url: "wss://gateway.example.test",
34+
});
35+
36+
gateway.connect(false);
37+
38+
expect(webSocketCtorCalls).toHaveLength(1);
39+
expect(webSocketCtorCalls[0]).toEqual({
40+
url: "wss://gateway.example.test/?v=10&encoding=json",
41+
options: { maxPayload: 16 * 1024 * 1024 },
42+
});
43+
});
44+
});

‎extensions/discord/src/internal/gateway.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,11 @@ type GatewayPluginOptions = {
4747
const READY_STATE_OPEN = 1;
4848
const DEFAULT_GATEWAY_URL = "wss://gateway.discord.gg/";
4949
const DISCORD_GATEWAY_PAYLOAD_LIMIT_BYTES = 4096;
50+
// Discord can send multi-megabyte member chunks. Keep generous headroom while
51+
// bounding ws's 100 MiB default before an inbound payload reaches JSON parsing.
52+
export const DISCORD_GATEWAY_WS_CLIENT_OPTIONS = Object.freeze({
53+
maxPayload: 16 * 1024 * 1024,
54+
}) satisfies ws.ClientOptions;
5055
const INVALID_SESSION_MIN_DELAY_MS = 1_000;
5156
const INVALID_SESSION_JITTER_MS = 4_000;
5257

@@ -171,7 +176,7 @@ export class GatewayPlugin extends Plugin {
171176
}
172177

173178
protected createWebSocket(url: string): ws.WebSocket {
174-
return new ws.WebSocket(url);
179+
return new ws.WebSocket(url, DISCORD_GATEWAY_WS_CLIENT_OPTIONS);
175180
}
176181

177182
private setupWebSocket(resume: boolean): void {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ const { GatewayIntents, GatewayPlugin } = vi.hoisted(() => {
5454
});
5555

5656
vi.mock("../internal/gateway.js", () => ({
57+
DISCORD_GATEWAY_WS_CLIENT_OPTIONS: { maxPayload: 16 * 1024 * 1024 },
5758
GatewayIntents,
5859
GatewayPlugin,
5960
}));

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

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,10 +37,7 @@ const DISCORD_GATEWAY_WS_RECEIVER_LIMIT_CODE = "WS_ERR_TOO_MANY_BUFFERED_PARTS";
3737
const DISCORD_GATEWAY_CLOSE_REASON_LOG_MAX_CHARS = 240;
3838
const discordDnsLookup = createDiscordDnsLookup();
3939

40-
type DiscordGatewayWebSocketCtor = new (
41-
url: string,
42-
options?: { agent?: unknown; handshakeTimeout?: number },
43-
) => ws.WebSocket;
40+
type DiscordGatewayWebSocketCtor = typeof ws.WebSocket;
4441
type DiscordGatewayWebSocketAgent = InstanceType<typeof HttpsAgent> | HttpAgent;
4542
const registrationPromises = new WeakMap<discordGateway.GatewayPlugin, Promise<void>>();
4643
type DiscordGatewayClient = Parameters<discordGateway.GatewayPlugin["registerClient"]>[0];
@@ -268,6 +265,7 @@ function createGatewayPlugin(params: {
268265
// already our proxy path and behaves predictably for lifecycle cleanup.
269266
const WebSocketCtor = params.testing?.webSocketCtor ?? ws.default;
270267
const socket = new WebSocketCtor(url, {
268+
...discordGateway.DISCORD_GATEWAY_WS_CLIENT_OPTIONS,
271269
handshakeTimeout: DISCORD_GATEWAY_HANDSHAKE_TIMEOUT_MS,
272270
...(params.wsAgent ? { agent: params.wsAgent } : {}),
273271
});

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

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -146,11 +146,7 @@ const {
146146

147147
// Unit test: don't import the real gateway just to check the prototype chain.
148148
vi.mock("../internal/gateway.js", () => ({
149-
GatewayIntents,
150-
GatewayPlugin,
151-
}));
152-
153-
vi.mock("../internal/gateway.js", () => ({
149+
DISCORD_GATEWAY_WS_CLIENT_OPTIONS: { maxPayload: 16 * 1024 * 1024 },
154150
GatewayIntents,
155151
GatewayPlugin,
156152
}));
@@ -162,7 +158,7 @@ vi.mock("node:https", () => ({
162158
vi.mock("ws", () => ({
163159
default: function MockWebSocket(
164160
url: string,
165-
options?: { agent?: unknown; handshakeTimeout?: number },
161+
options?: { agent?: unknown; handshakeTimeout?: number; maxPayload?: number },
166162
) {
167163
webSocketSpy(url, options);
168164
},
@@ -398,6 +394,7 @@ describe("createDiscordGatewayPlugin", () => {
398394
expect(webSocketSpy).toHaveBeenCalledWith("wss://gateway.discord.gg", {
399395
agent: getLastAgent(),
400396
handshakeTimeout: 30_000,
397+
maxPayload: 16 * 1024 * 1024,
401398
});
402399
expect(wsProxyAgentSpy).not.toHaveBeenCalled();
403400
});
@@ -515,6 +512,7 @@ describe("createDiscordGatewayPlugin", () => {
515512
expect(webSocketSpy).toHaveBeenCalledWith("wss://gateway.discord.gg", {
516513
agent: getLastProxyAgent(),
517514
handshakeTimeout: 30_000,
515+
maxPayload: 16 * 1024 * 1024,
518516
});
519517
expect(runtime.log).toHaveBeenCalledWith("discord: gateway proxy enabled");
520518
expect(runtime.error).not.toHaveBeenCalled();
@@ -537,6 +535,7 @@ describe("createDiscordGatewayPlugin", () => {
537535
expect(webSocketSpy).toHaveBeenCalledWith("wss://gateway.discord.gg", {
538536
agent: getLastProxyAgent(),
539537
handshakeTimeout: 30_000,
538+
maxPayload: 16 * 1024 * 1024,
540539
});
541540
expect(runtime.log).toHaveBeenCalledWith("discord: gateway proxy enabled");
542541
expect(runtime.error).not.toHaveBeenCalled();
@@ -559,6 +558,7 @@ describe("createDiscordGatewayPlugin", () => {
559558
expect(webSocketSpy).toHaveBeenCalledWith("wss://gateway.discord.gg", {
560559
agent: getLastProxyAgent(),
561560
handshakeTimeout: 30_000,
561+
maxPayload: 16 * 1024 * 1024,
562562
});
563563
expect(runtime.error).not.toHaveBeenCalled();
564564
expect(runtime.log).toHaveBeenCalledWith("discord: gateway proxy enabled");
@@ -580,6 +580,7 @@ describe("createDiscordGatewayPlugin", () => {
580580
expect(webSocketSpy).toHaveBeenCalledWith("wss://gateway.discord.gg", {
581581
agent: getLastAgent(),
582582
handshakeTimeout: 30_000,
583+
maxPayload: 16 * 1024 * 1024,
583584
});
584585
expect(runtime.log).not.toHaveBeenCalled();
585586
});
@@ -600,6 +601,7 @@ describe("createDiscordGatewayPlugin", () => {
600601
expect(webSocketSpy).toHaveBeenCalledWith("wss://gateway.discord.gg", {
601602
agent: getLastProxyAgent(),
602603
handshakeTimeout: 30_000,
604+
maxPayload: 16 * 1024 * 1024,
603605
});
604606
expect(runtime.log).toHaveBeenCalledTimes(1);
605607
expect(runtime.log).toHaveBeenCalledWith("discord: gateway proxy enabled");
@@ -691,6 +693,7 @@ describe("createDiscordGatewayPlugin", () => {
691693
expect(webSocketSpy).toHaveBeenCalledWith("wss://gateway.discord.gg", {
692694
agent: getLastProxyAgent(),
693695
handshakeTimeout: 30_000,
696+
maxPayload: 16 * 1024 * 1024,
694697
});
695698
expect(runtime.error).not.toHaveBeenCalled();
696699
expect(runtime.log).toHaveBeenCalledWith("discord: gateway proxy enabled");

0 commit comments

Comments
 (0)