feat(rpc)!: reject cross-origin WebSocket upgrades (#64)

Author: antfubot

packages/devframe/src/node/server.ts

@@ -56,6 +56,14 @@ export interface StartHttpAndWsOptions {
    * Default: `true`.
    */
   auth?: boolean
+  /**
+   * Extra origins to accept on the WS upgrade beyond the loopback default
+   * (`localhost`/`127.0.0.1`/`::1` and any `Origin`-less request from a
+   * native client). Add your LAN/tunnel origin here when reaching the tool
+   * from another host. Pass `false` to disable origin checking entirely
+   * (not recommended). Default: loopback-only.
+   */
+  allowedOrigins?: readonly string[] | false
   /**
    * Called once the WS server is bound so callers can mount static
    * handlers whose origin depends on the resolved port, or print their
@@ -136,6 +144,7 @@ export async function startHttpAndWs(options: StartHttpAndWsOptions): Promise<St
     // off-route attempts promptly. A shared (caller-owned) server may host
     // other sockets, so leave non-matching upgrades for them.
     destroyUnmatched: ownsHttpServer,
+    allowedOrigins: options.allowedOrigins,
     onDisconnected: (_peer, meta) => {
       rpcHost._emitSessionDisconnected(meta)
     },

packages/devframe/src/rpc/transports/ws-server.ts

@@ -64,6 +64,13 @@ export interface WsRpcTransportOptions {
   destroyUnmatched?: boolean
   /** When set, a new https.Server is created and the WS endpoint is attached to it. */
   https?: HttpsServerOptions
+  /**
+   * Extra origins to accept on the WS upgrade beyond the loopback default.
+   * Add your LAN/tunnel origin here when reaching the tool from another host.
+   * Pass `false` to disable origin checking entirely (not recommended).
+   * Default: loopback-only.
+   */
+  allowedOrigins?: readonly string[] | false
   /**
    * RPC function definitions, used by the per-call wire serializer to
    * dispatch between strict-JSON and structured-clone encoding based
@@ -109,6 +116,32 @@ function pathMatches(a: string, b: string): boolean {
   return strip(a) === strip(b)
 }
 
+export function isLoopbackHostname(hostname: string): boolean {
+  const h = hostname.replace(/^\[|\]$/g, '') // strip IPv6 brackets
+  return h === 'localhost' || h === '127.0.0.1' || h === '::1'
+    || h.endsWith('.localhost') || h.startsWith('127.')
+}
+
+/**
+ * Default origin policy for a localhost dev tool: allow requests with no
+ * `Origin` header (native, non-browser clients), allow any loopback origin
+ * (so cross-port localhost dev setups keep working), and allow explicitly
+ * configured origins. Everything else — a real remote page in the dev's
+ * browser — is rejected.
+ */
+export function isAllowedOrigin(origin: string | undefined, allowedOrigins: readonly string[]): boolean {
+  if (!origin)
+    return true
+  if (allowedOrigins.includes(origin))
+    return true
+  try {
+    return isLoopbackHostname(new URL(origin).hostname)
+  }
+  catch {
+    return false
+  }
+}
+
 /**
  * Route `upgrade` events on a server to the crossws adapter, optionally
  * filtered to a single `path`. Non-matching requests are left untouched so
@@ -121,20 +154,33 @@ function routeUpgrades(
   ws: NodeAdapter,
   path: string | undefined,
   destroyUnmatched: boolean,
+  allowedOrigins: readonly string[] | false | undefined,
 ): () => void {
   const listener = (req: IncomingMessage, socket: Duplex, head: Buffer) => {
+    socket.on('error', () => {
+      // Prevent unhandled ECONNRESET crashes when destroying the socket
+      // or when the client abruptly disconnects.
+    })
+
     if (path) {
       let pathname = req.url ?? '/'
       try {
         pathname = new URL(req.url ?? '/', 'http://localhost').pathname
       }
       catch {}
       if (!pathMatches(pathname, path)) {
-        if (destroyUnmatched)
+        if (destroyUnmatched) {
+          socket.write('HTTP/1.1 400 Bad Request\r\nConnection: close\r\n\r\n')
           socket.destroy()
+        }
         return
       }
     }
+    if (allowedOrigins !== false && !isAllowedOrigin(req.headers.origin, allowedOrigins ?? [])) {
+      socket.write('HTTP/1.1 403 Forbidden\r\nConnection: close\r\n\r\n')
+      socket.destroy()
+      return
+    }
     void ws.handleUpgrade(req, socket, head)
   }
   server.on('upgrade', listener)
@@ -165,6 +211,7 @@ export function attachWsRpcTransport<
     path,
     destroyUnmatched = false,
     https,
+    allowedOrigins,
     onConnected = NOOP,
     onDisconnected = NOOP,
     definitions = EMPTY_DEFS,
@@ -264,11 +311,11 @@ export function attachWsRpcTransport<
   if (server) {
     // Share an existing HTTP(S) server's port. Route upgrades ourselves so we
     // can coexist with the host's own upgrade handlers.
-    detach = routeUpgrades(server, ws, path, destroyUnmatched)
+    detach = routeUpgrades(server, ws, path, destroyUnmatched, allowedOrigins)
   }
   else if (https) {
     ownedServer = createHttpsServer(https)
-    detach = routeUpgrades(ownedServer, ws, path, true)
+    detach = routeUpgrades(ownedServer, ws, path, true, allowedOrigins)
     ownedServer.listen(port, host)
   }
   else {
@@ -278,7 +325,7 @@ export function attachWsRpcTransport<
       res.writeHead(426, { 'content-type': 'text/plain' })
       res.end('Upgrade Required')
     })
-    detach = routeUpgrades(ownedServer, ws, path, true)
+    detach = routeUpgrades(ownedServer, ws, path, true, allowedOrigins)
     ownedServer.listen(port, host)
   }
 

packages/devframe/src/rpc/transports/ws.test.ts

@@ -5,7 +5,7 @@ import { WebSocket, WebSocketServer } from 'ws'
 import { createRpcClient } from '../client'
 import { createRpcServer } from '../server'
 import { createWsRpcChannel } from './ws-client'
-import { attachWsRpcTransport } from './ws-server'
+import { attachWsRpcTransport, isAllowedOrigin, isLoopbackHostname } from './ws-server'
 
 vi.stubGlobal('WebSocket', WebSocket)
 
@@ -162,3 +162,92 @@ describe('devframe rpc', () => {
     }
   })
 })
+
+describe('ws origin check', () => {
+  it('isLoopbackHostname / isAllowedOrigin recognize loopback hosts', () => {
+    expect(isLoopbackHostname('localhost')).toBe(true)
+    expect(isLoopbackHostname('127.0.0.1')).toBe(true)
+    expect(isLoopbackHostname('127.5.5.5')).toBe(true)
+    expect(isLoopbackHostname('::1')).toBe(true)
+    expect(isLoopbackHostname('foo.localhost')).toBe(true)
+    expect(isLoopbackHostname('evil.example')).toBe(false)
+
+    expect(isAllowedOrigin(undefined, [])).toBe(true)
+    expect(isAllowedOrigin('http://localhost:5173', [])).toBe(true)
+    expect(isAllowedOrigin('http://evil.example', [])).toBe(false)
+    expect(isAllowedOrigin('http://evil.example', ['http://evil.example'])).toBe(true)
+  })
+
+  async function connectRaw(url: string, origin?: string): Promise<'open' | 'closed'> {
+    return await new Promise((resolve) => {
+      const ws = new WebSocket(url, origin ? { headers: { origin } } : undefined)
+      ws.on('open', () => {
+        ws.close()
+        resolve('open')
+      })
+      ws.on('error', () => resolve('closed'))
+      ws.on('unexpected-response', () => resolve('closed'))
+      ws.on('close', () => resolve('closed'))
+    })
+  }
+
+  it('rejects a cross-origin browser upgrade', async () => {
+    const HOST = '127.0.0.1'
+    const PORT = await getPort({ host: HOST, random: true })
+    const server = createRpcServer<Record<string, never>, Record<string, never>>({})
+    const { close } = attachWsRpcTransport(server, { port: PORT, host: HOST })
+
+    try {
+      const result = await connectRaw(`ws://${HOST}:${PORT}`, 'http://evil.example')
+      expect(result).toBe('closed')
+    }
+    finally {
+      await close()
+    }
+  })
+
+  it('allows a loopback origin', async () => {
+    const HOST = '127.0.0.1'
+    const PORT = await getPort({ host: HOST, random: true })
+    const server = createRpcServer<Record<string, never>, Record<string, never>>({})
+    const { close } = attachWsRpcTransport(server, { port: PORT, host: HOST })
+
+    try {
+      const result = await connectRaw(`ws://${HOST}:${PORT}`, 'http://localhost:12345')
+      expect(result).toBe('open')
+    }
+    finally {
+      await close()
+    }
+  })
+
+  it('allows a request with no Origin header (native client)', async () => {
+    const HOST = '127.0.0.1'
+    const PORT = await getPort({ host: HOST, random: true })
+    const server = createRpcServer<Record<string, never>, Record<string, never>>({})
+    const { close } = attachWsRpcTransport(server, { port: PORT, host: HOST })
+
+    try {
+      const result = await connectRaw(`ws://${HOST}:${PORT}`)
+      expect(result).toBe('open')
+    }
+    finally {
+      await close()
+    }
+  })
+
+  it('honors allowedOrigins for an otherwise-disallowed origin', async () => {
+    const HOST = '127.0.0.1'
+    const PORT = await getPort({ host: HOST, random: true })
+    const server = createRpcServer<Record<string, never>, Record<string, never>>({})
+    const { close } = attachWsRpcTransport(server, { port: PORT, host: HOST, allowedOrigins: ['http://evil.example'] })
+
+    try {
+      const result = await connectRaw(`ws://${HOST}:${PORT}`, 'http://evil.example')
+      expect(result).toBe('open')
+    }
+    finally {
+      await close()
+    }
+  })
+})