From ec6ec4ec02f8db46e3151e7334535e562dfc47e3 Mon Sep 17 00:00:00 2001
From: Mike Stop Continues <150434+mikestopcontinues@users.noreply.github.com>
Date: Sun, 26 May 2024 03:36:11 +0100
Subject: [PATCH] fix(node-ws): create only one WebSocketServer instead of per
 websocket request (#539)

---
 .changeset/yellow-lions-breathe.md |  5 +++++
 packages/node-ws/src/index.test.ts |  2 +-
 packages/node-ws/src/index.ts      | 33 ++++++++----------------------
 3 files changed, 15 insertions(+), 25 deletions(-)
 create mode 100644 .changeset/yellow-lions-breathe.md

diff --git a/.changeset/yellow-lions-breathe.md b/.changeset/yellow-lions-breathe.md
new file mode 100644
index 00000000..17a10966
--- /dev/null
+++ b/.changeset/yellow-lions-breathe.md
@@ -0,0 +1,5 @@
+---
+'@hono/node-ws': patch
+---
+
+create only one WebSocketServer instead of per websocket request
diff --git a/packages/node-ws/src/index.test.ts b/packages/node-ws/src/index.test.ts
index 4de803a9..42d951a0 100644
--- a/packages/node-ws/src/index.test.ts
+++ b/packages/node-ws/src/index.test.ts
@@ -32,7 +32,7 @@ describe('WebSocket helper', () => {
       )
     })
     injectWebSocket(server)
-    const ws = new WebSocket('ws://localhost:3030/')
+    new WebSocket('ws://localhost:3030/')
 
     expect(await mainPromise).toBe(true)
   })
diff --git a/packages/node-ws/src/index.ts b/packages/node-ws/src/index.ts
index 5cf94c28..34c4b1f0 100644
--- a/packages/node-ws/src/index.ts
+++ b/packages/node-ws/src/index.ts
@@ -2,7 +2,6 @@ import { Buffer } from 'buffer'
 import type { Server } from 'node:http'
 import type { Http2SecureServer, Http2Server } from 'node:http2'
 import type { Hono } from 'hono'
-import { createMiddleware } from 'hono/factory'
 import type { UpgradeWebSocket, WSContext } from 'hono/ws'
 import { WebSocketServer } from 'ws'
 
@@ -15,27 +14,17 @@ export interface NodeWebSocketInit {
   baseUrl?: string | URL
 }
 
-/**
- * Extended for telling WebSocket
- * @internal
- */
-class WSResponse extends Response {
-  readonly wss: WebSocketServer
-  constructor(wss: WebSocketServer) {
-    super()
-    this.wss = wss
-  }
-}
-
 /**
  * Create WebSockets for Node.js
  * @param init Options
  * @returns NodeWebSocket
  */
 export const createNodeWebSocket = (init: NodeWebSocketInit): NodeWebSocket => {
+  const wss = new WebSocketServer({noServer: true})
+
   return {
     injectWebSocket(server) {
-      ;(server as Server).on('upgrade', async (request, socket, head) => {
+      server.on('upgrade', async (request, socket, head) => {
         const url = new URL(request.url ?? '/', init.baseUrl ?? 'http://localhost')
         const headers = new Headers()
         for (const key in request.headers) {
@@ -45,15 +34,11 @@ export const createNodeWebSocket = (init: NodeWebSocketInit): NodeWebSocket => {
           }
           headers.append(key, Array.isArray(value) ? value[0] : value)
         }
-        const res = (await init.app.request(url, {
+        await init.app.request(url, {
           headers: headers,
-        })) as Response | WSResponse
-        if (!(res instanceof WSResponse)) {
-          socket.destroy()
-          return
-        }
-        res.wss.handleUpgrade(request, socket, head, (ws) => {
-          res.wss.emit('connection', ws, request)
+        })
+        wss.handleUpgrade(request, socket, head, (ws) => {
+          wss.emit('connection', ws, request)
         })
       })
     },
@@ -64,7 +49,6 @@ export const createNodeWebSocket = (init: NodeWebSocketInit): NodeWebSocket => {
           await next()
           return
         }
-        const wss = new WebSocketServer({ noServer: true })
         const events = await createEvents(c)
         wss.on('connection', (ws) => {
           const ctx: WSContext = {
@@ -109,7 +93,8 @@ export const createNodeWebSocket = (init: NodeWebSocketInit): NodeWebSocket => {
             )
           })
         })
-        return new WSResponse(wss)
+
+        return new Response()
       },
   }
 }