svelte-mud/src/websocket-server.js

import { WebSocketServer } from 'ws';
import * as net from 'net';
import * as tls from 'tls';
import http from 'http';
import { parse } from 'url';

// Configuration for connection persistence
const CONNECTION_PERSISTENCE_TIMEOUT = 5 * 60 * 1000; // 5 minutes in milliseconds
const HEARTBEAT_INTERVAL = 30 * 1000; // 30 seconds

// Create HTTP server
const server = http.createServer();

// Create WebSocket server
const wss = new WebSocketServer({ noServer: true });

// Active connections and their proxies
const connections = new Map();

// Persistent connections waiting for reconnection
// Key: sessionId, Value: { socket, mudHost, mudPort, useSSL, timeoutId, lastActivity }
const persistentConnections = new Map();

// Generate a unique session ID for persistent connections
function generateSessionId() {
  return `session-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
}

// Clean up a persistent connection
function cleanupPersistentConnection(sessionId) {
  const persistentConn = persistentConnections.get(sessionId);
  if (persistentConn) {
    console.log(`Cleaning up persistent connection for session ${sessionId}`);
    
    // Clear timeout
    if (persistentConn.timeoutId) {
      clearTimeout(persistentConn.timeoutId);
    }
    
    // Close MUD socket
    if (persistentConn.socket && !persistentConn.socket.destroyed) {
      persistentConn.socket.end();
    }
    
    // Remove from map
    persistentConnections.delete(sessionId);
  }
}

// Handle WebSocket connections
wss.on('connection', (ws, req, mudHost, mudPort, useSSL) => {
  console.log(`WebSocket connection established for ${mudHost}:${mudPort} (SSL: ${useSSL})`);
  
  // Create a unique ID for this connection
  const connectionId = `${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
  
  // Check for session ID in query parameters for reconnection
  const url = req.url || '';
  const urlParts = new URL(`http://localhost${url}`);
  const sessionId = urlParts.searchParams.get('sessionId');
  
  // Special handling for test connections
  if (mudHost === 'example.com' && mudPort === '23') {
    console.log('Test connection detected - using echo server mode');
    
    // Send welcome message
    ws.send('Hello from WebSocket test server! This is an echo server.');
    
    // Echo back messages
    ws.on('message', (message) => {
      console.log('Test server received:', message.toString());
      ws.send(`Echo: ${message.toString()}`);
    });
    
    // Handle close
    ws.on('close', () => {
      console.log('Test connection closed');
      connections.delete(connectionId);
    });
    
    // Store the connection (without a socket)
    connections.set(connectionId, { ws, testMode: true });
    return;
  }

  let socket;
  let currentSessionId = sessionId;
  
  // Check if this is a reconnection to an existing persistent session
  if (sessionId && persistentConnections.has(sessionId)) {
    console.log(`Reconnecting to existing session: ${sessionId}`);
    
    const persistentConn = persistentConnections.get(sessionId);
    socket = persistentConn.socket;
    
    // Clear the timeout since client reconnected
    if (persistentConn.timeoutId) {
      clearTimeout(persistentConn.timeoutId);
    }
    
    // Remove from persistent connections (now active again)
    persistentConnections.delete(sessionId);
    
    // Send reconnection notification with session ID in proper JSON format
    ws.send(`[SYSTEM]${JSON.stringify({ type: 'session_resumed', sessionId: sessionId })}`);
  } else {
    // Create new connection
    currentSessionId = generateSessionId();
    console.log(`Creating new session: ${currentSessionId}`);
    
    try {
      // Create a TCP socket connection to the MUD server
      // Use tls for SSL connections, net for regular connections
      socket = useSSL 
        ? tls.connect({ host: mudHost, port: parseInt(mudPort), rejectUnauthorized: false })
        : net.createConnection({ host: mudHost, port: parseInt(mudPort) });
      
      // Add error handler
      socket.on('error', (error) => {
        console.error(`Socket error for ${mudHost}:${mudPort}:`, error.message);
        // Send error to client
        if (ws.readyState === 1) {
          ws.send(Buffer.from(`ERROR: Connection to MUD server failed: ${error.message}\r\n`));
          setTimeout(() => {
            if (ws.readyState === 1) ws.close();
          }, 1000);
        }
        // Remove from connections map
        connections.delete(connectionId);
      });
      
      // Send session ID to client in proper JSON format
      ws.send(`[SYSTEM]${JSON.stringify({ sessionId: currentSessionId })}`);
      
    } catch (error) {
      console.error(`Error creating socket connection: ${error.message}`);
      if (ws.readyState === 1) {
        ws.send(Buffer.from(`ERROR: Failed to connect to MUD server: ${error.message}\r\n`));
        ws.close();
      }
      return;
    }
  }

  // Store the connection
  connections.set(connectionId, { ws, socket, sessionId: currentSessionId });
  
  // Handle data from the MUD server - only in regular mode, not test mode
  if (socket) {
    socket.on('data', (data) => {
      // Check for GMCP data (IAC SB GMCP) - very basic check for debugging
      // IAC = 255, SB = 250, GMCP = 201
      let isGmcp = false;
      for (let i = 0; i < data.length - 2; i++) {
        if (data[i] === 255 && data[i+1] === 250 && data[i+2] === 201) {
          isGmcp = true;
          console.log('WebSocket server: Detected GMCP data in server response');
          break;
        }
      }
      
      // Forward data to the WebSocket client if it's still open
      if (ws.readyState === 1) { // WebSocket.OPEN
        ws.send(data);
        console.log(`WebSocket server: Sent ${data.length} bytes to client${isGmcp ? ' (contains GMCP data)' : ''}`);
      }
    });
  }
  
  // Handle socket close from MUD server - this should trigger cleanup
  if (socket) {
    socket.on('close', () => {
      console.log(`MUD connection closed by server for ${mudHost}:${mudPort}`);
      // Close WebSocket if it's still open
      if (ws.readyState === 1) {
        ws.close();
      }
      // Remove from connections map
      connections.delete(connectionId);
      
      // Also cleanup any persistent connection
      if (currentSessionId) {
        cleanupPersistentConnection(currentSessionId);
      }
    });
  }

  // Handle WebSocket messages (data from client to server)
  ws.on('message', (message) => {
    try {
      // Skip if this is a test connection (already handled in the test mode section)
      const conn = connections.get(connectionId);
      if (conn && conn.testMode) return;
      
      // Check for system messages
      const messageStr = message.toString();
      if (messageStr.startsWith('[SYSTEM]')) {
        try {
          const jsonStr = messageStr.substring(8); // Remove "[SYSTEM]"
          const systemData = JSON.parse(jsonStr);
          
          if (systemData.type === 'explicit_disconnect') {
            console.log(`Received explicit disconnect command for session ${currentSessionId}`);
            // This is an explicit disconnect - don't persist the connection
            if (socket && socket.writable) {
              socket.end();
            }
            if (ws.readyState === 1) {
              ws.close();
            }
            connections.delete(connectionId);
            if (currentSessionId) {
              cleanupPersistentConnection(currentSessionId);
            }
            return;
          }
        } catch (error) {
          console.error('Error parsing system message:', error);
        }
        // Don't forward system messages to the MUD server
        return;
      }
      
      // Legacy support for old disconnect command
      if (messageStr.trim() === '[DISCONNECT]') {
        console.log(`Received legacy disconnect command for session ${currentSessionId}`);
        // This is an explicit disconnect - don't persist the connection
        if (socket && socket.writable) {
          socket.end();
        }
        if (ws.readyState === 1) {
          ws.close();
        }
        connections.delete(connectionId);
        if (currentSessionId) {
          cleanupPersistentConnection(currentSessionId);
        }
        return;
      }
      
      // Check for GMCP data (IAC SB GMCP) in client messages
      let isGmcp = false;
      if (message instanceof Buffer || message instanceof Uint8Array) {
        for (let i = 0; i < message.length - 2; i++) {
          if (message[i] === 255 && message[i+1] === 250 && message[i+2] === 201) {
            isGmcp = true;
            console.log('WebSocket server: Detected GMCP data in client message');
            break;
          }
        }
      }
      
      // Forward data to the MUD server
      // The message might be Buffer, ArrayBuffer, or string
      if (conn && conn.socket && conn.socket.writable) {
        conn.socket.write(message);
        console.log(`WebSocket server: Sent ${message.length} bytes to MUD server${isGmcp ? ' (contains GMCP data)' : ''}`);
      } else {
        console.error('Socket not writable, cannot send data to MUD server');
        if (ws.readyState === 1) { // WebSocket.OPEN
          ws.send(Buffer.from(`ERROR: Cannot send data to MUD server: Socket not connected\r\n`));
        }
      }
    } catch (error) {
      console.error('Error forwarding message to MUD server:', error);
      if (ws.readyState === 1) { // WebSocket.OPEN
        const errorMessage = error instanceof Error ? error.message : String(error);
        ws.send(Buffer.from(`ERROR: Failed to send data to MUD server: ${errorMessage}\r\n`));
      }
    }
  });
  
  // Handle WebSocket close - THIS IS THE KEY CHANGE FOR PERSISTENCE
  ws.on('close', () => {
    console.log(`WebSocket closed for ${mudHost}:${mudPort} (session: ${currentSessionId})`);
    
    const conn = connections.get(connectionId);
    if (conn && !conn.testMode && conn.socket && !conn.socket.destroyed) {
      console.log(`Moving connection to persistent state for ${CONNECTION_PERSISTENCE_TIMEOUT / 1000} seconds`);
      
      // Move the connection to persistent storage instead of closing it
      const timeoutId = setTimeout(() => {
        console.log(`Session ${currentSessionId} timed out, closing MUD connection`);
        cleanupPersistentConnection(currentSessionId);
      }, CONNECTION_PERSISTENCE_TIMEOUT);
      
      persistentConnections.set(currentSessionId, {
        socket: conn.socket,
        mudHost,
        mudPort,
        useSSL,
        timeoutId,
        lastActivity: Date.now()
      });
      
      console.log(`Session ${currentSessionId} will persist for ${CONNECTION_PERSISTENCE_TIMEOUT / 1000} seconds`);
    } else if (conn && conn.socket) {
      // Fallback to immediate cleanup if needed
      conn.socket.end();
    }
    
    // Remove from active connections map
    connections.delete(connectionId);
  });
  
  // Handle WebSocket errors
  ws.on('error', (error) => {
    console.error(`WebSocket error for ${mudHost}:${mudPort}:`, error.message);
    // Close socket on error - but only if it's not going to be persisted
    const conn = connections.get(connectionId);
    if (conn && conn.socket) {
      conn.socket.end();
    }
    // Remove from connections map
    connections.delete(connectionId);
  });
});

// Handle HTTP server upgrade (WebSocket handshake)
server.on('upgrade', (request, socket, head) => {
  // Parse URL to get query parameters
  const { pathname, query } = parse(request.url || '', true);
  
  // Only handle WebSocket connections to /mud-ws
  if (pathname === '/mud-ws') {
    // Extract MUD server details from query parameters
    const { host, port, useSSL } = query;
    
    if (!host || !port) {
      socket.write('HTTP/1.1 400 Bad Request\r\n\r\n');
      socket.destroy();
      return;
    }
    
    // Handle WebSocket upgrade
    wss.handleUpgrade(request, socket, head, (ws) => {
      wss.emit('connection', ws, request, host, port, useSSL === 'true');
    });
  } else {
    // For other upgrades (not to /mud-ws), close the connection
    socket.destroy();
  }
});

// Periodic cleanup of abandoned persistent connections
setInterval(() => {
  const now = Date.now();
  for (const [sessionId, persistentConn] of persistentConnections.entries()) {
    // Clean up connections that have been inactive for too long
    if (now - persistentConn.lastActivity > CONNECTION_PERSISTENCE_TIMEOUT * 2) {
      console.log(`Cleaning up abandoned session: ${sessionId}`);
      cleanupPersistentConnection(sessionId);
    }
  }
}, CONNECTION_PERSISTENCE_TIMEOUT);

// Start the WebSocket server
const PORT = process.env.WS_PORT || 3001;
server.listen(PORT, () => {
  console.log(`WebSocket server is running on port ${PORT}`);
  console.log(`Connection persistence timeout: ${CONNECTION_PERSISTENCE_TIMEOUT / 1000} seconds`);
});

export default server;