From c04f4773c52360d1d5c3bf7a70ffa722419b30bc Mon Sep 17 00:00:00 2001 From: Iro <155815508+Iro96@users.noreply.github.com> Date: Tue, 19 May 2026 13:42:23 +0700 Subject: [PATCH] ref(server-*): refactor and cleanup server stuffs (#1833) --- packages/server-runtime/src/index.ts | 118 +++++++++++--- packages/server-sdk/src/client.ts | 22 ++- packages/server-shared/src/errors.ts | 224 +++++++++++++-------------- 3 files changed, 219 insertions(+), 145 deletions(-) diff --git a/packages/server-runtime/src/index.ts b/packages/server-runtime/src/index.ts index 0250a23fb..9914359d4 100644 --- a/packages/server-runtime/src/index.ts +++ b/packages/server-runtime/src/index.ts @@ -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) { 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() const peers = peerStore.peers const peersByModule = new Map>() @@ -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') } } } diff --git a/packages/server-sdk/src/client.ts b/packages/server-sdk/src/client.ts index 03e1c2ca1..ecfe7d56e 100644 --- a/packages/server-sdk/src/client.ts +++ b/packages/server-sdk/src/client.ts @@ -396,14 +396,17 @@ export class Client { 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 { 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 { 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) => void | Promise)(data))), ) for (const result of results) { diff --git a/packages/server-shared/src/errors.ts b/packages/server-shared/src/errors.ts index 6830da8fa..597ef6c0b 100644 --- a/packages/server-shared/src/errors.ts +++ b/packages/server-shared/src/errors.ts @@ -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> = { + [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