FunConnect/server/src/relay.ts

import net from 'net';
import logger from './logger';
import { Room, RoomManager } from './room';
import config from './config';

export class RelayEngine {
  private server: net.Server | null = null;
  private roomManager: RoomManager;
  private hostConnections: Map<string, net.Socket> = new Map();
  private pendingConnections: Map<string, { roomId: string; playerName: string; resolve: (socket: net.Socket) => void }> = new Map();

  constructor(roomManager: RoomManager) {
    this.roomManager = roomManager;
  }

  start(port: number): Promise<void> {
    return new Promise((resolve, reject) => {
      this.server = net.createServer((clientSocket) => {
        this.handleConnection(clientSocket);
      });

      this.server.on('error', (err) => {
        logger.error(`Relay server error: ${err.message}`);
        reject(err);
      });

      this.server.listen(port, '0.0.0.0', () => {
        logger.info(`Relay TCP server listening on port ${port}`);
        resolve();
      });
    });
  }

  private handleConnection(clientSocket: net.Socket): void {
    const remoteAddr = `${clientSocket.remoteAddress}:${clientSocket.remotePort}`;
    logger.info(`New TCP connection from ${remoteAddr}`);

    let buffer = Buffer.alloc(0);
    let identified = false;

    const timeout = setTimeout(() => {
      if (!identified) {
        logger.warn(`Connection from ${remoteAddr} timed out waiting for identification`);
        clientSocket.destroy();
      }
    }, 15000);

    clientSocket.once('data', (data) => {
      clearTimeout(timeout);
      identified = true;
      buffer = Buffer.concat([buffer, data]);
      this.identifyAndRoute(clientSocket, buffer, remoteAddr);
    });

    clientSocket.on('error', (err) => {
      logger.debug(`Socket error from ${remoteAddr}: ${err.message}`);
    });
  }

  private identifyAndRoute(clientSocket: net.Socket, initialData: Buffer, remoteAddr: string): void {
    try {
      const headerStr = initialData.toString('utf8', 0, Math.min(initialData.length, 512));

      if (headerStr.startsWith('FUNMC_HOST:')) {
        this.handleHostConnection(clientSocket, headerStr, remoteAddr);
      } else if (headerStr.startsWith('FUNMC_JOIN:')) {
        this.handlePlayerJoin(clientSocket, headerStr, initialData, remoteAddr);
      } else {
        logger.debug(`Unknown protocol from ${remoteAddr}, closing`);
        clientSocket.destroy();
      }
    } catch (err: any) {
      logger.error(`Error routing connection from ${remoteAddr}: ${err.message}`);
      clientSocket.destroy();
    }
  }

  private handleHostConnection(hostSocket: net.Socket, header: string, remoteAddr: string): void {
    const parts = header.split('\n')[0].replace('FUNMC_HOST:', '').split('|');
    if (parts.length < 2) {
      hostSocket.write('ERROR:INVALID_FORMAT\n');
      hostSocket.destroy();
      return;
    }

    const roomId = parts[0].trim();
    const hostPort = parseInt(parts[1].trim());

    const room = this.roomManager.getRoom(roomId);
    if (!room) {
      hostSocket.write('ERROR:ROOM_NOT_FOUND\n');
      hostSocket.destroy();
      return;
    }

    room.hostSocket = hostSocket;
    this.hostConnections.set(roomId, hostSocket);
    hostSocket.write('OK:HOST_REGISTERED\n');

    logger.info(`Host registered for room ${roomId} from ${remoteAddr}, target port ${hostPort}`);

    hostSocket.on('close', () => {
      logger.info(`Host disconnected from room ${roomId}`);
      this.hostConnections.delete(roomId);
      this.roomManager.deleteRoom(roomId);
    });

    hostSocket.on('error', (err) => {
      logger.error(`Host socket error for room ${roomId}: ${err.message}`);
    });
  }

  private handlePlayerJoin(playerSocket: net.Socket, header: string, initialData: Buffer, remoteAddr: string): void {
    const parts = header.split('\n')[0].replace('FUNMC_JOIN:', '').split('|');
    if (parts.length < 2) {
      playerSocket.write('ERROR:INVALID_FORMAT\n');
      playerSocket.destroy();
      return;
    }

    const roomId = parts[0].trim();
    const playerName = parts[1].trim();
    const password = parts[2]?.trim();

    const room = this.roomManager.getRoom(roomId);
    if (!room) {
      playerSocket.write('ERROR:ROOM_NOT_FOUND\n');
      playerSocket.destroy();
      return;
    }

    if (room.password && room.password !== password) {
      playerSocket.write('ERROR:WRONG_PASSWORD\n');
      playerSocket.destroy();
      return;
    }

    if (room.isFull) {
      playerSocket.write('ERROR:ROOM_FULL\n');
      playerSocket.destroy();
      return;
    }

    const hostSocket = this.hostConnections.get(roomId);
    if (!hostSocket || hostSocket.destroyed) {
      playerSocket.write('ERROR:HOST_OFFLINE\n');
      playerSocket.destroy();
      return;
    }

    const player = room.addPlayer(playerName, playerSocket);
    playerSocket.write('OK:CONNECTED\n');

    this.setupRelay(roomId, player.id, playerSocket, hostSocket, room);
  }

  private setupRelay(roomId: string, playerId: string, playerSocket: net.Socket, hostSocket: net.Socket, room: Room): void {
    const hostAddr = hostSocket.remoteAddress;
    const hostPort = room.hostPort;

    const targetSocket = new net.Socket();

    targetSocket.connect(hostPort, hostAddr || '127.0.0.1', () => {
      logger.info(`Relay established: player ${playerId} <-> host ${hostAddr}:${hostPort} in room ${roomId}`);

      playerSocket.pipe(targetSocket);
      targetSocket.pipe(playerSocket);
    });

    targetSocket.on('error', (err) => {
      logger.error(`Target connection error for player ${playerId}: ${err.message}`);
      playerSocket.destroy();
    });

    targetSocket.on('close', () => {
      playerSocket.destroy();
      room.removePlayer(playerId);
    });

    playerSocket.on('close', () => {
      targetSocket.destroy();
      room.removePlayer(playerId);
    });

    playerSocket.on('error', (err) => {
      logger.debug(`Player socket error: ${err.message}`);
      targetSocket.destroy();
    });
  }

  registerHostDirect(roomId: string, hostSocket: net.Socket): void {
    this.hostConnections.set(roomId, hostSocket);
  }

  stop(): void {
    if (this.server) {
      this.server.close();
      logger.info('Relay TCP server stopped');
    }
    for (const [, socket] of this.hostConnections) {
      socket.destroy();
    }
    this.hostConnections.clear();
  }
}