From d05d809ce3e20712de8c3f05678b92481fb00672 Mon Sep 17 00:00:00 2001 From: Iro <155815508+Iro96@users.noreply.github.com> Date: Tue, 26 Aug 2025 11:52:26 +0700 Subject: [PATCH] optimize WebSocket handling in packages/server-runtime/src/index.ts (#459) --- packages/server-runtime/src/index.ts | 128 +++++++++++++++++---------- 1 file changed, 83 insertions(+), 45 deletions(-) diff --git a/packages/server-runtime/src/index.ts b/packages/server-runtime/src/index.ts index 08ad9bc46..919281d8a 100644 --- a/packages/server-runtime/src/index.ts +++ b/packages/server-runtime/src/index.ts @@ -1,17 +1,25 @@ import type { WebSocketEvent } from '@proj-airi/server-shared/types' - import type { AuthenticatedPeer, Peer } from './types' import { env } from 'node:process' - import { Format, LogLevel, setGlobalFormat, setGlobalLogLevel, useLogg } from '@guiiai/logg' import { createApp, createRouter, defineWebSocketHandler } from 'h3' setGlobalFormat(Format.Pretty) setGlobalLogLevel(LogLevel.Log) -function send(peer: Peer, event: WebSocketEvent>) { - peer.send(JSON.stringify(event)) +// cache token once +const AUTH_TOKEN = env.AUTHENTICATION_TOKEN || '' + +// pre-stringified responses +const RESPONSES = { + authenticated: JSON.stringify({ type: 'module:authenticated', data: { authenticated: true } }), + notAuthenticated: JSON.stringify({ type: 'error', data: { message: 'not authenticated' } }), +} + +// helper send function +function send(peer: Peer, event: WebSocketEvent> | string) { + peer.send(typeof event === 'string' ? event : JSON.stringify(event)) } function main() { @@ -26,81 +34,108 @@ function main() { app.use(router) const peers = new Map() + const peersByModule = new Map>() + + function registerModulePeer(p: AuthenticatedPeer, name: string, index?: number) { + if (!peersByModule.has(name)) { + peersByModule.set(name, new Map()) + } + peersByModule.get(name)!.set(index, p) + } + + function unregisterModulePeer(p: AuthenticatedPeer) { + if (!p.name) return + const group = peersByModule.get(p.name) + if (group) { + group.delete(p.index) + if (group.size === 0) { + peersByModule.delete(p.name) + } + } + } router.get('/ws', defineWebSocketHandler({ open: (peer) => { - const token = env.AUTHENTICATION_TOKEN || '' - if (token) { + if (AUTH_TOKEN) { peers.set(peer.id, { peer, authenticated: false, name: '' }) - } - else { - send(peer, { type: 'module:authenticated', data: { authenticated: true } }) + } else { + peer.send(RESPONSES.authenticated) peers.set(peer.id, { peer, authenticated: true, name: '' }) } websocketLogger.withFields({ peer: peer.id, activePeers: peers.size }).log('connected') }, message: (peer, message) => { - const token = env.AUTHENTICATION_TOKEN || '' - const event = message.json() as WebSocketEvent + let event: WebSocketEvent + try { + event = message.json() as WebSocketEvent + } catch (err) { + send(peer, { type: 'error', data: { message: 'invalid JSON' } }) + return + } switch (event.type) { - case 'module:authenticate': - if (!!token && event.data.token !== token) { + case 'module:authenticate': { + if (AUTH_TOKEN && event.data.token !== AUTH_TOKEN) { websocketLogger.withFields({ peer: peer.id }).debug('authentication failed') send(peer, { type: 'error', data: { message: 'invalid token' } }) return } - send(peer, { type: 'module:authenticated', data: { authenticated: true } }) - peers.set(peer.id, { peer, authenticated: true, name: '' }) + peer.send(RESPONSES.authenticated) +const p = peers.get(peer.id); +if (p) { + Object.assign(p, { authenticated: true }); +} return - case 'module:announce': - peers.set(peer.id, { peer, authenticated: true, name: event.data.name }) + } + case 'module:announce': { + const p = peers.get(peer.id) + if (p) { + unregisterModulePeer(p) + Object.assign(p, { authenticated: true, name: event.data.name }) + registerModulePeer(p, p.name, p.index) + } return - case 'ui:configure': - if (event.data.moduleName === '') { - send(peer, { type: 'error', data: { message: 'the field \'moduleName\' can\'t be empty for event \'ui:configure\'' } }) + } + case 'ui:configure': { + const { moduleName, moduleIndex, config } = event.data + + if (moduleName === '') { + send(peer, { type: 'error', data: { message: "the field 'moduleName' can't be empty for event 'ui:configure'" } }) return } - if (typeof event.data.moduleIndex !== 'undefined' && typeof event.data.moduleIndex !== 'number') { - send(peer, { type: 'error', data: { message: 'the field \'moduleIndex\' must be a number for event \'ui:configure\'' } }) + if (typeof moduleIndex !== 'undefined' && typeof moduleIndex !== 'number') { + send(peer, { type: 'error', data: { message: "the field 'moduleIndex' must be a number for event 'ui:configure'" } }) return } - if (typeof event.data.moduleIndex !== 'undefined' && event.data.moduleIndex < 0) { - send(peer, { type: 'error', data: { message: 'the field \'moduleIndex\' must be a positive number for event \'ui:configure\'' } }) + if (typeof moduleIndex !== 'undefined' && moduleIndex < 0) { + send(peer, { type: 'error', data: { message: "the field 'moduleIndex' must be a positive number for event 'ui:configure'" } }) return } - for (const [_id, p] of peers.entries()) { - if (p.name === '') { - continue - } - if (p.name === event.data.moduleName) { - if ((typeof p.index !== 'undefined' && typeof event.data.moduleIndex !== 'undefined' && p.index === event.data.moduleIndex)) { - send(p.peer, { type: 'module:configure', data: { config: event.data.config } }) - return - } - - send(p.peer, { type: 'module:configure', data: { config: event.data.config } }) - return - } - - continue + const target = peersByModule.get(moduleName)?.get(moduleIndex) + if (target) { + send(target.peer, { type: 'module:configure', data: { config } }) + } else { + send(peer, { type: 'error', data: { message: "module not found, it haven't announced it or the name was wrong" } }) } - - send(peer, { type: 'error', data: { message: 'module not found, it haven\'t announced it or the name was wrong' } }) return + } } - if (!peers.get(peer.id)?.authenticated) { + + // default case + const p = peers.get(peer.id) + if (!p?.authenticated) { websocketLogger.withFields({ peer: peer.id }).debug('not authenticated') - send(peer, { type: 'error', data: { message: 'not authenticated' } }) + peer.send(RESPONSES.notAuthenticated) return } - for (const [id, p] of peers.entries()) { + const payload = JSON.stringify(event) + for (const [id, other] of peers.entries()) { if (id !== peer.id) { - p.peer.send(JSON.stringify(event)) + other.peer.send(payload) } } }, @@ -108,6 +143,9 @@ function main() { websocketLogger.withFields({ peer: peer.id }).withError(error).error('an error occurred') }, close: (peer, details) => { + const p = peers.get(peer.id) + if (p) unregisterModulePeer(p) + websocketLogger.withFields({ peer: peer.id, details, activePeers: peers.size }).log('closed') peers.delete(peer.id) },