fix: The target websocket watch the specific url instead of upstream(… (#308)

Author: Dokki

index.js

@@ -133,11 +133,12 @@ class WebSocketProxy {
     this.prefixList = []
   }
 
-  addUpstream (prefix, rewritePrefix, upstream, wsClientOptions) {
+  addUpstream (prefix, rewritePrefix, upstream, wsUpstream, wsClientOptions) {
     this.prefixList.push({
       prefix: new URL(prefix, 'ws://127.0.0.1').pathname,
       rewritePrefix,
       upstream: convertUrlToWebSocket(upstream),
+      wsUpstream: wsUpstream ? convertUrlToWebSocket(wsUpstream) : '',
       wsClientOptions
     })
 
@@ -148,18 +149,15 @@ class WebSocketProxy {
   findUpstream (request) {
     const source = new URL(request.url, 'ws://127.0.0.1')
 
-    for (const { prefix, rewritePrefix, upstream, wsClientOptions } of this.prefixList) {
-      // If the "upstream" have path then need get the Base of url, otherwise the "target" path will be broken.
-      // Example: upstream is "ws://localhost:22/some/path" and after this code
-      // "target = new URL(source.pathname.replace(prefix, rewritePrefix), upstream)"
-      // The target.pathname will be "some/some/path"
-      const upstreamUrl = new URL(upstream)
-      const upstreamBase = upstreamUrl.pathname && upstreamUrl.pathname !== '/'
-        ? upstreamUrl.href.replace(upstreamUrl.pathname, '')
-        : upstream
+    for (const { prefix, rewritePrefix, upstream, wsUpstream, wsClientOptions } of this.prefixList) {
+      if (wsUpstream) {
+        const target = new URL(wsUpstream)
+        target.search = source.search
+        return { target, wsClientOptions }
+      }
 
       if (source.pathname.startsWith(prefix)) {
-        const target = new URL(source.pathname.replace(prefix, rewritePrefix), upstreamBase)
+        const target = new URL(source.pathname.replace(prefix, rewritePrefix), upstream)
         target.search = source.search
         return { target, wsClientOptions }
       }
@@ -212,7 +210,8 @@ function setupWebSocketProxy (fastify, options, rewritePrefix) {
     wsProxy.addUpstream(
       fastify.prefix,
       rewritePrefix,
-      options.wsUpstream ? options.wsUpstream : options.upstream,
+      options.upstream,
+      options.wsUpstream,
       options.wsClientOptions
     )
     // The else block is validate earlier in the code

test/websocket.js

@@ -421,27 +421,29 @@ test('Proxy websocket with custom upstream url', async (t) => {
   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.url])
+      serverMessages.push([message.toString(), binary, request.headers.host.split(':')[0], request.url])
       ws.send(message, { binary })
     })
   })
 
   await promisify(origin.listen.bind(origin))({ port: 0 })
+  // Host for wsUpstream and for later check.
+  const host = '127.0.0.1'
   // Path for wsUpstream and for later check.
   const path = '/some/path'
   const server = Fastify()
   server.register(proxy, {
     upstream: `ws://localhost:${origin.address().port}`,
     // Start proxy with different upstream, added path.
-    wsUpstream: `ws://localhost:${origin.address().port}${path}`,
+    wsUpstream: `ws://${host}:${origin.address().port}${path}`,
     websocket: true
   })
 
   await server.listen({ port: 0 })
   t.teardown(server.close.bind(server))
 
   // Start websocket with different upstream for connect, added path.
-  const ws = new WebSocket(`ws://localhost:${server.server.address().port}${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 }]
@@ -463,8 +465,8 @@ test('Proxy websocket with custom upstream url', async (t) => {
   }
   // Also check "path", must be the same.
   t.strictSame(serverMessages, [
-    ['hello', false, path],
-    ['fastify', true, path]
+    ['hello', false, host, path],
+    ['fastify', true, host, path]
   ])
 
   await Promise.all([