ref(server-*): refactor and cleanup server stuffs (#1833)

2026-05-19 14:42:23 +08:00
parent c2bf73c95d
commit c04f4773c5
3 changed files with 219 additions and 145 deletions
@@ -149,26 +149,45 @@ export function selectConsumerPeerId(options: {
 /**
 * Constant-time string comparison that prevents timing attacks (CWE-208).
 *
- * @param {string} a - the first string to compare
- * @param {string} b - the expected value (e.g., the real secret)
- * @returns {boolean} `true` if the strings are equal, `false` otherwise
+ * Compares two strings in constant time to prevent attackers from learning
+ * information about the target string through timing side-channels.
+ *
+ * Use when:
+ * - Comparing authentication tokens or secrets
+ * - Any security-sensitive string comparison
+ *
+ * Expects:
+ * - Both strings are available (no lazy evaluation)
+ *
+ * Returns:
+ * - `true` if the strings are equal, `false` otherwise
 */
 function timingSafeCompare(a: string, b: string): boolean {
  const bufA = Buffer.from(a)
  const bufB = Buffer.from(b)
-  if (bufA.length !== bufB.length) {
-    // Compare against itself to keep constant time, then return false
-    timingSafeEqual(bufA, bufA)
-    // To prevent leaking length information, we perform a dummy comparison on the
-    // expected value, making the execution time dependent on its length.
-    timingSafeEqual(bufB, bufB)
-    return false
-  }

-  return timingSafeEqual(bufA, bufB)
+  // Normalize attacker-controlled input to the expected length
+  // so timingSafeEqual always performs a real comparison.
+  const paddedA = Buffer.alloc(bufB.length)
+
+  bufA.copy(
+    paddedA,
+    0,
+    0,
+    Math.min(bufA.length, bufB.length),
+  )
+
+  return (
+    timingSafeEqual(paddedA, bufB)
+    && bufA.length === bufB.length
+  )
 }

-// helper send function
+/**
+ * Sends an event to a specific peer.
+ * Converts the event to JSON format before transmission.
+ * @internal
+ */
 function send(peer: Peer, event: WebSocketBaseEvent<string, unknown> | string) {
  peer.send(stringifyEvent(event))
 }
@@ -223,17 +242,33 @@ export function normalizeLoggerConfig(options?: AppOptions) {
 /**
 * Creates the H3 websocket application and its in-memory peer registry.
 *
+ * Sets up a complete websocket server with:
+ * - Peer authentication and lifecycle management
+ * - Module registration and discovery (registry sync)
+ * - Consumer-based event routing for load distribution
+ * - Health checking with automatic peer removal on timeout
+ * - Event routing with optional policy-based filtering
+ * - Heartbeat monitoring for liveness detection
+ *
 * Use when:
 * - Embedding the AIRI websocket runtime inside a server process
 * - Spinning up a testable application instance before binding a socket listener
 *
 * Expects:
 * - Caller lifecycle management to invoke `dispose` when the app is no longer needed
+ * - Auth token (if provided) must be validated for all clients
+ * - Routing middleware should be stateless and idempotent
 *
 * Returns:
- * - The H3 app plus cleanup helpers for peer shutdown and timer disposal
+ * - The H3 app at `/ws` endpoint plus cleanup helpers for peer shutdown and timer disposal
+ *
+ * Ownership:
+ * - Manages peer registry and module registry as internal mutable state
+ * - Owns all timers and intervals created during setup
+ * - Consumer orchestrator state is isolated within this function scope
 */
 export function setupApp(options?: AppOptions): { app: H3, closeAllPeers: () => void, dispose: () => void } {
+  // === Configuration & State Initialization ===
  const instanceId = options?.instanceId || optionOrEnv(undefined, 'SERVER_INSTANCE_ID', nanoid())
  const authToken = optionOrEnv(options?.auth?.token, 'AUTHENTICATION_TOKEN', '')

@@ -246,6 +281,7 @@ export function setupApp(options?: AppOptions): { app: H3, closeAllPeers: () =>
    onError: error => appLogger.withError(error).error('an error occurred'),
  })

+  // === Registries & Orchestrators ===
  const peerStore = createServerWsPeerStore<AuthenticatedPeer>()
  const peers = peerStore.peers
  const peersByModule = new Map<string, Map<number | undefined, AuthenticatedPeer>>()
@@ -261,6 +297,7 @@ export function setupApp(options?: AppOptions): { app: H3, closeAllPeers: () =>
  const healthCheckIntervalMs = resolveServerWsHealthCheckIntervalMs(heartbeatTtlMs)
  let disposed = false

+  // === Health Check & Peer Liveness ===
  function broadcastPeerHealthy(peerInfo: AuthenticatedPeer, parentId?: string) {
    if (!peerInfo.name || !peerInfo.identity) {
      return
@@ -284,7 +321,11 @@ export function setupApp(options?: AppOptions): { app: H3, closeAllPeers: () =>
    }
  }

-  function resetRoutingState() {
+  function resetRoutingState(force = false) {
+    if (!force && peers.size > 0) {
+      return
+    }
+
    peers.clear()
    peersByModule.clear()
    consumers.clear()
@@ -334,6 +375,7 @@ export function setupApp(options?: AppOptions): { app: H3, closeAllPeers: () =>
    healthCheckInterval.unref?.()
  }

+  // === Module Registry & Consumer Management ===
  function registerModulePeer(p: AuthenticatedPeer, name: string, index?: number) {
    if (!peersByModule.has(name)) {
      peersByModule.set(name, new Map())
@@ -440,6 +482,7 @@ export function setupApp(options?: AppOptions): { app: H3, closeAllPeers: () =>
      }))
  }

+  // === Broadcasting & Registry Synchronization ===
  function sendRegistrySync(peer: Peer, parentId?: string) {
    send(peer, {
      type: 'registry:modules:sync',
@@ -464,6 +507,8 @@ export function setupApp(options?: AppOptions): { app: H3, closeAllPeers: () =>
    }
  }

+  // === WebSocket Gateway Handler ===
+  // Handles peer lifecycle: open, message, error, close
  const websocketGateway = createGateway({
    handler: {
      open: (peer) => {
@@ -873,7 +918,7 @@ export function setupApp(options?: AppOptions): { app: H3, closeAllPeers: () =>
    dispose: () => {
      clearInterval(healthCheckInterval)
      closeAllPeers()
-      resetRoutingState()
+      resetRoutingState(true)
    },
  })

@@ -881,13 +926,46 @@ export function setupApp(options?: AppOptions): { app: H3, closeAllPeers: () =>

  function closeAllPeers() {
    logger.withFields({ totalPeers: peers.size }).log('closing all peers')
-    for (const peer of Array.from(peers.values())) {
-      logger.withFields({ peer: peer.peer.id, peerName: peer.name }).debug('closing peer')
+
+    for (const peerInfo of Array.from(peers.values())) {
+      logger.withFields({
+        peer: peerInfo.peer.id,
+        peerName: peerInfo.name,
+      }).debug('closing peer')
+
      try {
-        peer.peer.close?.()
+        peerInfo.peer.close?.()
      }
      catch (error) {
-        logger.withFields({ peer: peer.peer.id, peerName: peer.name }).withError(error as Error).debug('failed to close peer during shutdown')
+        logger
+          .withFields({
+            peer: peerInfo.peer.id,
+            peerName: peerInfo.name,
+          })
+          .withError(error as Error)
+          .debug('failed to close peer during shutdown')
+
+        // Leave the peer registered until forced disposal cleanup.
+        continue
+      }
+
+      // Some websocket runtimes may never emit `close`
+      // during abrupt shutdown sequences. Remove peers
+      // synchronously after initiating a successful close
+      // so shutdown cleanup is deterministic.
+      peers.delete(peerInfo.peer.id)
+
+      try {
+        unregisterModulePeer(peerInfo, 'server shutdown')
+      }
+      catch (error) {
+        logger
+          .withFields({
+            peer: peerInfo.peer.id,
+            peerName: peerInfo.name,
+          })
+          .withError(error as Error)
+          .debug('failed to unregister peer during shutdown')
      }
    }
  }
@@ -396,14 +396,17 @@ export class Client<C = undefined> {
      void this.handleMessage(event)
    }

-    ws.onerror = (event: any) => {
+    ws.onerror = (event: unknown) => {
      clearConnectTimer()

      if (!isCurrentSocket()) {
        return
      }

-      const error = event?.error instanceof Error ? event.error : new Error('WebSocket error')
+      // Extract error from WebSocket error event which may vary in shape
+      const error = (event as any)?.error instanceof Error
+        ? (event as any).error
+        : new Error('WebSocket error')
      if (this.connectionAttempt) {
        this.handleSocketFailure(error, ws)
      }
@@ -715,10 +718,13 @@ export class Client<C = undefined> {
          return
        }

-        const modules = (data.data as any)?.modules as Array<{
-          name: string
-          identity?: { id?: string }
-        }> ?? []
+        const syncData = data.data as {
+          modules?: Array<{
+            name: string
+            identity?: { id?: string }
+          }>
+        } | unknown
+        const modules = Array.isArray((syncData as any)?.modules) ? (syncData as any).modules : []

        const selfRegistered = modules.some(
          m => m.name === this.opts.name
@@ -758,8 +764,10 @@ export class Client<C = undefined> {
      return
    }

+    // Cast is necessary here because the Set stores callbacks from potentially different event types,
+    // but we're only calling listeners registered for this specific event type
    const results = await Promise.allSettled(
-      Array.from(listeners).map(listener => Promise.resolve(listener(data as any))),
+      Array.from(listeners).map(listener => Promise.resolve((listener as (data: WebSocketEvent<C>) => void | Promise<void>)(data))),
    )

    for (const result of results) {
@@ -41,45 +41,106 @@ export function createInvalidJsonServerErrorMessage(errorMessage: string) {
  return `invalid JSON, error: ${errorMessage}`
 }

+/**
+ * Error metadata registry for predictable error code classification.
+ * Maps error messages to their error code and classification properties.
+ * @internal
+ */
+const errorMetadataRegistry: Record<string, Omit<ParsedServerErrorMessage, 'message'>> = {
+  [ServerErrorMessages.invalidToken]: {
+    authentication: true,
+    code: 'invalid-token',
+    recoverable: false,
+    terminal: true,
+  },
+  [ServerErrorMessages.notAuthenticated]: {
+    authentication: true,
+    code: 'not-authenticated',
+    recoverable: true,
+    terminal: false,
+  },
+  [ServerErrorMessages.mustAuthenticateBeforeAnnouncing]: {
+    authentication: true,
+    code: 'must-authenticate-before-announcing',
+    recoverable: true,
+    terminal: false,
+  },
+  [ServerErrorMessages.invalidEventFormat]: {
+    authentication: false,
+    code: 'invalid-event-format',
+    recoverable: false,
+    terminal: false,
+  },
+  [ServerErrorMessages.moduleAnnounceNameInvalid]: {
+    authentication: false,
+    code: 'module-announce-name-invalid',
+    recoverable: false,
+    terminal: false,
+  },
+  [ServerErrorMessages.moduleAnnounceIndexInvalid]: {
+    authentication: false,
+    code: 'module-announce-index-invalid',
+    recoverable: false,
+    terminal: false,
+  },
+  [ServerErrorMessages.moduleAnnounceIdentityInvalid]: {
+    authentication: false,
+    code: 'module-announce-identity-invalid',
+    recoverable: false,
+    terminal: false,
+  },
+  [ServerErrorMessages.moduleNotFound]: {
+    authentication: false,
+    code: 'module-not-found',
+    recoverable: false,
+    terminal: false,
+  },
+  [ServerErrorMessages.moduleConsumerEventInvalid]: {
+    authentication: false,
+    code: 'module-consumer-event-invalid',
+    recoverable: false,
+    terminal: false,
+  },
+  [ServerErrorMessages.noConsumerRegistered]: {
+    authentication: false,
+    code: 'no-consumer-registered',
+    recoverable: true,
+    terminal: false,
+  },
+  [ServerErrorMessages.uiConfigureModuleNameInvalid]: {
+    authentication: false,
+    code: 'ui-configure-module-name-invalid',
+    recoverable: false,
+    terminal: false,
+  },
+  [ServerErrorMessages.uiConfigureModuleIndexInvalid]: {
+    authentication: false,
+    code: 'ui-configure-module-index-invalid',
+    recoverable: false,
+    terminal: false,
+  },
+}
+
+/**
+ * Parses a server error message and classifies it.
+ *
+ * Use when:
+ * - Receiving error messages from the server
+ * - Determining whether to retry or give up
+ * - Checking if the error is authentication-related
+ *
+ * Expects:
+ * - Message string that matches one of ServerErrorMessages or starts with 'invalid JSON, error: '
+ *
+ * Returns:
+ * - Parsed error with classification (code, authentication, recoverable, terminal)
+ */
 export function parseServerErrorMessage(message: string): ParsedServerErrorMessage {
-  if (message === ServerErrorMessages.invalidToken) {
-    return {
-      authentication: true,
-      code: 'invalid-token',
-      message,
-      recoverable: false,
-      terminal: true,
-    }
-  }
-
-  if (message === ServerErrorMessages.notAuthenticated) {
-    return {
-      authentication: true,
-      code: 'not-authenticated',
-      message,
-      recoverable: true,
-      terminal: false,
-    }
-  }
-
-  if (message === ServerErrorMessages.mustAuthenticateBeforeAnnouncing) {
-    return {
-      authentication: true,
-      code: 'must-authenticate-before-announcing',
-      message,
-      recoverable: true,
-      terminal: false,
-    }
-  }
-
-  if (message === ServerErrorMessages.invalidEventFormat) {
-    return {
-      authentication: false,
-      code: 'invalid-event-format',
-      message,
-      recoverable: false,
-      terminal: false,
-    }
+  const metadata = Object.hasOwn(errorMetadataRegistry, message)
+    ? errorMetadataRegistry[message]
+    : undefined
+  if (metadata) {
+    return { ...metadata, message }
  }

  if (message.startsWith('invalid JSON, error: ')) {
@@ -92,86 +153,6 @@ export function parseServerErrorMessage(message: string): ParsedServerErrorMessa
    }
  }

-  if (message === ServerErrorMessages.moduleAnnounceNameInvalid) {
-    return {
-      authentication: false,
-      code: 'module-announce-name-invalid',
-      message,
-      recoverable: false,
-      terminal: false,
-    }
-  }
-
-  if (message === ServerErrorMessages.moduleAnnounceIndexInvalid) {
-    return {
-      authentication: false,
-      code: 'module-announce-index-invalid',
-      message,
-      recoverable: false,
-      terminal: false,
-    }
-  }
-
-  if (message === ServerErrorMessages.moduleAnnounceIdentityInvalid) {
-    return {
-      authentication: false,
-      code: 'module-announce-identity-invalid',
-      message,
-      recoverable: false,
-      terminal: false,
-    }
-  }
-
-  if (message === ServerErrorMessages.moduleNotFound) {
-    return {
-      authentication: false,
-      code: 'module-not-found',
-      message,
-      recoverable: false,
-      terminal: false,
-    }
-  }
-
-  if (message === ServerErrorMessages.moduleConsumerEventInvalid) {
-    return {
-      authentication: false,
-      code: 'module-consumer-event-invalid',
-      message,
-      recoverable: false,
-      terminal: false,
-    }
-  }
-
-  if (message === ServerErrorMessages.noConsumerRegistered) {
-    return {
-      authentication: false,
-      code: 'no-consumer-registered',
-      message,
-      recoverable: true,
-      terminal: false,
-    }
-  }
-
-  if (message === ServerErrorMessages.uiConfigureModuleNameInvalid) {
-    return {
-      authentication: false,
-      code: 'ui-configure-module-name-invalid',
-      message,
-      recoverable: false,
-      terminal: false,
-    }
-  }
-
-  if (message === ServerErrorMessages.uiConfigureModuleIndexInvalid) {
-    return {
-      authentication: false,
-      code: 'ui-configure-module-index-invalid',
-      message,
-      recoverable: false,
-      terminal: false,
-    }
-  }
-
  return {
    authentication: false,
    code: 'unknown',
@@ -181,10 +162,17 @@ export function parseServerErrorMessage(message: string): ParsedServerErrorMessa
  }
 }

+/**
+ * Checks if a server error message is authentication-related.
+ */
 export function isAuthenticationServerErrorMessage(message: string) {
  return parseServerErrorMessage(message).authentication
 }

+/**
+ * Checks if a server error message is a terminal authentication error.
+ * Terminal errors should not be retried.
+ */
 export function isTerminalAuthenticationServerErrorMessage(message: string) {
  const parsed = parseServerErrorMessage(message)
  return parsed.authentication && parsed.terminal