From 6f90a574c465b4d0ecadbe605bdf434b9f3c95f3 Mon Sep 17 00:00:00 2001
From: Dmytro Kulyk <lnkvisitor.ts@gmail.com>
Date: Sun, 23 Feb 2025 08:31:44 +0200
Subject: [PATCH] feat(node-ws): Reject unexpected WebSocket connections (#973)

* Added rejection of WebSocket connections when the app does not expect them

* added changeset

* Updated waiter names; removed strict option
---
 .changeset/selfish-lies-sing.md    |  5 ++++
 packages/node-ws/src/index.test.ts | 34 ++++++++++++++++++++++++++
 packages/node-ws/src/index.ts      | 38 +++++++++++++++++++++---------
 3 files changed, 66 insertions(+), 11 deletions(-)
 create mode 100644 .changeset/selfish-lies-sing.md

diff --git a/.changeset/selfish-lies-sing.md b/.changeset/selfish-lies-sing.md
new file mode 100644
index 00000000..62a91376
--- /dev/null
+++ b/.changeset/selfish-lies-sing.md
@@ -0,0 +1,5 @@
+---
+'@hono/node-ws': minor
+---
+
+Added rejection of WebSocket connections when the app does not expect them
diff --git a/packages/node-ws/src/index.test.ts b/packages/node-ws/src/index.test.ts
index d7da9eea..771f80d0 100644
--- a/packages/node-ws/src/index.test.ts
+++ b/packages/node-ws/src/index.test.ts
@@ -51,6 +51,40 @@ describe('WebSocket helper', () => {
     expect(await mainPromise).toBe(true)
   })
 
+  it('Should be rejected if upgradeWebSocket is not used', async () => {
+    app.get(
+      '/', (c)=>c.body('')
+    )
+
+    {
+      const ws = new WebSocket('ws://localhost:3030/')
+      const mainPromise = new Promise<boolean>((resolve) => {
+        ws.onerror = () => {
+          resolve(true)
+        }
+        ws.onopen = () => {
+          resolve(false)
+        }
+      })
+
+      expect(await mainPromise).toBe(true)
+    }
+
+    { //also should rejected on fallback
+      const ws = new WebSocket('ws://localhost:3030/notFound')
+      const mainPromise = new Promise<boolean>((resolve) => {
+        ws.onerror = () => {
+          resolve(true)
+        }
+        ws.onopen = () => {
+          resolve(false)
+        }
+      })
+
+      expect(await mainPromise).toBe(true)
+    }
+  })
+
   it('Should be able to connect', async () => {
     const mainPromise = new Promise<boolean>((resolve) =>
       app.get(
diff --git a/packages/node-ws/src/index.ts b/packages/node-ws/src/index.ts
index ef1ed60d..bea96bb9 100644
--- a/packages/node-ws/src/index.ts
+++ b/packages/node-ws/src/index.ts
@@ -5,6 +5,7 @@ import { WebSocketServer } from 'ws'
 import type { IncomingMessage } from 'http'
 import type { Server } from 'node:http'
 import type { Http2SecureServer, Http2Server } from 'node:http2'
+import type { Duplex } from 'node:stream'
 import { CloseEvent } from './events'
 
 export interface NodeWebSocket {
@@ -24,25 +25,25 @@ export interface NodeWebSocketInit {
  */
 export const createNodeWebSocket = (init: NodeWebSocketInit): NodeWebSocket => {
   const wss = new WebSocketServer({ noServer: true })
-  const waiter = new Map<IncomingMessage, (ws: WebSocket) => void>()
+  const waiterMap = new Map<IncomingMessage, { resolve: (ws: WebSocket) => void, response: Response }>()
 
   wss.on('connection', (ws, request) => {
-    const waiterFn = waiter.get(request)
-    if (waiterFn) {
-      waiterFn(ws)
-      waiter.delete(request)
+    const waiter = waiterMap.get(request)
+    if (waiter) {
+      waiter.resolve(ws)
+      waiterMap.delete(request)
     }
   })
 
-  const nodeUpgradeWebSocket = (request: IncomingMessage) => {
+  const nodeUpgradeWebSocket = (request: IncomingMessage, response: Response) => {
     return new Promise<WebSocket>((resolve) => {
-      waiter.set(request, resolve)
+      waiterMap.set(request, { resolve, response })
     })
   }
 
   return {
     injectWebSocket(server) {
-      server.on('upgrade', async (request, socket, head) => {
+      server.on('upgrade', async (request, socket: Duplex, head) => {
         const url = new URL(request.url ?? '/', init.baseUrl ?? 'http://localhost')
         const headers = new Headers()
         for (const key in request.headers) {
@@ -52,11 +53,25 @@ export const createNodeWebSocket = (init: NodeWebSocketInit): NodeWebSocket => {
           }
           headers.append(key, Array.isArray(value) ? value[0] : value)
         }
-        await init.app.request(
+
+        const response = await init.app.request(
           url,
           { headers: headers },
           { incoming: request, outgoing: undefined }
         )
+
+        const waiter = waiterMap.get(request)
+        if (!waiter || waiter.response !== response) {
+          socket.end(
+            'HTTP/1.1 400 Bad Request\r\n' +
+            'Connection: close\r\n' +
+            'Content-Length: 0\r\n' +
+            '\r\n'
+          )
+          waiterMap.delete(request)
+          return
+        }
+
         wss.handleUpgrade(request, socket, head, (ws) => {
           wss.emit('connection', ws, request)
         })
@@ -70,8 +85,9 @@ export const createNodeWebSocket = (init: NodeWebSocketInit): NodeWebSocket => {
           return
         }
 
+        const response = new Response()
         ;(async () => {
-          const ws = await nodeUpgradeWebSocket(c.env.incoming)
+          const ws = await nodeUpgradeWebSocket(c.env.incoming, response)
           const events = await createEvents(c)
 
           const ctx: WSContext<WebSocket> = {
@@ -116,7 +132,7 @@ export const createNodeWebSocket = (init: NodeWebSocketInit): NodeWebSocket => {
           })
         })()
 
-        return new Response()
+        return response
       },
   }
 }