fix: websocket upstream with rewrite prefix (#389)

Author: ivan-tymoshenko

index.js

@@ -150,12 +150,11 @@ class WebSocketProxy {
   }
 
   findUpstream (request, dest) {
-    const { search, pathname } = new URL(request.url, 'ws://127.0.0.1')
+    const { search } = new URL(request.url, 'ws://127.0.0.1')
 
     if (typeof this.wsUpstream === 'string' && this.wsUpstream !== '') {
-      const target = new URL(this.wsUpstream)
+      const target = new URL(dest, this.wsUpstream)
       target.search = search
-      target.pathname = target.pathname === '/' ? pathname : target.pathname
       return target
     }
 

test/websocket.js

@@ -543,6 +543,75 @@ test('Proxy websocket with custom upstream url', async (t) => {
   ])
 })
 
+test('Proxy websocket with custom upstream url', async (t) => {
+  t.plan(5)
+
+  const origin = createServer()
+  const wss = new WebSocket.Server({ server: origin })
+
+  t.teardown(wss.close.bind(wss))
+  t.teardown(origin.close.bind(origin))
+
+  const serverMessages = []
+  wss.on('connection', (ws, request) => {
+    ws.on('message', (message, binary) => {
+      // Also need save request.url for check from what url the message is coming.
+      serverMessages.push([message.toString(), binary, request.headers.host.split(':', 1)[0], request.url])
+      ws.send(message, { binary })
+    })
+  })
+
+  await promisify(origin.listen.bind(origin))({ port: 0, host: '127.0.0.1' })
+  // Host for wsUpstream and for later check.
+  const host = '127.0.0.1'
+  // Path for wsUpstream and for later check.
+  const prefix = '/prefix'
+  const rewritePrefix = '/rewrite'
+  const server = Fastify()
+  server.register(proxy, {
+    prefix,
+    rewritePrefix,
+    wsUpstream: `ws://${host}:${origin.address().port}`,
+    websocket: true
+  })
+
+  await server.listen({ port: 0, host: '127.0.0.1' })
+  t.teardown(server.close.bind(server))
+
+  const path = '/some/path'
+  // Start websocket with different upstream for connect, added path.
+  const ws = new WebSocket(`ws://${host}:${server.server.address().port}${prefix}${path}`)
+  await once(ws, 'open')
+
+  const data = [{ message: 'hello', binary: false }, { message: 'fastify', binary: true, isBuffer: true }]
+  const dataLength = data.length
+  let dataIndex = 0
+
+  for (; dataIndex < dataLength; dataIndex++) {
+    const { message: msg, binary, isBuffer } = data[dataIndex]
+    const message = isBuffer
+      ? Buffer.from(msg)
+      : msg
+
+    ws.send(message, { binary })
+
+    const [reply, binaryAnswer] = await once(ws, 'message')
+
+    t.equal(reply.toString(), msg)
+    t.equal(binaryAnswer, binary)
+  }
+  // Also check "path", must be the same.
+  t.strictSame(serverMessages, [
+    ['hello', false, host, rewritePrefix + path],
+    ['fastify', true, host, rewritePrefix + path]
+  ])
+
+  await Promise.all([
+    once(ws, 'close'),
+    server.close()
+  ])
+})
+
 test('multiple websocket upstreams with host constraints', async (t) => {
   t.plan(4)