File size: 2,464 Bytes
4f33245
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
//
// 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
      }));
    }
  });
};