improve game websockt (re)connection (#2584)

Previously, the connection and reconnection logic were identical in
Worker.ts, so clients would need to be re-authorized for cosmetics etc
even when reconnecting. Now, on reconnect, Worker.ts only does
authentication - verifying the jwt is valid.

This will allow clients to require a valid turnstile token when first
connecting, and not when reconnecting after a broken ws connection.

## 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

## Please put your Discord username so you can be contacted if a bug or
regression is found:

evan

src/client/ClientGameRunner.ts

@@ -79,9 +79,17 @@ export function joinLobby(
 
   const transport = new Transport(lobbyConfig, eventBus);
 
+  let hasJoined = false;
+
   const onconnect = () => {
-    console.log(`Joined game lobby ${lobbyConfig.gameID}`);
-    transport.joinGame(0);
+    if (hasJoined) {
+      console.log("rejoining game");
+      transport.rejoinGame(0);
+    } else {
+      hasJoined = true;
+      console.log(`Joining game lobby ${lobbyConfig.gameID}`);
+      transport.joinGame();
+    }
   };
   let terrainLoad: Promise<TerrainMapData> | null = null;
 
@@ -198,7 +206,6 @@ export class ClientGameRunner {
   private isActive = false;
 
   private turnsSeen = 0;
-  private hasJoined = false;
   private lastMousePosition: { x: number; y: number } | null = null;
 
   private lastMessageTime: number = 0;
@@ -322,13 +329,12 @@ export class ClientGameRunner {
 
     const onconnect = () => {
       console.log("Connected to game server!");
-      this.transport.joinGame(this.turnsSeen);
+      this.transport.rejoinGame(this.turnsSeen);
     };
     const onmessage = (message: ServerMessage) => {
       this.lastMessageTime = Date.now();
       if (message.type === "start") {
-        this.hasJoined = true;
-        console.log("starting game!");
+        console.log("starting game! in client game runner");
 
         if (this.gameView.config().isRandomSpawn()) {
           const goToPlayer = () => {
@@ -403,10 +409,6 @@ export class ClientGameRunner {
         );
       }
       if (message.type === "turn") {
-        if (!this.hasJoined) {
-          this.transport.joinGame(0);
-          return;
-        }
         // Track when we receive the turn to calculate delay
         const now = Date.now();
         if (this.lastTickReceiveTime > 0) {
@@ -425,7 +427,10 @@ export class ClientGameRunner {
         }
       }
     };
-    this.transport.connect(onconnect, onmessage);
+    this.transport.updateCallback(onconnect, onmessage);
+    console.log("sending join game");
+    // Rejoin game from the start so we don't miss any turns.
+    this.transport.rejoinGame(0);
   }
 
   public stop() {

src/client/Transport.ts

@@ -17,6 +17,7 @@ import {
   ClientJoinMessage,
   ClientMessage,
   ClientPingMessage,
+  ClientRejoinMessage,
   ClientSendWinnerMessage,
   Intent,
   ServerMessage,
@@ -287,6 +288,14 @@ export class Transport {
     }
   }
 
+  public updateCallback(
+    onconnect: () => void,
+    onmessage: (message: ServerMessage) => void,
+  ) {
+    this.onconnect = onconnect;
+    this.onmessage = onmessage;
+  }
+
   private connectLocal(
     onconnect: () => void,
     onmessage: (message: ServerMessage) => void,
@@ -376,18 +385,27 @@ export class Transport {
     }
   }
 
-  joinGame(numTurns: number) {
+  joinGame() {
     this.sendMsg({
       type: "join",
       gameID: this.lobbyConfig.gameID,
       clientID: this.lobbyConfig.clientID,
-      lastTurn: numTurns,
       token: this.lobbyConfig.token,
       username: this.lobbyConfig.playerName,
       cosmetics: this.lobbyConfig.cosmetics,
     } satisfies ClientJoinMessage);
   }
 
+  rejoinGame(lastTurn: number) {
+    this.sendMsg({
+      type: "rejoin",
+      gameID: this.lobbyConfig.gameID,
+      clientID: this.lobbyConfig.clientID,
+      lastTurn: lastTurn,
+      token: this.lobbyConfig.token,
+    } satisfies ClientRejoinMessage);
+  }
+
   leaveGame() {
     if (this.isLocal) {
       this.localServer.endGame();

src/core/Schemas.ts

@@ -88,6 +88,7 @@ export type ClientMessage =
   | ClientPingMessage
   | ClientIntentMessage
   | ClientJoinMessage
+  | ClientRejoinMessage
   | ClientLogMessage
   | ClientHashMessage;
 export type ServerMessage =
@@ -110,6 +111,7 @@ export type ClientSendWinnerMessage = z.infer<typeof ClientSendWinnerSchema>;
 export type ClientPingMessage = z.infer<typeof ClientPingMessageSchema>;
 export type ClientIntentMessage = z.infer<typeof ClientIntentMessageSchema>;
 export type ClientJoinMessage = z.infer<typeof ClientJoinMessageSchema>;
+export type ClientRejoinMessage = z.infer<typeof ClientRejoinMessageSchema>;
 export type ClientLogMessage = z.infer<typeof ClientLogMessageSchema>;
 export type ClientHashMessage = z.infer<typeof ClientHashSchema>;
 
@@ -529,17 +531,25 @@ export const ClientJoinMessageSchema = z.object({
   clientID: ID,
   token: TokenSchema, // WARNING: PII
   gameID: ID,
-  lastTurn: z.number(), // The last turn the client saw.
   username: UsernameSchema,
   // Server replaces the refs with the actual cosmetic data.
   cosmetics: PlayerCosmeticRefsSchema.optional(),
 });
 
+export const ClientRejoinMessageSchema = z.object({
+  type: z.literal("rejoin"),
+  gameID: ID,
+  clientID: ID,
+  lastTurn: z.number(),
+  token: TokenSchema,
+});
+
 export const ClientMessageSchema = z.discriminatedUnion("type", [
   ClientSendWinnerSchema,
   ClientPingMessageSchema,
   ClientIntentMessageSchema,
   ClientJoinMessageSchema,
+  ClientRejoinMessageSchema,
   ClientLogMessageSchema,
   ClientHashSchema,
 ]);

src/server/Client.ts

@@ -18,7 +18,8 @@ export class Client {
     public readonly flares: string[] | undefined,
     public readonly ip: string,
     public readonly username: string,
-    public readonly ws: WebSocket,
+    public ws: WebSocket,
     public readonly cosmetics: PlayerCosmetics | undefined,
+    public readonly isRejoin: boolean = false,
   ) {}
 }

src/server/GameManager.ts

@@ -1,4 +1,5 @@
 import { Logger } from "winston";
+import WebSocket from "ws";
 import { ServerConfig } from "../core/configuration/Config";
 import {
   Difficulty,
@@ -7,7 +8,7 @@ import {
   GameMode,
   GameType,
 } from "../core/game/Game";
-import { GameConfig, GameID } from "../core/Schemas";
+import { ClientRejoinMessage, GameConfig, GameID } from "../core/Schemas";
 import { Client } from "./Client";
 import { GamePhase, GameServer } from "./GameServer";
 
@@ -25,10 +26,23 @@ export class GameManager {
     return this.games.get(id) ?? null;
   }
 
-  addClient(client: Client, gameID: GameID, lastTurn: number): boolean {
+  joinClient(client: Client, gameID: GameID): boolean {
     const game = this.games.get(gameID);
     if (game) {
-      game.addClient(client, lastTurn);
+      game.joinClient(client);
+      return true;
+    }
+    return false;
+  }
+
+  rejoinClient(
+    ws: WebSocket,
+    persistentID: string,
+    msg: ClientRejoinMessage,
+  ): boolean {
+    const game = this.games.get(msg.gameID);
+    if (game) {
+      game.rejoinClient(ws, persistentID, msg);
       return true;
     }
     return false;

src/server/GameServer.ts

@@ -7,6 +7,7 @@ import { GameType } from "../core/game/Game";
 import {
   ClientID,
   ClientMessageSchema,
+  ClientRejoinMessage,
   ClientSendWinnerMessage,
   GameConfig,
   GameInfo,
@@ -129,7 +130,7 @@ export class GameServer {
     }
   }
 
-  public addClient(client: Client, lastTurn: number) {
+  public joinClient(client: Client) {
     this.websockets.add(client.ws);
     if (this.kickedClients.has(client.clientID)) {
       this.log.warn(`cannot add client, already kicked`, {
@@ -137,6 +138,14 @@ export class GameServer {
       });
       return;
     }
+
+    if (this.allClients.has(client.clientID)) {
+      this.log.warn("cannot add client, already in game", {
+        clientID: client.clientID,
+      });
+      return;
+    }
+
     // Log when lobby creator joins private game
     if (client.clientID === this.lobbyCreatorID) {
       this.log.info("Lobby creator joined", {
@@ -144,11 +153,10 @@ export class GameServer {
         creatorID: this.lobbyCreatorID,
       });
     }
-    this.log.info("client (re)joining game", {
+    this.log.info("client joining game", {
       clientID: client.clientID,
       persistentID: client.persistentID,
       clientIP: ipAnonymize(client.ip),
-      isRejoin: lastTurn > 0,
     });
 
     if (
@@ -185,36 +193,67 @@ export class GameServer {
       }
     }
 
-    // Remove stale client if this is a reconnect
-    const existing = this.activeClients.find(
-      (c) => c.clientID === client.clientID,
-    );
-    if (existing !== undefined) {
-      if (client.persistentID !== existing.persistentID) {
-        this.log.error("persistent ids do not match", {
-          clientID: client.clientID,
-          clientIP: ipAnonymize(client.ip),
-          clientPersistentID: client.persistentID,
-          existingIP: ipAnonymize(existing.ip),
-          existingPersistentID: existing.persistentID,
-        });
-        return;
-      }
-
-      client.lastPing = existing.lastPing;
-      client.reportedWinner = existing.reportedWinner;
-
-      this.activeClients = this.activeClients.filter((c) => c !== existing);
-    }
-
     // Client connection accepted
     this.activeClients.push(client);
     client.lastPing = Date.now();
-
     this.markClientDisconnected(client.clientID, false);
-
     this.allClients.set(client.clientID, client);
+    this.addListeners(client);
 
+    // In case a client joined the game late and missed the start message.
+    if (this._hasStarted) {
+      this.sendStartGameMsg(client.ws, 0);
+    }
+  }
+
+  public rejoinClient(
+    ws: WebSocket,
+    persistentID: string,
+    msg: ClientRejoinMessage,
+  ): void {
+    this.websockets.add(ws);
+
+    if (this.kickedClients.has(msg.clientID)) {
+      this.log.warn("cannot rejoin client, client has been kicked", {
+        clientID: msg.clientID,
+      });
+      return;
+    }
+
+    const client = this.allClients.get(msg.clientID);
+    if (!client) {
+      this.log.warn("cannot rejoin client, existing client not found", {
+        clientID: msg.clientID,
+      });
+      return;
+    }
+
+    if (client.persistentID !== persistentID) {
+      this.log.error("persistent ids do not match", {
+        clientID: msg.clientID,
+        clientPersistentID: persistentID,
+        existingIP: ipAnonymize(client.ip),
+        existingPersistentID: client.persistentID,
+      });
+      return;
+    }
+
+    this.activeClients = this.activeClients.filter(
+      (c) => c.clientID !== msg.clientID,
+    );
+    this.activeClients.push(client);
+    client.lastPing = Date.now();
+    this.markClientDisconnected(msg.clientID, false);
+
+    client.ws = ws;
+    this.addListeners(client);
+
+    if (this._hasStarted) {
+      this.sendStartGameMsg(client.ws, msg.lastTurn);
+    }
+  }
+
+  private addListeners(client: Client) {
     client.ws.removeAllListeners("message");
     client.ws.on("message", async (message: string) => {
       try {
@@ -236,6 +275,13 @@ export class GameServer {
         }
         const clientMsg = parsed.data;
         switch (clientMsg.type) {
+          case "rejoin": {
+            // Client is already connected, no auth required, send start game message if game has started
+            if (this._hasStarted) {
+              this.sendStartGameMsg(client.ws, clientMsg.lastTurn);
+            }
+            break;
+          }
           case "intent": {
             if (clientMsg.intent.clientID !== client.clientID) {
               this.log.warn(
@@ -333,11 +379,6 @@ export class GameServer {
         client.ws.close(1002, "WS_ERR_UNEXPECTED_RSV_1");
       }
     });
-
-    // In case a client joined the game late and missed the start message.
-    if (this._hasStarted) {
-      this.sendStartGameMsg(client.ws, lastTurn);
-    }
   }
 
   public numClients(): number {

src/server/Worker.ts

@@ -317,7 +317,7 @@ export async function startWorker() {
         if (clientMsg.type === "ping") {
           // Ignore ping
           return;
-        } else if (clientMsg.type !== "join") {
+        } else if (clientMsg.type !== "join" && clientMsg.type !== "rejoin") {
           log.warn(
             `Invalid message before join: ${JSON.stringify(clientMsg, replacer)}`,
           );
@@ -342,6 +342,23 @@ export async function startWorker() {
         }
         const { persistentId, claims } = result;
 
+        if (clientMsg.type === "rejoin") {
+          log.info("rejoining game", {
+            gameID: clientMsg.gameID,
+            clientID: clientMsg.clientID,
+            persistentID: persistentId,
+          });
+          const wasFound = gm.rejoinClient(ws, persistentId, clientMsg);
+
+          if (!wasFound) {
+            log.warn(
+              `game ${clientMsg.gameID} not found on worker ${workerId}`,
+            );
+            ws.close(1002, "Game not found");
+          }
+          return;
+        }
+
         let roles: string[] | undefined;
         let flares: string[] | undefined;
 
@@ -402,11 +419,7 @@ export async function startWorker() {
           cosmeticResult.cosmetics,
         );
 
-        const wasFound = gm.addClient(
-          client,
-          clientMsg.gameID,
-          clientMsg.lastTurn,
-        );
+        const wasFound = gm.joinClient(client, clientMsg.gameID);
 
         if (!wasFound) {
           log.info(`game ${clientMsg.gameID} not found on worker ${workerId}`);