Configure gateway heartbeat timeout

.lintstagedrc.mjs

@@ -1,4 +1,9 @@
 export default {
     "*.{j,t}s": [() => "npm run build:src:tsgo", "eslint --concurrency 4" /* sweet spot it seems */, "prettier --write"],
-    "src/schemas/{*,**/*}.ts": [() => "npm run build:src:tsgo", () => "node scripts/schema.js", () => "node scripts/openapi.js", () => "git add assets/schemas.json assets/openapi.json"],
+    "src/schemas/{*,**/*}.ts": [
+        () => "npm run build:src:tsgo",
+        () => "node scripts/schema.js",
+        () => "node scripts/openapi.js",
+        () => "git add assets/schemas.json assets/openapi.json",
+    ],
 };

eslint.config.mjs

@@ -66,7 +66,7 @@ export default defineConfig([
             // "sort-imports": ["error", {}],
             "default-case": "error",
             "default-case-last": "error",
-            "yoda": "error",
+            yoda: "error",
             // unsure what the defaults are here, but we want them to error
             "for-direction": "error",
             "constructor-super": "error",

src/cdn/util/Storage.ts

@@ -86,7 +86,9 @@ if (process.env.STORAGE_PROVIDER === "file" || !process.env.STORAGE_PROVIDER) {
     const forcePathStyle = process.env.STORAGE_FORCE_PATH_STYLE === "true";
 
     if (process.env.STORAGE_FORCE_PATH_STYLE === undefined) {
-        console.warn(`[CDN] STORAGE_FORCE_PATH_STYLE is not set for S3 provider; defaulting to virtual-hosted style. Set STORAGE_FORCE_PATH_STYLE=true to enable path-style addressing.`);
+        console.warn(
+            `[CDN] STORAGE_FORCE_PATH_STYLE is not set for S3 provider; defaulting to virtual-hosted style. Set STORAGE_FORCE_PATH_STYLE=true to enable path-style addressing.`,
+        );
     }
 
     const { S3Storage } = require("./S3Storage");

src/gateway/events/Connection.test.ts

@@ -0,0 +1,74 @@
+import assert from "node:assert/strict";
+import { afterEach, describe, it, mock } from "node:test";
+import { GATEWAY_HEARTBEAT_INTERVAL } from "../../util/config/types/GatewayConfiguration";
+
+afterEach(() => {
+    mock.restoreAll();
+});
+
+function createSocket() {
+    const closes: { code: number; reason?: string }[] = [];
+    const socket = {
+        on() {
+            return socket;
+        },
+        close(code: number, reason?: string) {
+            closes.push({ code, reason });
+            return socket;
+        },
+    };
+
+    return { socket, closes };
+}
+
+describe("Connection", () => {
+    it("uses configured heartbeat timeout and advertises the shared heartbeat interval", async () => {
+        process.env.DATABASE ??= "postgres://spacebar:spacebar@localhost:5432/spacebar";
+        const { Config } = require("../../util/index.js") as typeof import("../../util/index.js");
+        const HeartbeatUtil = require("../util/Heartbeat.js") as typeof import("../util/Heartbeat.js");
+        const SendUtil = require("../util/Send.js") as typeof import("../util/Send.js");
+        const heartbeatTimeout = 65_000;
+        const { socket } = createSocket();
+        const sentPayloads: unknown[] = [];
+        const setHeartbeat = mock.method(HeartbeatUtil, "setHeartbeat", () => undefined);
+        mock.method(SendUtil, "Send", async (_socket: unknown, payload: unknown) => {
+            sentPayloads.push(payload);
+        });
+        mock.method(
+            Config,
+            "get",
+            () =>
+                ({
+                    gateway: { heartbeatTimeout },
+                    security: {
+                        cdnSignatureIncludeIp: false,
+                        cdnSignatureIncludeUserAgent: false,
+                        forwardedFor: null,
+                    },
+                }) as ReturnType<typeof Config.get>,
+        );
+        const { Connection } = require("./Connection.js") as typeof import("./Connection.js");
+
+        await Connection.call(
+            { clients: { size: 1 } } as never,
+            socket as never,
+            {
+                headers: { "user-agent": "spacebar-test" },
+                socket: { remoteAddress: "127.0.0.1" },
+                url: "/?encoding=json&version=8",
+            } as never,
+        );
+        clearTimeout((socket as never as { readyTimeout?: NodeJS.Timeout }).readyTimeout);
+
+        assert.equal(setHeartbeat.mock.callCount(), 1);
+        assert.deepEqual(setHeartbeat.mock.calls[0].arguments, [socket, heartbeatTimeout]);
+        assert.deepEqual(sentPayloads, [
+            {
+                op: 10,
+                d: {
+                    heartbeat_interval: GATEWAY_HEARTBEAT_INTERVAL,
+                },
+            },
+        ]);
+    });
+});

src/gateway/events/Connection.ts

@@ -27,7 +27,7 @@ import { Close } from "./Close";
 import { Message } from "./Message";
 import { Deflate, Inflate } from "fast-zlib";
 import { URL } from "node:url";
-import { Config, ErlpackType } from "@spacebar/util";
+import { Config, ErlpackType, GATEWAY_HEARTBEAT_INTERVAL } from "@spacebar/util";
 import { Decoder, Encoder } from "@toondepauw/node-zstd";
 
 let erlpack: ErlpackType | null = null;
@@ -137,12 +137,12 @@ export async function Connection(this: WS.Server, socket: WebSocket, request: In
         socket.permissions = {};
         socket.sequence = 0;
 
-        setHeartbeat(socket);
+        setHeartbeat(socket, Config.get().gateway.heartbeatTimeout);
 
         await Send(socket, {
             op: OPCODES.Hello,
             d: {
-                heartbeat_interval: 1000 * 30,
+                heartbeat_interval: GATEWAY_HEARTBEAT_INTERVAL,
             },
         });
 

src/gateway/opcodes/Heartbeat.test.ts

@@ -0,0 +1,28 @@
+import assert from "node:assert/strict";
+import { afterEach, describe, it, mock } from "node:test";
+import { OPCODES } from "../util/Constants";
+
+afterEach(() => {
+    mock.restoreAll();
+});
+
+describe("onHeartbeat", () => {
+    it("reschedules heartbeat timeouts with the configured gateway timeout", async () => {
+        process.env.DATABASE ??= "postgres://spacebar:spacebar@localhost:5432/spacebar";
+        const { Config, Session } = require("../../util/index.js") as typeof import("../../util/index.js");
+        const HeartbeatUtil = require("../util/Heartbeat.js") as typeof import("../util/Heartbeat.js");
+        const SendUtil = require("../util/Send.js") as typeof import("../util/Send.js");
+        const heartbeatTimeout = 65_000;
+        const socket = { session_id: "session", user_id: "user" };
+        const setHeartbeat = mock.method(HeartbeatUtil, "setHeartbeat", () => undefined);
+        mock.method(SendUtil, "Send", async () => undefined);
+        mock.method(Session, "update", async () => undefined);
+        mock.method(Config, "get", () => ({ gateway: { heartbeatTimeout } }) as ReturnType<typeof Config.get>);
+        const { onHeartbeat } = require("./Heartbeat.js") as typeof import("./Heartbeat.js");
+
+        await onHeartbeat.call(socket as never, { op: OPCODES.Heartbeat, d: null });
+
+        assert.equal(setHeartbeat.mock.callCount(), 1);
+        assert.deepEqual(setHeartbeat.mock.calls[0].arguments, [socket, heartbeatTimeout]);
+    });
+});

src/gateway/opcodes/Heartbeat.ts

@@ -19,7 +19,7 @@
 import { OPCODES, Payload, WebSocket } from "@spacebar/gateway";
 import { setHeartbeat } from "../util/Heartbeat";
 import { Send } from "../util/Send";
-import { Session } from "@spacebar/util";
+import { Config, Session } from "@spacebar/util";
 import { FindOptionsWhere } from "typeorm";
 
 interface QoSData {
@@ -36,7 +36,7 @@ export interface QoSPayload {
 export async function onHeartbeat(this: WebSocket, data: Payload) {
     // TODO: validate payload
 
-    setHeartbeat(this);
+    setHeartbeat(this, Config.get().gateway.heartbeatTimeout);
 
     if (data.op === OPCODES.SetQoS) {
         this.qos = (data.d as QoSData).qos;

src/gateway/util/Heartbeat.test.ts

@@ -0,0 +1,59 @@
+import assert from "node:assert/strict";
+import { afterEach, describe, it, mock } from "node:test";
+import { CLOSECODES } from "./Constants";
+import { DEFAULT_GATEWAY_HEARTBEAT_TIMEOUT, setHeartbeat } from "./Heartbeat";
+import type { WebSocket } from "./WebSocket";
+
+afterEach(() => {
+    mock.timers.reset();
+});
+
+function createSocket() {
+    const closes: { code: number; reason?: string }[] = [];
+    const socket = {
+        close(code: number, reason?: string) {
+            closes.push({ code, reason });
+        },
+    } as WebSocket;
+
+    return { socket, closes };
+}
+
+describe("setHeartbeat", () => {
+    it("uses the default gateway heartbeat timeout", () => {
+        mock.timers.enable({ apis: ["setTimeout"] });
+        const { socket, closes } = createSocket();
+
+        setHeartbeat(socket);
+        mock.timers.tick(DEFAULT_GATEWAY_HEARTBEAT_TIMEOUT - 1);
+        assert.deepEqual(closes, []);
+
+        mock.timers.tick(1);
+        assert.deepEqual(closes, [{ code: CLOSECODES.Session_timed_out, reason: undefined }]);
+    });
+
+    it("closes the socket after the configured timeout", () => {
+        mock.timers.enable({ apis: ["setTimeout"] });
+        const { socket, closes } = createSocket();
+
+        setHeartbeat(socket, 50);
+        mock.timers.tick(49);
+        assert.deepEqual(closes, []);
+
+        mock.timers.tick(1);
+        assert.deepEqual(closes, [{ code: CLOSECODES.Session_timed_out, reason: undefined }]);
+    });
+
+    it("clears the previous heartbeat timeout before scheduling a new one", () => {
+        mock.timers.enable({ apis: ["setTimeout"] });
+        const { socket, closes } = createSocket();
+
+        setHeartbeat(socket, 50);
+        setHeartbeat(socket, 400);
+        mock.timers.tick(399);
+
+        assert.deepEqual(closes, []);
+        mock.timers.tick(1);
+        assert.deepEqual(closes, [{ code: CLOSECODES.Session_timed_out, reason: undefined }]);
+    });
+});

src/gateway/util/Heartbeat.ts

@@ -18,10 +18,12 @@
 
 import { CLOSECODES } from "./Constants";
 import { WebSocket } from "./WebSocket";
+import { DEFAULT_GATEWAY_HEARTBEAT_TIMEOUT } from "../../util/config/types/GatewayConfiguration";
 
-// TODO: make heartbeat timeout configurable
-export function setHeartbeat(socket: WebSocket) {
+export { DEFAULT_GATEWAY_HEARTBEAT_TIMEOUT } from "../../util/config/types/GatewayConfiguration";
+
+export function setHeartbeat(socket: WebSocket, timeout = DEFAULT_GATEWAY_HEARTBEAT_TIMEOUT) {
     if (socket.heartbeatTimeout) clearTimeout(socket.heartbeatTimeout);
 
-    socket.heartbeatTimeout = setTimeout(() => socket.close(CLOSECODES.Session_timed_out), 1000 * 45);
+    socket.heartbeatTimeout = setTimeout(() => socket.close(CLOSECODES.Session_timed_out), timeout);
 }

src/util/config/Config.ts

@@ -25,6 +25,7 @@ import {
     EmbedConfiguration,
     EndpointConfiguration,
     ExternalTokensConfiguration,
+    GatewayConfiguration,
     GeneralConfiguration,
     GifConfiguration,
     GuildConfiguration,
@@ -42,7 +43,7 @@ import {
 
 export class ConfigValue {
     admin: EndpointConfiguration = new EndpointConfiguration();
-    gateway: EndpointConfiguration = new EndpointConfiguration();
+    gateway: GatewayConfiguration = new GatewayConfiguration();
     cdn: CdnConfiguration = new CdnConfiguration();
     api: ApiConfiguration = new ApiConfiguration();
     general: GeneralConfiguration = new GeneralConfiguration();

src/util/config/index.ts

@@ -17,3 +17,4 @@
 */
 
 export * from "./Config";
+export * from "./types";

src/util/config/types/GatewayConfiguration.test.ts

@@ -0,0 +1,26 @@
+import assert from "node:assert/strict";
+import { describe, it } from "node:test";
+import { DEFAULT_GATEWAY_HEARTBEAT_TIMEOUT, GATEWAY_HEARTBEAT_INTERVAL, GatewayConfiguration, isValidGatewayHeartbeatTimeout } from "./GatewayConfiguration";
+
+describe("GatewayConfiguration", () => {
+    it("keeps endpoint settings and adds heartbeat timeout defaults", () => {
+        const config = new GatewayConfiguration();
+
+        assert.equal(config.endpointPrivate, null);
+        assert.equal(config.endpointPublic, null);
+        assert.equal(config.heartbeatTimeout, DEFAULT_GATEWAY_HEARTBEAT_TIMEOUT);
+        assert.equal(DEFAULT_GATEWAY_HEARTBEAT_TIMEOUT, 45_000);
+        assert.equal(GATEWAY_HEARTBEAT_INTERVAL, 30_000);
+    });
+
+    it("rejects timeout values that would close before the advertised heartbeat interval", () => {
+        assert.equal(isValidGatewayHeartbeatTimeout(GATEWAY_HEARTBEAT_INTERVAL + 1), true);
+        assert.equal(isValidGatewayHeartbeatTimeout(GATEWAY_HEARTBEAT_INTERVAL), false);
+        assert.equal(isValidGatewayHeartbeatTimeout(0), false);
+        assert.equal(isValidGatewayHeartbeatTimeout(-1), false);
+        assert.equal(isValidGatewayHeartbeatTimeout(null), false);
+        assert.equal(isValidGatewayHeartbeatTimeout(Number.NaN), false);
+        assert.equal(isValidGatewayHeartbeatTimeout(Number.POSITIVE_INFINITY), false);
+        assert.equal(isValidGatewayHeartbeatTimeout("45000"), false);
+    });
+});

src/util/config/types/GatewayConfiguration.ts

@@ -0,0 +1,12 @@
+import { EndpointConfiguration } from "./EndpointConfiguration";
+
+export const GATEWAY_HEARTBEAT_INTERVAL = 30_000;
+export const DEFAULT_GATEWAY_HEARTBEAT_TIMEOUT = 45_000;
+
+export function isValidGatewayHeartbeatTimeout(timeout: unknown): timeout is number {
+    return typeof timeout === "number" && Number.isFinite(timeout) && timeout > GATEWAY_HEARTBEAT_INTERVAL;
+}
+
+export class GatewayConfiguration extends EndpointConfiguration {
+    heartbeatTimeout: number = DEFAULT_GATEWAY_HEARTBEAT_TIMEOUT;
+}

src/util/config/types/index.ts

@@ -22,6 +22,7 @@ export * from "./DefaultsConfiguration";
 export * from "./EmailConfiguration";
 export * from "./EndpointConfiguration";
 export * from "./ExternalTokensConfiguration";
+export * from "./GatewayConfiguration";
 export * from "./GeneralConfiguration";
 export * from "./GifConfiguration";
 export * from "./GuildConfiguration";

src/util/util/Config.ts

@@ -19,7 +19,7 @@
 import { existsSync } from "node:fs";
 import fs from "node:fs/promises";
 import { OrmUtils } from "..";
-import { ConfigValue } from "../config";
+import { DEFAULT_GATEWAY_HEARTBEAT_TIMEOUT, GATEWAY_HEARTBEAT_INTERVAL, ConfigValue, isValidGatewayHeartbeatTimeout } from "../config";
 import { ConfigEntity } from "../entities";
 import { JsonValue } from "@protobuf-ts/runtime";
 import { bold, red, redBright } from "picocolors";
@@ -226,6 +226,11 @@ function validateFinalConfig(config: ConfigValue) {
     assertConfig("cdn_endpointPublic", (v) => v != null, 'A valid public CDN endpoint URL, eg. "http://localhost:3003/"');
     assertConfig("cdn_endpointPrivate", (v) => v != null, 'A valid private CDN endpoint URL, eg. "http://localhost:3003/" - must be routable from the API server!');
     assertConfig("gateway_endpointPublic", (v) => v != null, 'A valid public gateway endpoint URL, eg. "ws://localhost:3002/"');
+    assertConfig(
+        "gateway_heartbeatTimeout",
+        isValidGatewayHeartbeatTimeout,
+        `${DEFAULT_GATEWAY_HEARTBEAT_TIMEOUT} (must be greater than the advertised heartbeat interval of ${GATEWAY_HEARTBEAT_INTERVAL}ms)`,
+    );
 
     if (hasErrors) {
         console.error("[Config] Your config has invalid values. Fix them first https://docs.spacebar.chat/setup/server/configuration");