src/services/websocketManager.js

//
// SPDX-FileCopyrightText: Hadad <[email protected]>
// SPDX-License-Identifier: Apache-2.0
//

import config from '../../config.js';
import { generateId } from '../utils/idGenerator.js';

const clients = new Map();

export const setupWebSocket = (wss) => {
  wss.on('connection', (ws) => {
    const clientId = generateId();
    let clientSessionId = null;
    
    ws.clientId = clientId;
    ws.isAlive = true;
    
    ws.on('pong', () => {
      ws.isAlive = true;
    });
    
    ws.on('message', (data) => {
      try {
        const message = JSON.parse(data.toString());
        
        if (message.type === 'register' && message.sessionId) {
          clientSessionId = message.sessionId;
          
          if (clients.has(clientSessionId)) {
            const oldClient = clients.get(clientSessionId);
            if (oldClient && oldClient.readyState === 1) {
              oldClient.close();
            }
          }
          
          clients.set(clientSessionId, ws);
          ws.sessionId = clientSessionId;
          
          ws.send(JSON.stringify({
            type: 'registered',
            sessionId: clientSessionId
          }));
        }
        
        handleWebSocketMessage(ws, message);
      } catch (error) {}
    });
    
    ws.on('close', () => {
      if (clientSessionId) {
        clients.delete(clientSessionId);
      }
    });
    
    ws.on('error', () => {
      if (clientSessionId) {
        clients.delete(clientSessionId);
      }
    });
    
    ws.send(JSON.stringify({
      type: 'connected',
      clientId
    }));
  });
  
  const interval = setInterval(() => {
    wss.clients.forEach((ws) => {
      if (ws.isAlive === false) {
        return ws.terminate();
      }
      ws.isAlive = false;
      ws.ping();
    });
  }, config.websocket.heartbeat);
  
  wss.on('close', () => {
    clearInterval(interval);
  });
};

const handleWebSocketMessage = (ws, message) => {
  if (message.type === 'ping') {
    ws.send(JSON.stringify({ type: 'pong' }));
  }
};

export const sendToSession = (sessionId, data) => {
  const client = clients.get(sessionId);
  if (client && client.readyState === 1) {
    client.send(JSON.stringify({
      ...data,
      sessionId
    }));
  }
};

export const broadcastToAll = (data) => {
  clients.forEach((client, sessionId) => {
    if (client.readyState === 1) {
      client.send(JSON.stringify({
        ...data,
        sessionId
      }));
    }
  });
};