From 91dfcbc6aafa019011a3e52740ff555c996cc51c Mon Sep 17 00:00:00 2001
From: Valentinas Janeiko <valentinas.janeiko@shares.io>
Date: Mon, 22 Jan 2024 19:51:09 +0000
Subject: [PATCH] Make sure query params are included in WebSocket request

---
 index.js          | 9 +++++++--
 test/websocket.js | 5 +++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/index.js b/index.js
index 2283584..1db3276 100644
--- a/index.js
+++ b/index.js
@@ -150,20 +150,25 @@ class WebSocketProxy {
   }
 
   findUpstream (request, dest) {
+    const search = new URL(request.url, 'ws://127.0.0.1').search
+
     if (typeof this.wsUpstream === 'string' && this.wsUpstream !== '') {
       const target = new URL(this.wsUpstream)
-      target.search = new URL(request.url, 'ws://127.0.0.1').search
+      target.search = search
       return target
     }
 
     if (typeof this.upstream === 'string' && this.upstream !== '') {
-      return new URL(dest, this.upstream)
+      const target = new URL(dest, this.upstream)
+      target.search = search
+      return target
     }
 
     const upstream = this.getUpstream(request, '')
     const target = new URL(dest, upstream)
     /* istanbul ignore next */
     target.protocol = upstream.indexOf('http:') === 0 ? 'ws:' : 'wss'
+    target.search = search
     return target
   }
 
diff --git a/test/websocket.js b/test/websocket.js
index bbb253d..1a43a53 100644
--- a/test/websocket.js
+++ b/test/websocket.js
@@ -522,7 +522,7 @@ test('multiple websocket upstreams with host constraints', async (t) => {
 })
 
 test('multiple websocket upstreams with distinct server options', async (t) => {
-  t.plan(2)
+  t.plan(4)
 
   const server = Fastify()
 
@@ -532,7 +532,8 @@ test('multiple websocket upstreams with distinct server options', async (t) => {
     t.teardown(wss.close.bind(wss))
     t.teardown(origin.close.bind(origin))
 
-    wss.once('connection', (ws) => {
+    wss.once('connection', (ws, req) => {
+      t.equal(req.url, `/?q=${name}`)
       ws.once('message', message => {
         // echo
         ws.send(message)