optimize WebSocket handling in packages/server-runtime/src/index.ts (#459)

This commit is contained in:
Iro
2025-08-26 12:52:26 +08:00
committed by GitHub
parent f6d6349528
commit d05d809ce3
+83 -45
View File
@@ -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<Record<string, unknown>>) {
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<Record<string, unknown>> | 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<string, AuthenticatedPeer>()
const peersByModule = new Map<string, Map<number | undefined, AuthenticatedPeer>>()
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)
},