fix: scope websocket upgrades to the proxy prefix

index.js

@@ -16,6 +16,7 @@ const kWs = Symbol('ws')
 const kWsHead = Symbol('wsHead')
 const kWsRewritePrefix = Symbol('wsRewritePrefix')
 const kWsUpgradeListener = Symbol('wsUpgradeListener')
+const kWsPrefixes = Symbol('wsPrefixes')
 
 function liftErrorCode (code) {
   /* c8 ignore start */
@@ -412,6 +413,19 @@ function proxyWebSocketsWithReconnection (logger, source, target, options, hooks
   })
 }
 
+function isUpgradeWithinPrefixes (rawRequest, prefixes) {
+  const pathname = rawRequest.url.split('?', 1)[0]
+  for (const prefix of prefixes) {
+    // Normalise the prefix to its path without a trailing slash. A proxy
+    // mounted at the root ('' or '/') owns every upgrade.
+    const base = prefix.endsWith('/') ? prefix.slice(0, -1) : prefix
+    if (base === '' || pathname === base || pathname.startsWith(`${base}/`)) {
+      return true
+    }
+  }
+  return false
+}
+
 function handleUpgrade (fastify, rawRequest, socket, head) {
   // Save a reference to the socket and then dispatch the request through the normal fastify router so that it will invoke hooks and then eventually a route handler that might upgrade the socket.
   rawRequest[kWs] = socket
@@ -445,10 +459,25 @@ class WebSocketProxy {
     })
 
     if (!fastify.server[kWsUpgradeListener]) {
-      fastify.server[kWsUpgradeListener] = (rawRequest, socket, head) =>
-        handleUpgrade(fastify, rawRequest, socket, head)
+      // A single 'upgrade' listener is shared by every proxy registered on
+      // this server. When the proxy coexists with other 'upgrade' listeners
+      // (e.g. @fastify/websocket), it must only handle upgrades that target
+      // one of the registered proxy prefixes, otherwise it would hijack
+      // WebSocket endpoints it does not own. When the proxy is the only
+      // 'upgrade' listener, every upgrade is dispatched as before so that
+      // out-of-prefix requests are still routed (and rejected) normally.
+      const prefixes = fastify.server[kWsPrefixes] = []
+      fastify.server[kWsUpgradeListener] = (rawRequest, socket, head) => {
+        if (
+          fastify.server.listenerCount('upgrade') === 1 ||
+          isUpgradeWithinPrefixes(rawRequest, prefixes)
+        ) {
+          handleUpgrade(fastify, rawRequest, socket, head)
+        }
+      }
       fastify.server.on('upgrade', fastify.server[kWsUpgradeListener])
     }
+    fastify.server[kWsPrefixes].push(fastify.prefix)
 
     this.handleUpgrade = (request, dest, cb) => {
       wss.handleUpgrade(request.raw, request.raw[kWs], request.raw[kWsHead], (socket) => {

test/ws-upgrade-prefix.js

@@ -0,0 +1,122 @@
+'use strict'
+
+const { test } = require('node:test')
+const assert = require('node:assert')
+const Fastify = require('fastify')
+const proxy = require('../')
+const WebSocket = require('ws')
+const { createServer } = require('node:http')
+const { promisify } = require('node:util')
+const { once } = require('node:events')
+
+// Spin up an upstream echo server that prefixes every message it receives, so
+// the proxied payload can be told apart from a locally handled one.
+async function createEchoUpstream (t) {
+  const origin = createServer()
+  const wss = new WebSocket.Server({ server: origin })
+  wss.on('connection', (ws) => {
+    ws.on('message', (message) => ws.send(`proxied:${message}`))
+  })
+  t.after(() => { wss.close() })
+  t.after(() => { origin.close() })
+  await promisify(origin.listen.bind(origin))({ port: 0, host: '127.0.0.1' })
+  return `ws://127.0.0.1:${origin.address().port}`
+}
+
+// A second, independent WebSocket endpoint mounted directly on the raw server.
+// Its 'upgrade' listener is registered after the proxy's, so in the buggy
+// behaviour the proxy hijacks (and 404s) the socket before this one can run.
+function mountStandalone (t, server, ownedPath) {
+  const wss = new WebSocket.Server({ noServer: true })
+  wss.on('connection', (ws) => {
+    ws.on('message', (message) => ws.send(`standalone:${message}`))
+  })
+  t.after(() => { wss.close() })
+  server.server.on('upgrade', (rawRequest, socket, head) => {
+    if (rawRequest.url.split('?', 1)[0] === ownedPath) {
+      wss.handleUpgrade(rawRequest, socket, head, (ws) => {
+        wss.emit('connection', ws, rawRequest)
+      })
+    }
+  })
+}
+
+async function connect (port, path) {
+  const ws = new WebSocket(`ws://127.0.0.1:${port}${path}`)
+  await once(ws, 'open')
+  ws.send('hello')
+  const [reply] = await once(ws, 'message')
+  ws.close()
+  await once(ws, 'close')
+  return reply.toString()
+}
+
+test('does not hijack websocket upgrades outside a prefixed proxy', async (t) => {
+  const upstream = await createEchoUpstream(t)
+
+  const server = Fastify()
+  server.register(proxy, { prefix: '/proxied', upstream, websocket: true })
+  await server.listen({ port: 0, host: '127.0.0.1' })
+  t.after(() => { server.close() })
+  const port = server.server.address().port
+
+  mountStandalone(t, server, '/standalone')
+
+  // Out of prefix: must reach the standalone handler, untouched by the proxy.
+  assert.strictEqual(await connect(port, '/standalone'), 'standalone:hello')
+  // In prefix (exact match): still proxied.
+  assert.strictEqual(await connect(port, '/proxied'), 'proxied:hello')
+  // In prefix (nested path): still proxied.
+  assert.strictEqual(await connect(port, '/proxied/nested'), 'proxied:hello')
+})
+
+test('a prefix with a trailing slash is scoped correctly', async (t) => {
+  const upstream = await createEchoUpstream(t)
+
+  const server = Fastify()
+  server.register(proxy, { prefix: '/pub/', upstream, websocket: true })
+  await server.listen({ port: 0, host: '127.0.0.1' })
+  t.after(() => { server.close() })
+  const port = server.server.address().port
+
+  mountStandalone(t, server, '/standalone')
+
+  // Out of prefix: handled by the standalone endpoint.
+  assert.strictEqual(await connect(port, '/standalone'), 'standalone:hello')
+  // In prefix: proxied to the upstream.
+  assert.strictEqual(await connect(port, '/pub/nested'), 'proxied:hello')
+})
+
+test('a root proxy still owns upgrades when coexisting with another listener', async (t) => {
+  const upstream = await createEchoUpstream(t)
+
+  const server = Fastify()
+  server.register(proxy, { upstream, websocket: true })
+  await server.listen({ port: 0, host: '127.0.0.1' })
+  t.after(() => { server.close() })
+  const port = server.server.address().port
+
+  // Present only so the proxy is not the sole 'upgrade' listener; it claims a
+  // path the root proxy is never asked about.
+  mountStandalone(t, server, '/reserved')
+
+  // The root proxy owns every path, so an arbitrary path is still proxied.
+  assert.strictEqual(await connect(port, '/anything'), 'proxied:hello')
+})
+
+test('multiple proxies each own their own prefix', async (t) => {
+  const upstream = await createEchoUpstream(t)
+
+  const server = Fastify()
+  server.register(proxy, { prefix: '/pub', upstream, websocket: true })
+  server.register(proxy, { prefix: '/api', upstream, websocket: true })
+  await server.listen({ port: 0, host: '127.0.0.1' })
+  t.after(() => { server.close() })
+  const port = server.server.address().port
+
+  mountStandalone(t, server, '/standalone')
+
+  assert.strictEqual(await connect(port, '/standalone'), 'standalone:hello')
+  assert.strictEqual(await connect(port, '/pub/x'), 'proxied:hello')
+  assert.strictEqual(await connect(port, '/api/x'), 'proxied:hello')
+})