fix

src/client.ts

@@ -35,7 +35,7 @@ export abstract class Client {
     this.disconnect$ = merge(
       this.disconnectSubject.asObservable(),
       this._onDisconnect(),
-    ).pipe(take(1));
+    ).pipe(take(1), share());
     this.receive$ = this._receive().pipe(
       YGOProProtoPipe(YGOProCtos, {
         onError: (error) => {
@@ -108,7 +108,10 @@ export abstract class Client {
     return this.ip || this.physicalIp() || 'unknown';
   }
 
+  hostname = '';
   name = '';
   vpass = '';
   name_vpass = '';
+
+  established = false;
 }

src/services/client-handler.ts

 import {
+  YGOProCtosBase,
   YGOProCtosExternalAddress,
+  YGOProCtosJoinGame,
   YGOProCtosPlayerInfo,
 } from 'ygopro-msg-encode';
 import { Context } from '../app';
 import { Client } from '../client';
 import { IpResolver } from './ip-resolver';
 import { WsClient } from '../transport/ws/client';
+import { forkJoin, filter, takeUntil, timeout, firstValueFrom } from 'rxjs';
 
 export class ClientHandler {
   constructor(private ctx: Context) {
     this.ctx.middleware(
       YGOProCtosExternalAddress,
       async (msg, client, next) => {
-        if (client instanceof WsClient) {
+        if (client instanceof WsClient || client.ip) {
+          // ws should tell real IP and hostname in http headers, so we skip this step for ws clients
           return next();
         }
         this.ctx
@@ -21,33 +25,67 @@ export class ClientHandler {
             client,
             msg.real_ip === '0.0.0.0' ? undefined : msg.real_ip,
           );
+        client.hostname = msg.hostname?.split(':')[0] || '';
         return next();
       },
     );
 
     this.ctx.middleware(YGOProCtosPlayerInfo, async (msg, client, next) => {
+      if (!client.ip) {
+        this.ctx.get(IpResolver).setClientIp(client);
+      }
       const [name, vpass] = msg.name.split('$');
       client.name = name;
       client.vpass = vpass || '';
       return next();
     });
+    this.ctx.middleware(YGOProCtosBase, async (msg, client, next) => {
+      const isPreHandshakeMsg = [
+        YGOProCtosExternalAddress,
+        YGOProCtosPlayerInfo,
+        YGOProCtosJoinGame,
+      ].some((allowed) => msg instanceof allowed);
+      if (client.established !== isPreHandshakeMsg) {
+        // disallow any messages before handshake is complete, except for the ones needed for handshake
+        return undefined;
+      }
+      return next();
+    });
   }
 
   private logger = this.ctx.createLogger('ClientHandler');
 
   async handleClient(client: Client): Promise<void> {
-    try {
-      client.init().receive$.subscribe(async (msg) => {
-        try {
-          await this.ctx.dispatch(msg, client);
-        } catch (e) {
-          this.logger.warn(
-            `Error dispatching message ${msg.constructor.name} from ${client.loggingIp()}: ${(e as Error).message}`,
-          );
-        }
+    client.init();
+    const receive$ = client.receive$;
+
+    receive$.subscribe(async (msg) => {
+      try {
+        await this.ctx.dispatch(msg, client);
+      } catch (e) {
+        this.logger.warn(
+          `Error dispatching message ${msg.constructor.name} from ${client.loggingIp()}: ${(e as Error).message}`,
+        );
+      }
+    });
+
+    const handshake$ = forkJoin([
+      receive$.pipe(
+        filter((msg) => msg instanceof YGOProCtosPlayerInfo),
+        takeUntil(client.disconnect$),
+      ),
+      receive$.pipe(
+        filter((msg) => msg instanceof YGOProCtosJoinGame),
+        takeUntil(client.disconnect$),
+      ),
+    ]).pipe(timeout(5000), takeUntil(client.disconnect$));
+
+    firstValueFrom(handshake$)
+      .then(() => {
+        client.established = true;
+      })
+      .catch(() => {
+        client.disconnect();
       });
-    } catch {
-      client.disconnect();
-    }
   }
 }

src/transport/ws/server.ts

@@ -68,6 +68,7 @@ export class WsServer {
   private handleConnection(ws: WebSocket, req: IncomingMessage): void {
     const client = new WsClient(this.ctx, ws, req);
     if (this.ctx.get(IpResolver).setClientIp(client, client.xffIp())) return;
+    client.hostname = req.headers.host?.split(':')[0] || '';
     const handler = this.ctx.get(ClientHandler);
     handler.handleClient(client).catch((err) => {
       this.logger.error({ err }, 'Error handling client');