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';

// Default configuration for connection persistence (fallback values)
const DEFAULT_PERSISTENCE_TIMEOUT = 5 * 60 * 1000; // 5 minutes in milliseconds
const DEFAULT_MAX_BUFFER_MESSAGES = 100; // Maximum number of messages to buffer
const DEFAULT_MAX_BUFFER_SIZE_KB = 10; // Maximum buffer size in KB

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
// Key: connectionId, Value: { ws, socket, sessionId, settings }
const connections = new Map();

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

// Parse connection settings from URL parameters with defaults
function parseConnectionSettings(urlParts) {
  const persistenceTimeoutParam = urlParts.searchParams.get('persistenceTimeout');
  const maxBufferMessagesParam = urlParts.searchParams.get('maxBufferMessages');
  const maxBufferSizeKBParam = urlParts.searchParams.get('maxBufferSizeKB');
  
  return {
    persistenceTimeoutMs: persistenceTimeoutParam ? 
      parseInt(persistenceTimeoutParam) * 60 * 1000 : // Convert minutes to milliseconds
      DEFAULT_PERSISTENCE_TIMEOUT,
    maxBufferMessages: maxBufferMessagesParam ? 
      parseInt(maxBufferMessagesParam) : 
      DEFAULT_MAX_BUFFER_MESSAGES,
    maxBufferSizeKB: maxBufferSizeKBParam ? 
      parseInt(maxBufferSizeKBParam) : 
      DEFAULT_MAX_BUFFER_SIZE_KB
  };
}

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

// Buffer a message for a persistent connection
function bufferMessage(sessionId, data) {
  const persistentConn = persistentConnections.get(sessionId);
  if (!persistentConn) {
    return; // No persistent connection to buffer for
  }
  
  if (!persistentConn.messageBuffer) {
    persistentConn.messageBuffer = [];
    persistentConn.bufferSize = 0;
  }
  
  // Add timestamp to the message
  const bufferedMessage = {
    data: data,
    timestamp: Date.now()
  };
  
  persistentConn.messageBuffer.push(bufferedMessage);
  persistentConn.bufferSize += data.length;
  
  // Use this connection's specific settings for buffer limits
  const settings = persistentConn.settings || {
    maxBufferMessages: DEFAULT_MAX_BUFFER_MESSAGES,
    maxBufferSizeKB: DEFAULT_MAX_BUFFER_SIZE_KB
  };
  
  // Trim buffer if it gets too large
  while (persistentConn.messageBuffer.length > settings.maxBufferMessages || 
         persistentConn.bufferSize > settings.maxBufferSizeKB * 1000) {
    const removed = persistentConn.messageBuffer.shift();
    if (removed) {
      persistentConn.bufferSize -= removed.data.length;
    }
  }
  
  console.log(`Buffered ${data.length} bytes for session ${sessionId} (${persistentConn.messageBuffer.length} messages, ${persistentConn.bufferSize} bytes total, limits: ${settings.maxBufferMessages} msgs/${settings.maxBufferSizeKB}KB)`);
}

// Replay buffered messages to a reconnected client
function replayBufferedMessages(ws, sessionId) {
  const persistentConn = persistentConnections.get(sessionId);
  if (!persistentConn || !persistentConn.messageBuffer) {
    return 0; // No messages to replay
  }
  
  const messages = persistentConn.messageBuffer;
  console.log(`Replaying ${messages.length} buffered messages for session ${sessionId}`);
  
  // Send a notification about message replay
  const replayNotification = `[SYSTEM]${JSON.stringify({ 
    type: 'message_replay_start', 
    messageCount: messages.length,
    timespan: messages.length > 0 ? Date.now() - messages[0].timestamp : 0
  })}`;
  ws.send(replayNotification);
  
  // Send all buffered messages
  for (const message of messages) {
    if (ws.readyState === 1) { // WebSocket.OPEN
      ws.send(message.data);
    }
  }
  
  // Send replay complete notification
  const replayComplete = `[SYSTEM]${JSON.stringify({ type: 'message_replay_complete' })}`;
  ws.send(replayComplete);
  
  // Clear the buffer since messages have been replayed
  const messageCount = messages.length;
  persistentConn.messageBuffer = [];
  persistentConn.bufferSize = 0;
  
  return messageCount;
}

// 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 and settings in query parameters
  const url = req.url || '';
  const urlParts = new URL(`http://localhost${url}`);
  const sessionId = urlParts.searchParams.get('sessionId');
  
  // Parse connection settings for this specific connection
  const connectionSettings = parseConnectionSettings(urlParts);
  
  console.log(`Connection settings for ${connectionId}: timeout=${connectionSettings.persistenceTimeoutMs/1000/60}min, maxMessages=${connectionSettings.maxBufferMessages}, maxSizeKB=${connectionSettings.maxBufferSizeKB}`);
  
  // 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);
    }
    
    // Replay any buffered messages first
    const replayedCount = replayBufferedMessages(ws, sessionId);
    
    // Remove from persistent connections (now active again) - do this after replay
    persistentConnections.delete(sessionId);
    
    // Send reconnection notification with session ID in proper JSON format
    ws.send(`[SYSTEM]${JSON.stringify({ 
      type: 'session_resumed', 
      sessionId: sessionId,
      messagesReplayed: replayedCount
    })}`);
  } 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 with its settings
  connections.set(connectionId, { 
    ws, 
    socket, 
    sessionId: currentSessionId, 
    settings: connectionSettings 
  });
  
  // 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)' : ''}`);
      } else {
        // WebSocket is not open, buffer the message if we have a session
        if (currentSessionId) {
          bufferMessage(currentSessionId, 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 ${conn.settings.persistenceTimeoutMs / 1000} seconds`);
      
      // Move the connection to persistent storage instead of closing it
      // Use this connection's specific timeout setting
      const timeoutId = setTimeout(() => {
        console.log(`Session ${currentSessionId} timed out, closing MUD connection`);
        cleanupPersistentConnection(currentSessionId);
      }, conn.settings.persistenceTimeoutMs);
      
      persistentConnections.set(currentSessionId, {
        socket: conn.socket,
        mudHost,
        mudPort,
        useSSL,
        timeoutId,
        lastActivity: Date.now(),
        messageBuffer: [],
        bufferSize: 0,
        settings: conn.settings  // Store the connection's settings
      });
      
      console.log(`Session ${currentSessionId} will persist for ${conn.settings.persistenceTimeoutMs / 1000} seconds with settings: ${conn.settings.maxBufferMessages} msgs/${conn.settings.maxBufferSizeKB}KB`);
    } 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
    // Use double the connection's specific timeout or default
    const timeoutThreshold = (persistentConn.settings?.persistenceTimeoutMs || DEFAULT_PERSISTENCE_TIMEOUT) * 2;
    if (now - persistentConn.lastActivity > timeoutThreshold) {
      console.log(`Cleaning up abandoned session: ${sessionId}`);
      cleanupPersistentConnection(sessionId);
    }
  }
}, DEFAULT_PERSISTENCE_TIMEOUT); // Run cleanup every default timeout period

// 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(`Default connection persistence timeout: ${DEFAULT_PERSISTENCE_TIMEOUT / 1000} seconds (configurable per connection)`);
});

export default server;