From 37c939efefa4fe7bed0c8395e55ad083b4a21dac Mon Sep 17 00:00:00 2001
From: silnex <24711610+silnex@users.noreply.github.com>
Date: Wed, 20 Mar 2024 02:10:10 +0800
Subject: [PATCH] websocket pass through path (#341)

---
 index.js          |  3 ++-
 test/websocket.js | 67 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/index.js b/index.js
index 1db3276..4f24ca1 100644
--- a/index.js
+++ b/index.js
@@ -150,11 +150,12 @@ class WebSocketProxy {
   }
 
   findUpstream (request, dest) {
-    const search = new URL(request.url, 'ws://127.0.0.1').search
+    const { search, pathname } = new URL(request.url, 'ws://127.0.0.1')
 
     if (typeof this.wsUpstream === 'string' && this.wsUpstream !== '') {
       const target = new URL(this.wsUpstream)
       target.search = search
+      target.pathname = target.pathname === '/' ? pathname : target.pathname
       return target
     }
 
diff --git a/test/websocket.js b/test/websocket.js
index b08395e..a711895 100644
--- a/test/websocket.js
+++ b/test/websocket.js
@@ -574,3 +574,70 @@ test('multiple websocket upstreams with distinct server options', async (t) => {
     server.close()
   ])
 })
+
+test('keep proxy websocket pathname', 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 path = '/keep/path'
+  const server = Fastify()
+  server.register(proxy, {
+    upstream: `ws://127.0.0.1:${origin.address().port}`,
+    // Start proxy with different upstream, without path
+    wsUpstream: `ws://${host}:${origin.address().port}`,
+    websocket: true
+  })
+
+  await server.listen({ port: 0, host: '127.0.0.1' })
+  t.teardown(server.close.bind(server))
+
+  // Start websocket with different upstream for connect, added path.
+  const ws = new WebSocket(`ws://${host}:${server.server.address().port}${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, path],
+    ['fastify', true, host, path]
+  ])
+
+  await Promise.all([
+    once(ws, 'close'),
+    server.close()
+  ])
+})