onilib/src/realtime/index.js

A modular Node.js library for real-time online integration in games and web applications

const WebSocket = require('ws');
const { v4: uuidv4 } = require('uuid');
const EventEmitter = require('events');

class RealtimeModule extends EventEmitter {
  constructor(config = {}) {
    super();
    this.config = {
      port: config.port || 8080,
      host: config.host || '0.0.0.0',
      maxConnections: config.maxConnections || 1000,
      heartbeatInterval: config.heartbeatInterval || 30000,
      authRequired: config.authRequired !== false,
      ...config
    };
    
    this.server = null;
    this.clients = new Map();
    this.rooms = new Map();
    this.middlewares = [];
    this.messageHandlers = new Map();
    this.heartbeatInterval = null;
  }

  async initialize(noi) {
    this.noi = noi;
    this.core = noi.core;
    this.auth = noi.getModule('auth');
    
    await this.startServer();
    this.setupHeartbeat();
    this.registerDefaultHandlers();
    
    this.core.log('info', `Realtime module initialized on port ${this.config.port}`);
  }

  async startServer() {
    this.server = new WebSocket.Server({
      port: this.config.port,
      host: this.config.host,
      maxPayload: 1024 * 1024 // 1MB
    });

    this.server.on('connection', this.handleConnection.bind(this));
    this.server.on('error', this.handleServerError.bind(this));
    
    this.core.log('info', `WebSocket server listening on ${this.config.host}:${this.config.port}`);
  }

  async handleConnection(ws, request) {
    const clientId = uuidv4();
    const client = {
      id: clientId,
      ws,
      authenticated: false,
      auth: null,
      rooms: new Set(),
      lastPing: Date.now(),
      metadata: {}
    };

    this.clients.set(clientId, client);
    
    // Set up WebSocket event handlers
    ws.on('message', (data) => this.handleMessage(client, data));
    ws.on('close', () => this.handleDisconnection(client));
    ws.on('error', (error) => this.handleClientError(client, error));
    ws.on('pong', () => {
      client.lastPing = Date.now();
    });

    // Send welcome message
    this.sendToClient(client, {
      type: 'connection',
      clientId,
      timestamp: Date.now()
    });

    this.core.log('debug', `Client connected: ${clientId}`);
    this.emit('client:connected', client);

    // Authenticate if required
    if (this.config.authRequired) {
      setTimeout(() => {
        if (!client.authenticated) {
          this.disconnectClient(client, 'Authentication timeout');
        }
      }, 10000); // 10 second auth timeout
    }
  }

  async handleMessage(client, data) {
    try {
      const message = JSON.parse(data.toString());
      
      // Apply middlewares
      for (const middleware of this.middlewares) {
        const result = await middleware(client, message);
        if (result === false) {
          return; // Middleware blocked the message
        }
      }

      // Handle authentication messages
      if (message.type === 'auth' && !client.authenticated) {
        await this.handleAuthentication(client, message);
        return;
      }

      // Require authentication for other messages
      if (this.config.authRequired && !client.authenticated) {
        this.sendError(client, 'Authentication required');
        return;
      }

      // Handle the message
      const handler = this.messageHandlers.get(message.type);
      if (handler) {
        await handler(client, message);
      } else {
        this.core.log('warn', `Unknown message type: ${message.type}`);
        this.sendError(client, `Unknown message type: ${message.type}`);
      }

      this.emit('message', { client, message });
    } catch (error) {
      this.core.log('error', `Message handling error: ${error.message}`);
      this.sendError(client, 'Invalid message format');
    }
  }

  async handleAuthentication(client, message) {
    try {
      const { strategy = 'jwt', credentials } = message.data;
      
      if (!this.auth) {
        throw new Error('Auth module not available');
      }

      const authResult = await this.auth.authenticate(credentials, strategy);
      
      client.authenticated = true;
      client.auth = authResult;
      
      this.sendToClient(client, {
        type: 'auth_success',
        data: {
          clientId: client.id,
          authenticated: true
        }
      });

      this.core.log('debug', `Client authenticated: ${client.id}`);
      this.emit('client:authenticated', client);
    } catch (error) {
      this.sendError(client, `Authentication failed: ${error.message}`);
      setTimeout(() => {
        this.disconnectClient(client, 'Authentication failed');
      }, 1000);
    }
  }

  handleDisconnection(client) {
    // Remove from all rooms
    for (const roomId of client.rooms) {
      this.leaveRoom(client, roomId);
    }
    
    this.clients.delete(client.id);
    
    this.core.log('debug', `Client disconnected: ${client.id}`);
    this.emit('client:disconnected', client);
  }

  handleClientError(client, error) {
    this.core.log('error', `Client error (${client.id}): ${error.message}`);
    this.emit('client:error', { client, error });
  }

  handleServerError(error) {
    this.core.log('error', `WebSocket server error: ${error.message}`);
    this.emit('server:error', error);
  }

  registerDefaultHandlers() {
    this.registerHandler('ping', (client, message) => {
      this.sendToClient(client, { type: 'pong', timestamp: Date.now() });
    });

    this.registerHandler('join_room', (client, message) => {
      const { roomId } = message.data;
      this.joinRoom(client, roomId);
    });

    this.registerHandler('leave_room', (client, message) => {
      const { roomId } = message.data;
      this.leaveRoom(client, roomId);
    });

    this.registerHandler('room_message', (client, message) => {
      const { roomId, data } = message.data;
      this.sendToRoom(roomId, {
        type: 'room_message',
        from: client.id,
        data
      }, client.id);
    });

    this.registerHandler('direct_message', (client, message) => {
      const { targetId, data } = message.data;
      const target = this.clients.get(targetId);
      if (target) {
        this.sendToClient(target, {
          type: 'direct_message',
          from: client.id,
          data
        });
      }
    });
  }

  registerHandler(type, handler) {
    this.messageHandlers.set(type, handler);
    this.core.log('debug', `Message handler registered: ${type}`);
  }

  use(middleware) {
    this.middlewares.push(middleware);
  }

  joinRoom(client, roomId) {
    if (!this.rooms.has(roomId)) {
      this.rooms.set(roomId, {
        id: roomId,
        clients: new Set(),
        createdAt: Date.now(),
        metadata: {}
      });
    }

    const room = this.rooms.get(roomId);
    room.clients.add(client.id);
    client.rooms.add(roomId);

    this.sendToClient(client, {
      type: 'room_joined',
      data: {
        roomId,
        clientCount: room.clients.size
      }
    });

    // Notify other clients in the room
    this.sendToRoom(roomId, {
      type: 'client_joined',
      data: {
        clientId: client.id,
        roomId
      }
    }, client.id);

    this.core.log('debug', `Client ${client.id} joined room ${roomId}`);
    this.emit('room:joined', { client, roomId });
  }

  leaveRoom(client, roomId) {
    const room = this.rooms.get(roomId);
    if (!room) return;

    room.clients.delete(client.id);
    client.rooms.delete(roomId);

    // Notify other clients in the room
    this.sendToRoom(roomId, {
      type: 'client_left',
      data: {
        clientId: client.id,
        roomId
      }
    });

    // Remove empty rooms
    if (room.clients.size === 0) {
      this.rooms.delete(roomId);
    }

    this.sendToClient(client, {
      type: 'room_left',
      data: { roomId }
    });

    this.core.log('debug', `Client ${client.id} left room ${roomId}`);
    this.emit('room:left', { client, roomId });
  }

  sendToClient(client, message) {
    if (client.ws.readyState === WebSocket.OPEN) {
      client.ws.send(JSON.stringify(message));
    }
  }

  sendToRoom(roomId, message, excludeClientId = null) {
    const room = this.rooms.get(roomId);
    if (!room) return;

    for (const clientId of room.clients) {
      if (clientId !== excludeClientId) {
        const client = this.clients.get(clientId);
        if (client) {
          this.sendToClient(client, message);
        }
      }
    }
  }

  broadcast(message, excludeClientId = null) {
    for (const client of this.clients.values()) {
      if (client.id !== excludeClientId) {
        this.sendToClient(client, message);
      }
    }
  }

  sendError(client, error) {
    this.sendToClient(client, {
      type: 'error',
      error: typeof error === 'string' ? error : error.message,
      timestamp: Date.now()
    });
  }

  disconnectClient(client, reason = 'Disconnected') {
    this.sendToClient(client, {
      type: 'disconnect',
      reason,
      timestamp: Date.now()
    });
    
    setTimeout(() => {
      client.ws.terminate();
    }, 1000);
  }

  setupHeartbeat() {
    this.heartbeatInterval = setInterval(() => {
      const now = Date.now();
      const timeout = this.config.heartbeatInterval * 2;
      
      for (const client of this.clients.values()) {
        if (now - client.lastPing > timeout) {
          this.core.log('debug', `Client ${client.id} heartbeat timeout`);
          client.ws.terminate();
        } else {
          client.ws.ping();
        }
      }
    }, this.config.heartbeatInterval);
  }

  getStats() {
    return {
      connectedClients: this.clients.size,
      activeRooms: this.rooms.size,
      totalRoomClients: Array.from(this.rooms.values())
        .reduce((sum, room) => sum + room.clients.size, 0)
    };
  }

  getRooms() {
    return Array.from(this.rooms.values()).map(room => ({
      id: room.id,
      clientCount: room.clients.size,
      createdAt: room.createdAt,
      metadata: room.metadata
    }));
  }

  getClients() {
    return Array.from(this.clients.values()).map(client => ({
      id: client.id,
      authenticated: client.authenticated,
      rooms: Array.from(client.rooms),
      lastPing: client.lastPing
    }));
  }

  async stop() {
    if (this.heartbeatInterval) {
      clearInterval(this.heartbeatInterval);
    }
    
    if (this.server) {
      this.server.close();
    }
    
    this.core.log('info', 'Realtime module stopped');
  }
}

module.exports = { RealtimeModule };