refactor(runtime-rpc): extract the Node WebSocket lifecycle; opt-in pinned port

src/main/runtime/rpc/node-websocket-lifecycle.ts

@@ -0,0 +1,153 @@
+import type { Server as HttpServer } from 'node:http'
+import type { Server as HttpsServer } from 'node:https'
+import type { WebSocket, WebSocketServer } from 'ws'
+import type { RemoteRuntimeServerHeartbeat } from './remote-runtime-server-heartbeat'
+
+type WebSocketMessagePayload = string | Uint8Array<ArrayBufferLike>
+export type WebSocketMessageHandler = {
+  bivarianceHack(
+    msg: WebSocketMessagePayload,
+    reply: (response: string) => void,
+    ws: WebSocket
+  ): void
+}['bivarianceHack']
+
+export type WebSocketConnectionCloseHandler = (
+  clientId: string | null,
+  ws: WebSocket,
+  hasOtherConnections: boolean
+) => void
+
+// Why: WS connections are long-lived and multiplex many RPCs by `id`; auth and dispatch are delegated to the message handler.
+export function attachNodeWebSocketLifecycle(args: {
+  ws: WebSocket
+  heartbeat: RemoteRuntimeServerHeartbeat
+  preAuthTimeoutMs: number
+  preAuthTimers: WeakMap<WebSocket, ReturnType<typeof setTimeout>>
+  clientIds: Map<WebSocket, string>
+  heartbeatConnections: Set<WebSocket>
+  // Why: read lazily so a handler registered after start still reaches sockets accepted earlier.
+  getMessageHandler: () => WebSocketMessageHandler | null
+  getConnectionCloseHandler: () => WebSocketConnectionCloseHandler | null
+}): void {
+  const { ws } = args
+  let finalized = false
+  const onPong = (): void => args.heartbeat.noteAlive(ws)
+  const onMessage = (data: WebSocket.RawData, isBinary: boolean): void => {
+    // Why: any inbound frame counts as proof of life, so an actively-talking client isn't reaped mid-request.
+    args.heartbeat.noteAlive(ws)
+    const message =
+      typeof data === 'string' ? data : isBinary ? toBinaryPayload(data) : data.toString()
+    args.getMessageHandler()?.(
+      message,
+      (response) => {
+        // Why: mobile clients disconnect often; guard the write so we don't throw on a dead socket.
+        if (ws.readyState === ws.OPEN) {
+          ws.send(response)
+        }
+      },
+      ws
+    )
+  }
+  const finalize = (): void => {
+    if (finalized) {
+      return
+    }
+    finalized = true
+    ws.off('pong', onPong)
+    ws.off('message', onMessage)
+    ws.off('close', finalize)
+    ws.off('error', onError)
+    clearNodeWebSocketPreAuthTimer(ws, args.preAuthTimers)
+    args.heartbeatConnections.delete(ws)
+    if (args.heartbeatConnections.size === 0) {
+      args.heartbeat.stop()
+    }
+    const clientId = args.clientIds.get(ws) ?? null
+    args.clientIds.delete(ws)
+    const hasOtherConnections =
+      clientId !== null && Array.from(args.clientIds.values()).includes(clientId)
+    args.getConnectionCloseHandler()?.(clientId, ws, hasOtherConnections)
+  }
+  const onError = (): void => {
+    // Why: close isn't guaranteed after every error path; finalize here too so pre-auth E2EE state and connection ids can't leak.
+    finalize()
+    ws.close()
+  }
+  const preAuthTimer = setTimeout(() => {
+    if (!args.clientIds.has(ws)) {
+      // Why: a silent auto-ponging client would otherwise hold a finite mobile slot forever without starting the E2EE handshake.
+      ws.terminate()
+    }
+  }, args.preAuthTimeoutMs)
+  preAuthTimer.unref?.()
+  args.preAuthTimers.set(ws, preAuthTimer)
+  ws.on('pong', onPong)
+  ws.on('message', onMessage)
+  // Why: clean up connection-scoped state (e.g. mobile-fit overrides) so a dropped phone doesn't leave orphaned phone-fit on desktop.
+  ws.on('close', finalize)
+  ws.on('error', onError)
+  // Why: install lifecycle ownership before periodic heartbeat ticks can observe this socket.
+  args.heartbeatConnections.add(ws)
+  args.heartbeat.noteAlive(ws)
+  if (args.heartbeatConnections.size === 1) {
+    // Unauthenticated sockets are protected by the pre-auth timeout; heartbeat probes begin only
+    // after E2EE binds a client id, avoiding control frames during the handshake.
+    args.heartbeat.start(() => args.clientIds.keys())
+  }
+}
+
+function toBinaryPayload(data: Exclude<WebSocket.RawData, string>): Uint8Array {
+  return Array.isArray(data) ? Buffer.concat(data) : new Uint8Array(data)
+}
+
+export function clearNodeWebSocketPreAuthTimer(
+  ws: WebSocket,
+  preAuthTimers: WeakMap<WebSocket, ReturnType<typeof setTimeout>>
+): void {
+  const timer = preAuthTimers.get(ws)
+  if (timer) {
+    clearTimeout(timer)
+    preAuthTimers.delete(ws)
+  }
+}
+
+export async function stopNodeWebSocketTransport(args: {
+  wss: WebSocketServer | null
+  httpServer: HttpServer | HttpsServer | null
+  heartbeat: RemoteRuntimeServerHeartbeat
+  heartbeatConnections: Set<WebSocket>
+}): Promise<void> {
+  args.heartbeat.stop()
+  args.heartbeatConnections.clear()
+  if (args.wss) {
+    for (const client of args.wss.clients) {
+      // Why: a half-open mobile socket may never answer a close frame, which keeps httpServer.close pending.
+      client.terminate()
+    }
+    args.wss.close()
+  }
+  const httpServer = args.httpServer
+  if (httpServer) {
+    await new Promise<void>((resolve, reject) => {
+      httpServer.close((error) => {
+        if (error) {
+          reject(error)
+          return
+        }
+        resolve()
+      })
+      // Why: idle keep-alive static-web connections would otherwise hold close() open.
+      httpServer.closeAllConnections()
+    })
+  }
+}
+
+// Why: force-terminate soon after the 1013 close since a half-open phone may never ack and would hold the descriptor past the WS cap; the 'error' listener absorbs a reset while closing.
+export function rejectNodeWebSocketOverCapacity(ws: WebSocket): void {
+  ws.on('error', () => {})
+  ws.close(1013, 'Maximum connections reached')
+  const terminateTimer = setTimeout(() => ws.terminate(), 1_000)
+  terminateTimer.unref?.()
+  ws.once('close', () => clearTimeout(terminateTimer))
+}

src/main/runtime/rpc/websocket-transport-limits.ts

@@ -0,0 +1,5 @@
+export const WEBSOCKET_TRANSPORT_MAX_MESSAGE_BYTES = 1024 * 1024
+// Why: one desktop remote-host client can hold many concurrent streams, so keep the cap high enough that stale streams don't starve control RPCs.
+export const WEBSOCKET_TRANSPORT_MAX_CONNECTIONS = 128
+// Why: bound pre-upgrade descriptor use above the WS cap so raw sockets can't grow without bound.
+export const WEBSOCKET_TRANSPORT_MAX_TCP_CONNECTIONS = WEBSOCKET_TRANSPORT_MAX_CONNECTIONS * 2

src/main/runtime/rpc/ws-transport.test.ts

@@ -1,10 +1,12 @@
 import { EventEmitter } from 'node:events'
 import { mkdtempSync } from 'node:fs'
+import { connect } from 'node:net'
 import { tmpdir } from 'node:os'
 import { join } from 'node:path'
 import { describe, expect, it, afterEach, vi } from 'vitest'
 import WebSocket from 'ws'
 import { WebSocketTransport } from './ws-transport'
+import { rejectNodeWebSocketOverCapacity } from './node-websocket-lifecycle'
 import { loadOrCreateTlsCertificate } from '../tls-certificate'
 
 // Why: disable TLS verification for self-signed certs in tests.
@@ -433,9 +435,7 @@ describe('WebSocketTransport', () => {
 
     vi.useFakeTimers()
     try {
-      ;(transport as unknown as { rejectOverCapacity(ws: WebSocket): void }).rejectOverCapacity(
-        serverSocket!
-      )
+      rejectNodeWebSocketOverCapacity(serverSocket!)
       vi.advanceTimersByTime(1_000)
     } finally {
       vi.useRealTimers()
@@ -455,6 +455,24 @@ describe('WebSocketTransport', () => {
     await transport.stop()
   })
 
+  it('stops while an HTTP client holds an unanswered request open', async () => {
+    const transport = new WebSocketTransport({ host: '127.0.0.1', port: 0 })
+    transports.push(transport)
+    await transport.start()
+    const socket = connect(transport.resolvedPort, '127.0.0.1')
+    await new Promise<void>((resolve) => socket.once('connect', () => resolve()))
+    socket.on('error', () => {})
+    // Why: the server's automatic 100 Continue proves the request is parsed and in flight.
+    const continued = new Promise<void>((resolve) => socket.once('data', () => resolve()))
+    socket.write(
+      'POST /unanswered HTTP/1.1\r\nHost: 127.0.0.1\r\nContent-Length: 1\r\nExpect: 100-continue\r\n\r\n'
+    )
+    await continued
+
+    await transport.stop()
+    socket.destroy()
+  })
+
   it('is safe to stop without starting', async () => {
     const { transport } = await createTransport()
     await transport.stop()
@@ -669,6 +687,23 @@ describe('WebSocketTransport', () => {
       expect(transport.resolvedPort).toBe(fallbackPort)
     })
 
+    it('fails closed when a managed tunnel requires the configured port', async () => {
+      const preferredHolder = new WebSocketTransport({ host: '127.0.0.1', port: 0 })
+      transports.push(preferredHolder)
+      await preferredHolder.start()
+      const transport = new WebSocketTransport({
+        host: '127.0.0.1',
+        port: preferredHolder.resolvedPort,
+        fallbackPort: await reserveFreePort(),
+        preferPinnedPort: true,
+        strictPort: true
+      })
+      transports.push(transport)
+
+      await expect(transport.start()).rejects.toMatchObject({ code: 'EADDRINUSE' })
+      expect(transport.resolvedHost).toBeNull()
+    })
+
     it('binds the preferred port when the persisted fallback is taken', async () => {
       const fallbackHolder = new WebSocketTransport({ host: '127.0.0.1', port: 0 })
       transports.push(fallbackHolder)