From d60cc081e3baca1d92bd4c5920036ed692c1aa5b Mon Sep 17 00:00:00 2001
From: Iro <155815508+Iro96@users.noreply.github.com>
Date: Sun, 21 Sep 2025 12:49:11 +0700
Subject: [PATCH] feat(server-runtime): harden WebSocket peer handling and
 validation (#598)

---
 packages/server-runtime/src/index.ts      | 55 +++++++++++++++--------
 packages/server-runtime/src/types/conn.ts | 11 +++++
 2 files changed, 48 insertions(+), 18 deletions(-)

diff --git a/packages/server-runtime/src/index.ts b/packages/server-runtime/src/index.ts
index f54af95b9..5daebb317 100644
--- a/packages/server-runtime/src/index.ts
+++ b/packages/server-runtime/src/index.ts
@@ -1,12 +1,14 @@
 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'
 
+import { env } from 'node:process'
+
+import type { AuthenticatedPeer, Peer } from './types'
+
+import { WebSocketReadyState } from './types'
+
 setGlobalFormat(Format.Pretty)
 setGlobalLogLevel(LogLevel.Log)
 
@@ -42,7 +44,12 @@ function main() {
     if (!peersByModule.has(name)) {
       peersByModule.set(name, new Map())
     }
-    peersByModule.get(name)!.set(index, p)
+    const group = peersByModule.get(name)!
+    if (group.has(index)) {
+      // log instead of silent overwrite
+      websocketLogger.withFields({ name, index }).debug('peer replaced for module')
+    }
+    group.set(index, p)
   }
 
   function unregisterModulePeer(p: AuthenticatedPeer) {
@@ -75,7 +82,8 @@ function main() {
         event = message.json() as WebSocketEvent
       }
       catch (err) {
-        send(peer, { type: 'error', data: { message: `invalid JSON, error: ${err.message}` } })
+        const errorMessage = err instanceof Error ? err.message : String(err)
+        send(peer, { type: 'error', data: { message: `invalid JSON, error: ${errorMessage}` } })
         return
       }
 
@@ -104,12 +112,16 @@ function main() {
               return
             }
             if (typeof index !== 'undefined') {
-              if (typeof index !== 'number' || index < 0) {
-                send(peer, { type: 'error', data: { message: 'the field \'index\' must be a non-negative number for event \'module:announce\'' } })
+              if (!Number.isInteger(index) || index < 0) {
+                send(peer, { type: 'error', data: { message: 'the field \'index\' must be a non-negative integer for event \'module:announce\'' } })
                 return
               }
             }
-            Object.assign(p, { authenticated: true, name, index })
+            if (AUTH_TOKEN && !p.authenticated) {
+              send(peer, { type: 'error', data: { message: 'must authenticate before announcing' } })
+              return
+            }
+            Object.assign(p, { name, index })
             registerModulePeer(p, p.name, p.index)
           }
           return
@@ -122,13 +134,14 @@ function main() {
             send(peer, { type: 'error', data: { message: 'the field \'moduleName\' can\'t be empty for event \'ui:configure\'' } })
             return
           }
-          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 moduleIndex !== 'undefined' && moduleIndex < 0) {
-            send(peer, { type: 'error', data: { message: 'the field \'moduleIndex\' must be a positive number for event \'ui:configure\'' } })
-            return
+          if (typeof moduleIndex !== 'undefined') {
+            if (!Number.isInteger(moduleIndex) || moduleIndex < 0) {
+              send(peer, {
+                type: 'error',
+                data: { message: 'the field \'moduleIndex\' must be a non-negative integer for event \'ui:configure\'' },
+              })
+              return
+            }
           }
 
           const target = peersByModule.get(moduleName)?.get(moduleIndex)
@@ -136,7 +149,7 @@ function main() {
             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 hasn\'t announced itself or the name is incorrect' } })
           }
           return
         }
@@ -152,9 +165,15 @@ function main() {
 
       const payload = JSON.stringify(event)
       for (const [id, other] of peers.entries()) {
-        if (id !== peer.id) {
+        if (id === peer.id)
+          continue
+        if (other.peer.readyState === WebSocketReadyState.OPEN) {
           other.peer.send(payload)
         }
+        else {
+          peers.delete(id)
+          unregisterModulePeer(other)
+        }
       }
     },
     error: (peer, error) => {
diff --git a/packages/server-runtime/src/types/conn.ts b/packages/server-runtime/src/types/conn.ts
index 4a080922c..59ce0a7d6 100644
--- a/packages/server-runtime/src/types/conn.ts
+++ b/packages/server-runtime/src/types/conn.ts
@@ -6,6 +6,10 @@ export interface Peer {
   send: (data: unknown, options?: {
     compress?: boolean
   }) => number | void | undefined
+  /**
+   * WebSocket lifecycle state (mirrors WebSocket.readyState)
+   */
+  readyState?: number
 }
 
 export interface NamedPeer {
@@ -14,6 +18,13 @@ export interface NamedPeer {
   peer: Peer
 }
 
+export enum WebSocketReadyState {
+  CONNECTING = 0,
+  OPEN = 1,
+  CLOSING = 2,
+  CLOSED = 3,
+}
+
 export interface AuthenticatedPeer extends NamedPeer {
   authenticated: boolean
 }