refactor: extract websocket handler functions (#1751)

## Description:

The websocket message handler functions have gotten quite large. This
change extracts them to functions in their own file, and extracts the
`"join"` message acceptance logic into its own function allowing for all
cases to be accounted for when we add error messages in #1447.

## Please complete the following:

- [x] I have added screenshots for all UI updates
- [x] I process any text displayed to the user through translateText()
and I've added it to the en.json file
- [x] I have added relevant tests to the test directory
- [x] I confirm I have thoroughly tested these changes and take full
responsibility for any bugs introduced
- [ ] I have read and accepted the CLA agreement (only required once).

src/server/GameServer.ts

@@ -4,7 +4,6 @@ import WebSocket from "ws";
 import { z } from "zod";
 import {
   ClientID,
-  ClientMessageSchema,
   ClientSendWinnerMessage,
   GameConfig,
   GameInfo,
@@ -25,6 +24,7 @@ import { GameType } from "../core/game/Game";
 import { archive } from "./Archive";
 import { Client } from "./Client";
 import { gatekeeper } from "./Gatekeeper";
+import { postJoinMessageHandler } from "./worker/websocket/handler/message/PostJoinHandler";
 export enum GamePhase {
   Lobby = "LOBBY",
   Active = "ACTIVE",
@@ -41,7 +41,7 @@ export class GameServer {
   private turns: Turn[] = [];
   private intents: Intent[] = [];
   public activeClients: Client[] = [];
-  private LobbyCreatorID: string | undefined;
+  lobbyCreatorID: string | undefined;
   private allClients: Map<ClientID, Client> = new Map();
   private clientsDisconnectedStatus: Map<ClientID, boolean> = new Map();
   private _hasStarted = false;
@@ -49,9 +49,9 @@ export class GameServer {
 
   private endTurnIntervalID: ReturnType<typeof setInterval> | undefined;
 
-  private lastPingUpdate = 0;
+  lastPingUpdate = 0;
 
-  private winner: ClientSendWinnerMessage | null = null;
+  winner: ClientSendWinnerMessage | null = null;
 
   // Note: This can be undefined if accessed before the game starts.
   private gameStartInfo!: GameStartInfo;
@@ -60,8 +60,8 @@ export class GameServer {
 
   private _hasPrestarted = false;
 
-  private kickedClients: Set<ClientID> = new Set();
-  private outOfSyncClients: Set<ClientID> = new Set();
+  kickedClients: Set<ClientID> = new Set();
+  outOfSyncClients: Set<ClientID> = new Set();
 
   private websockets: Set<WebSocket> = new Set();
 
@@ -74,7 +74,7 @@ export class GameServer {
     lobbyCreatorID?: string,
   ) {
     this.log = log_.child({ gameID: id });
-    this.LobbyCreatorID = lobbyCreatorID ?? undefined;
+    this.lobbyCreatorID = lobbyCreatorID ?? undefined;
   }
 
   public updateGameConfig(gameConfig: Partial<GameConfig>): void {
@@ -121,9 +121,9 @@ export class GameServer {
       return;
     }
     // Log when lobby creator joins private game
-    if (client.clientID === this.LobbyCreatorID) {
+    if (client.clientID === this.lobbyCreatorID) {
       this.log.info("Lobby creator joined", {
-        creatorID: this.LobbyCreatorID,
+        creatorID: this.lobbyCreatorID,
         gameID: this.id,
       });
     }
@@ -200,122 +200,9 @@ export class GameServer {
     client.ws.removeAllListeners("message");
     client.ws.on(
       "message",
-      gatekeeper.wsHandler(client.ip, async (message: string) => {
-        try {
-          const parsed = ClientMessageSchema.safeParse(JSON.parse(message));
-          if (!parsed.success) {
-            const error = z.prettifyError(parsed.error);
-            this.log.error("Failed to parse client message", error, {
-              clientID: client.clientID,
-            });
-            client.ws.send(
-              JSON.stringify({
-                error,
-                message,
-                type: "error",
-              } satisfies ServerErrorMessage),
-            );
-            client.ws.close(1002, "ClientMessageSchema");
-            return;
-          }
-          const clientMsg = parsed.data;
-          switch (clientMsg.type) {
-            case "intent": {
-              if (clientMsg.intent.clientID !== client.clientID) {
-                this.log.warn(
-                  `client id mismatch, client: ${client.clientID}, intent: ${clientMsg.intent.clientID}`,
-                );
-                return;
-              }
-              switch (clientMsg.intent.type) {
-                case "mark_disconnected": {
-                  this.log.warn(
-                    `Should not receive mark_disconnected intent from client`,
-                  );
-                  return;
-                }
-
-                // Handle kick_player intent via WebSocket
-                case "kick_player": {
-                  const authenticatedClientID = client.clientID;
-
-                  // Check if the authenticated client is the lobby creator
-                  if (authenticatedClientID !== this.LobbyCreatorID) {
-                    this.log.warn(`Only lobby creator can kick players`, {
-                      clientID: authenticatedClientID,
-                      creatorID: this.LobbyCreatorID,
-                      gameID: this.id,
-                      target: clientMsg.intent.target,
-                    });
-                    return;
-                  }
-
-                  // Don't allow lobby creator to kick themselves
-                  if (authenticatedClientID === clientMsg.intent.target) {
-                    this.log.warn(`Cannot kick yourself`, {
-                      clientID: authenticatedClientID,
-                    });
-                    return;
-                  }
-
-                  // Log and execute the kick
-                  this.log.info(`Lobby creator initiated kick of player`, {
-                    creatorID: authenticatedClientID,
-                    gameID: this.id,
-                    kickMethod: "websocket",
-                    target: clientMsg.intent.target,
-                  });
-
-                  this.kickClient(clientMsg.intent.target);
-                  return;
-                }
-                default: {
-                  this.addIntent(clientMsg.intent);
-                  break;
-                }
-              }
-              break;
-            }
-            case "ping": {
-              this.lastPingUpdate = Date.now();
-              client.lastPing = Date.now();
-              break;
-            }
-            case "hash": {
-              client.hashes.set(clientMsg.turnNumber, clientMsg.hash);
-              break;
-            }
-            case "winner": {
-              if (
-                this.outOfSyncClients.has(client.clientID) ||
-                this.kickedClients.has(client.clientID) ||
-                this.winner !== null
-              ) {
-                return;
-              }
-              this.winner = clientMsg;
-              this.archiveGame();
-              break;
-            }
-            default: {
-              this.log.warn(
-                `Unknown message type: ${(clientMsg as any).type}`,
-                {
-                  clientID: client.clientID,
-                },
-              );
-              break;
-            }
-          }
-        } catch (error) {
-          this.log.info(
-            `error handline websocket request in game server: ${error}`,
-            {
-              clientID: client.clientID,
-            },
-          );
-        }
-      }),
+      gatekeeper.wsHandler(client.ip, (message) =>
+        postJoinMessageHandler(this, this.log, client, message),
+      ),
     );
     client.ws.on("close", () => {
       this.log.info("client disconnected", {
@@ -422,7 +309,7 @@ export class GameServer {
     });
   }
 
-  private addIntent(intent: Intent) {
+  addIntent(intent: Intent) {
     this.intents.push(intent);
   }
 
@@ -518,7 +405,7 @@ export class GameServer {
   }
 
   public isPrivateLobbyCreator(clientID: string): boolean {
-    return this.LobbyCreatorID === clientID;
+    return this.lobbyCreatorID === clientID;
   }
 
   phase(): GamePhase {
@@ -666,7 +553,7 @@ export class GameServer {
     });
   }
 
-  private archiveGame() {
+  archiveGame() {
     this.log.info("archiving game", {
       gameID: this.id,
       winner: this.winner?.winner,

src/server/Worker.ts

@@ -9,22 +9,14 @@ import { z } from "zod";
 import { GameEnv } from "../core/configuration/Config";
 import { getServerConfigFromServer } from "../core/configuration/ConfigLoader";
 import { GameType } from "../core/game/Game";
-import {
-  ClientMessageSchema,
-  GameRecord,
-  GameRecordSchema,
-  ID,
-  ServerErrorMessage,
-} from "../core/Schemas";
+import { GameRecord, GameRecordSchema, ID } from "../core/Schemas";
 import { CreateGameInputSchema, GameInputSchema } from "../core/WorkerSchemas";
 import { archive, readGameRecord } from "./Archive";
-import { Client } from "./Client";
 import { GameManager } from "./GameManager";
 import { gatekeeper, LimiterType } from "./Gatekeeper";
-import { getUserMe, verifyClientToken } from "./jwt";
 import { logger } from "./Logger";
-
 import { PrivilegeRefresher } from "./PrivilegeRefresher";
+import { preJoinMessageHandler } from "./worker/websocket/handler/message/PreJoinHandler";
 import { initWorkerMetrics } from "./WorkerMetrics";
 
 const config = getServerConfigFromServer();
@@ -307,159 +299,9 @@ export async function startWorker() {
   wss.on("connection", (ws: WebSocket, req) => {
     ws.on(
       "message",
-      gatekeeper.wsHandler(req, async (message: string) => {
-        const forwarded = req.headers["x-forwarded-for"];
-        const ip = Array.isArray(forwarded)
-          ? forwarded[0]
-          : // eslint-disable-next-line @typescript-eslint/prefer-nullish-coalescing
-            forwarded || req.socket.remoteAddress || "unknown";
-
-        try {
-          // Parse and handle client messages
-          const parsed = ClientMessageSchema.safeParse(
-            JSON.parse(message.toString()),
-          );
-          if (!parsed.success) {
-            const error = z.prettifyError(parsed.error);
-            log.warn("Error parsing client message", error);
-            ws.send(
-              JSON.stringify({
-                error: error.toString(),
-                type: "error",
-              } satisfies ServerErrorMessage),
-            );
-            ws.close(1002, "ClientJoinMessageSchema");
-            return;
-          }
-          const clientMsg = parsed.data;
-
-          if (clientMsg.type === "ping") {
-            // Ignore ping
-            return;
-          } else if (clientMsg.type !== "join") {
-            log.warn(
-              `Invalid message before join: ${JSON.stringify(clientMsg)}`,
-            );
-            return;
-          }
-
-          // Verify this worker should handle this game
-          const expectedWorkerId = config.workerIndex(clientMsg.gameID);
-          if (expectedWorkerId !== workerId) {
-            log.warn(
-              `Worker mismatch: Game ${clientMsg.gameID} should be on worker ${expectedWorkerId}, but this is worker ${workerId}`,
-            );
-            return;
-          }
-
-          // Verify token signature
-          const result = await verifyClientToken(clientMsg.token, config);
-          if (result === false) {
-            log.warn("Unauthorized: Invalid token");
-            ws.close(1002, "Unauthorized");
-            return;
-          }
-          const { persistentId, claims } = result;
-
-          let roles: string[] | undefined;
-          let flares: string[] | undefined;
-
-          const allowedFlares = config.allowedFlares();
-          if (claims === null) {
-            if (allowedFlares !== undefined) {
-              log.warn("Unauthorized: Anonymous user attempted to join game");
-              ws.close(1002, "Unauthorized");
-              return;
-            }
-          } else {
-            // Verify token and get player permissions
-            const result = await getUserMe(clientMsg.token, config);
-            if (result === false) {
-              log.warn("Unauthorized: Invalid session");
-              ws.close(1002, "Unauthorized");
-              return;
-            }
-            roles = result.player.roles;
-            flares = result.player.flares;
-
-            if (allowedFlares !== undefined) {
-              const allowed =
-                allowedFlares.length === 0 ||
-                allowedFlares.some((f) => flares?.includes(f));
-              if (!allowed) {
-                log.warn(
-                  "Forbidden: player without an allowed flare attempted to join game",
-                );
-                ws.close(1002, "Forbidden");
-                return;
-              }
-            }
-          }
-
-          // Check if the flag is allowed
-          if (clientMsg.flag !== undefined) {
-            if (clientMsg.flag.startsWith("!")) {
-              const allowed = privilegeRefresher
-                .get()
-                .isCustomFlagAllowed(clientMsg.flag, flares);
-              if (allowed !== true) {
-                log.warn(`Custom flag ${allowed}: ${clientMsg.flag}`);
-                ws.close(1002, `Custom flag ${allowed}`);
-                return;
-              }
-            }
-          }
-
-          // Check if the pattern is allowed
-          if (clientMsg.pattern !== undefined) {
-            const allowed = privilegeRefresher
-              .get()
-              .isPatternAllowed(clientMsg.pattern, flares);
-            if (allowed !== true) {
-              log.warn(`Pattern ${allowed}: ${clientMsg.pattern}`);
-              ws.close(1002, `Pattern ${allowed}`);
-              return;
-            }
-          }
-
-          // Create client and add to game
-          const client = new Client(
-            clientMsg.clientID,
-            persistentId,
-            claims,
-            roles,
-            flares,
-            ip,
-            clientMsg.username,
-            ws,
-            clientMsg.flag,
-            clientMsg.pattern,
-          );
-
-          const wasFound = gm.addClient(
-            client,
-            clientMsg.gameID,
-            clientMsg.lastTurn,
-          );
-
-          if (!wasFound) {
-            log.info(
-              `game ${clientMsg.gameID} not found on worker ${workerId}`,
-            );
-            // Handle game not found case
-          }
-
-          // Handle other message types
-        } catch (error) {
-          ws.close(1011, "Internal server error");
-          log.warn(
-            `error handling websocket message for ${ipAnonymize(ip)}: ${error}`.substring(
-              0,
-              250,
-            ),
-          );
-        }
-      }),
+      gatekeeper.wsHandler(req, (message) =>
+        preJoinMessageHandler(req, ws, privilegeRefresher, gm, message),
+      ),
     );
 
     ws.on("error", (error: Error) => {

src/server/worker/websocket/handler/message/PostJoinHandler.ts

@@ -0,0 +1,122 @@
+import { Logger } from "winston";
+import { z } from "zod";
+import {
+  ClientMessageSchema,
+  ServerErrorMessage,
+} from "../../../../../core/Schemas";
+import { Client } from "../../../../Client";
+import { GameServer } from "../../../../GameServer";
+
+export async function postJoinMessageHandler(
+  gs: GameServer,
+  log: Logger,
+  client: Client,
+  message: string,
+) {
+  try {
+    const parsed = ClientMessageSchema.safeParse(JSON.parse(message));
+    if (!parsed.success) {
+      const error = z.prettifyError(parsed.error);
+      log.error("Failed to parse client message", error, {
+        clientID: client.clientID,
+      });
+      client.ws.send(
+        JSON.stringify({
+          error,
+          message,
+          type: "error",
+        } satisfies ServerErrorMessage),
+      );
+      client.ws.close(1002, "ClientMessageSchema");
+      return;
+    }
+    const clientMsg = parsed.data;
+    switch (clientMsg.type) {
+      case "intent": {
+        if (clientMsg.intent.clientID !== client.clientID) {
+          log.warn(
+            `client id mismatch, client: ${client.clientID}, intent: ${clientMsg.intent.clientID}`,
+          );
+          return;
+        }
+        switch (clientMsg.intent.type) {
+          case "mark_disconnected": {
+            log.warn(`Should not receive mark_disconnected intent from client`);
+            return;
+          }
+
+          // Handle kick_player intent via WebSocket
+          case "kick_player": {
+            const authenticatedClientID = client.clientID;
+
+            // Check if the authenticated client is the lobby creator
+            if (authenticatedClientID !== gs.lobbyCreatorID) {
+              log.warn(`Only lobby creator can kick players`, {
+                clientID: authenticatedClientID,
+                creatorID: gs.lobbyCreatorID,
+                gameID: gs.id,
+                target: clientMsg.intent.target,
+              });
+              return;
+            }
+
+            // Don't allow lobby creator to kick themselves
+            if (authenticatedClientID === clientMsg.intent.target) {
+              log.warn(`Cannot kick yourself`, {
+                clientID: authenticatedClientID,
+              });
+              return;
+            }
+
+            // Log and execute the kick
+            log.info(`Lobby creator initiated kick of player`, {
+              creatorID: authenticatedClientID,
+              gameID: gs.id,
+              kickMethod: "websocket",
+              target: clientMsg.intent.target,
+            });
+
+            gs.kickClient(clientMsg.intent.target);
+            return;
+          }
+          default: {
+            gs.addIntent(clientMsg.intent);
+            break;
+          }
+        }
+        break;
+      }
+      case "ping": {
+        gs.lastPingUpdate = Date.now();
+        client.lastPing = Date.now();
+        break;
+      }
+      case "hash": {
+        client.hashes.set(clientMsg.turnNumber, clientMsg.hash);
+        break;
+      }
+      case "winner": {
+        if (
+          gs.outOfSyncClients.has(client.clientID) ||
+          gs.kickedClients.has(client.clientID) ||
+          gs.winner !== null
+        ) {
+          return;
+        }
+        gs.winner = clientMsg;
+        gs.archiveGame();
+        break;
+      }
+      default: {
+        log.warn(`Unknown message type: ${(clientMsg as any).type}`, {
+          clientID: client.clientID,
+        });
+        break;
+      }
+    }
+  } catch (error) {
+    log.info(`error handline websocket request in game server: ${error}`, {
+      clientID: client.clientID,
+    });
+  }
+}

src/server/worker/websocket/handler/message/PreJoinHandler.ts

@@ -0,0 +1,257 @@
+import http from "http";
+import ipAnonymize from "ip-anonymize";
+import { WebSocket } from "ws";
+import { z } from "zod";
+import { getServerConfigFromServer } from "../../../../../core/configuration/ConfigLoader";
+import {
+  ClientMessageSchema,
+  ServerErrorMessage,
+} from "../../../../../core/Schemas";
+import { Client } from "../../../../Client";
+import { GameManager } from "../../../../GameManager";
+import { getUserMe, verifyClientToken } from "../../../../jwt";
+import { logger } from "../../../../Logger";
+import { PrivilegeRefresher } from "../../../../PrivilegeRefresher";
+
+const config = getServerConfigFromServer();
+
+const workerId = parseInt(process.env.WORKER_ID ?? "0");
+const log = logger.child({ comp: `w_${workerId}` });
+
+export async function preJoinMessageHandler(
+  req: http.IncomingMessage,
+  ws: WebSocket,
+  privilegeRefresher: PrivilegeRefresher,
+  gm: GameManager,
+  message: string,
+): Promise<void> {
+  const result = await handleJoinMessage(
+    req,
+    ws,
+    privilegeRefresher,
+    gm,
+    message,
+  );
+  if (result === undefined) {
+    // The message was ignored, because it wasn't a "join" message
+    // TODO: Rate limit this
+    return;
+  } else if (result.success === false) {
+    // Join failure
+    const { code, description, error, reason } = result;
+    log.warn(`${reason}: ${description}`, error);
+    if (error) {
+      ws.send(
+        JSON.stringify({
+          error,
+          type: "error",
+        } satisfies ServerErrorMessage),
+      );
+    }
+    ws.close(code, reason);
+  } else {
+    // Join success
+  }
+}
+
+async function handleJoinMessage(
+  req: http.IncomingMessage,
+  ws: WebSocket,
+  privilegeRefresher: PrivilegeRefresher,
+  gm: GameManager,
+  message: string,
+): Promise<
+  | undefined
+  | {
+      success: true;
+    }
+  | {
+      success: false;
+      code: 1002;
+      description: string;
+      error?: string;
+      reason:
+        | "ClientJoinMessageSchema"
+        | "Flag invalid"
+        | "Flag restricted"
+        | "Forbidden"
+        | "Not found"
+        | "Pattern invalid"
+        | "Pattern restricted"
+        | "Pattern unlisted"
+        | "Unauthorized";
+    }
+  | {
+      success: false;
+      code: 1011;
+      reason: "Internal server error";
+      error: string;
+      description: string;
+    }
+> {
+  const forwarded = req.headers["x-forwarded-for"];
+  const ip = Array.isArray(forwarded)
+    ? forwarded[0]
+    : // eslint-disable-next-line @typescript-eslint/prefer-nullish-coalescing
+      forwarded || req.socket.remoteAddress || "unknown";
+
+  try {
+    // Parse and handle client messages
+    const parsed = ClientMessageSchema.safeParse(
+      JSON.parse(message.toString()),
+    );
+    if (!parsed.success) {
+      const error = z.prettifyError(parsed.error);
+      return {
+        code: 1002,
+        description: "Error parsing client message",
+        error,
+        reason: "ClientJoinMessageSchema",
+        success: false,
+      };
+    }
+    const clientMsg = parsed.data;
+
+    if (clientMsg.type === "ping") {
+      // Ignore ping
+      return;
+    } else if (clientMsg.type !== "join") {
+      log.warn(`Invalid message before join: ${JSON.stringify(clientMsg)}`);
+      return;
+    }
+
+    // Verify this worker should handle this game
+    const expectedWorkerId = config.workerIndex(clientMsg.gameID);
+    if (expectedWorkerId !== workerId) {
+      log.warn(
+        `Worker mismatch: Game ${clientMsg.gameID} should be on worker ${expectedWorkerId}, but this is worker ${workerId}`,
+      );
+      return;
+    }
+
+    // Verify token signature
+    const result = await verifyClientToken(clientMsg.token, config);
+    if (result === false) {
+      return {
+        code: 1002,
+        description: "Unauthorized: Invalid token",
+        reason: "Unauthorized",
+        success: false,
+      };
+    }
+    const { persistentId, claims } = result;
+
+    let roles: string[] | undefined;
+    let flares: string[] | undefined;
+
+    const allowedFlares = config.allowedFlares();
+    if (claims === null) {
+      if (allowedFlares !== undefined) {
+        return {
+          code: 1002,
+          description: "Unauthorized: Anonymous user attempted to join game",
+          reason: "Unauthorized",
+          success: false,
+        };
+      }
+    } else {
+      // Verify token and get player permissions
+      const result = await getUserMe(clientMsg.token, config);
+      if (result === false) {
+        return {
+          code: 1002,
+          description: "Unauthorized: Anonymous user attempted to join game",
+          reason: "Unauthorized",
+          success: false,
+        };
+      }
+      roles = result.player.roles;
+      flares = result.player.flares;
+
+      if (allowedFlares !== undefined) {
+        const allowed =
+          allowedFlares.length === 0 ||
+          allowedFlares.some((f) => flares?.includes(f));
+        if (!allowed) {
+          return {
+            code: 1002,
+            description:
+              "Forbidden: player without an allowed flare attempted to join game",
+            reason: "Forbidden",
+            success: false,
+          };
+        }
+      }
+    }
+
+    // Check if the flag is allowed
+    if (clientMsg.flag !== undefined) {
+      if (clientMsg.flag.startsWith("!")) {
+        const allowed = privilegeRefresher
+          .get()
+          .isCustomFlagAllowed(clientMsg.flag, flares);
+        if (allowed !== true) {
+          return {
+            code: 1002,
+            description: clientMsg.flag,
+            reason: `Flag ${allowed}`,
+            success: false,
+          };
+        }
+      }
+    }
+
+    // Check if the pattern is allowed
+    if (clientMsg.pattern !== undefined) {
+      const allowed = privilegeRefresher
+        .get()
+        .isPatternAllowed(clientMsg.pattern, flares);
+      if (allowed !== true) {
+        return {
+          code: 1002,
+          description: clientMsg.pattern,
+          reason: `Pattern ${allowed}`,
+          success: false,
+        };
+      }
+    }
+
+    // Create client
+    const client = new Client(
+      clientMsg.clientID,
+      persistentId,
+      claims,
+      roles,
+      flares,
+      ip,
+      clientMsg.username,
+      ws,
+      clientMsg.flag,
+      clientMsg.pattern,
+    );
+
+    const wasFound = gm.addClient(client, clientMsg.gameID, clientMsg.lastTurn);
+
+    if (!wasFound) {
+      return {
+        code: 1002,
+        description: `game ${clientMsg.gameID} not found on worker ${workerId}`,
+        reason: "Not found",
+        success: false,
+      };
+    }
+
+    // Success
+    return {
+      success: true,
+    };
+  } catch (error) {
+    return {
+      code: 1011,
+      description: `error handling websocket message for ${ipAnonymize(ip)}`,
+      error: error instanceof Error ? error.message : String(error),
+      reason: "Internal server error",
+      success: false,
+    };
+  }
+}