optimize WebSocket handling in packages/server-runtime/src/index.ts (#459)
This commit is contained in:
@@ -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)
|
||||
},
|
||||
|
||||
Reference in New Issue
Block a user