From 6b9fb24874d04017bc6c449dbc75b330345f43aa Mon Sep 17 00:00:00 2001
From: EdamAmex <121654029+EdamAme-x@users.noreply.github.com>
Date: Sun, 3 Nov 2024 16:14:03 +0900
Subject: [PATCH] fix(helper/websocket): allow custom events to be passed to
 generics (#3616)

* fix(helper/websocket): allow custom events to be passed to generics

* chore: disable any

* add tests

* patch

* test patch
---
 .../cloudflare-workers/websocket.test.ts      |  7 +-
 src/adapter/cloudflare-workers/websocket.ts   | 88 ++++++++++---------
 src/helper/websocket/index.ts                 |  4 +-
 3 files changed, 50 insertions(+), 49 deletions(-)

diff --git a/src/adapter/cloudflare-workers/websocket.test.ts b/src/adapter/cloudflare-workers/websocket.test.ts
index a065d4ab7..e07a612db 100644
--- a/src/adapter/cloudflare-workers/websocket.test.ts
+++ b/src/adapter/cloudflare-workers/websocket.test.ts
@@ -21,14 +21,13 @@ describe('upgradeWebSocket middleware', () => {
     app.get(
       '/ws',
       upgradeWebSocket(() => ({
-        // eslint-disable-next-line @typescript-eslint/no-unused-vars
         onMessage(evt, ws) {
-          resolve(evt.data)
+          resolve([evt.data, ws.readyState || 1])
         },
       }))
     )
   )
-  it('Should receive message is valid', async () => {
+  it('Should receive message and readyState is valid', async () => {
     const sendingData = Math.random().toString()
     await app.request('/ws', {
       headers: {
@@ -41,7 +40,7 @@ describe('upgradeWebSocket middleware', () => {
       })
     )
 
-    expect(sendingData).toBe(await wsPromise)
+    expect([sendingData, 1]).toStrictEqual(await wsPromise)
   })
   it('Should call next() when header does not have upgrade', async () => {
     const next = vi.fn()
diff --git a/src/adapter/cloudflare-workers/websocket.ts b/src/adapter/cloudflare-workers/websocket.ts
index cbf468c70..072a5c487 100644
--- a/src/adapter/cloudflare-workers/websocket.ts
+++ b/src/adapter/cloudflare-workers/websocket.ts
@@ -1,51 +1,53 @@
 import { WSContext, defineWebSocketHelper } from '../../helper/websocket'
-import type { UpgradeWebSocket, WSReadyState } from '../../helper/websocket'
+import type { UpgradeWebSocket, WSEvents, WSReadyState } from '../../helper/websocket'
 
 // Based on https://github.com/honojs/hono/issues/1153#issuecomment-1767321332
-export const upgradeWebSocket: UpgradeWebSocket<WebSocket> = defineWebSocketHelper(
-  async (c, events) => {
-    const upgradeHeader = c.req.header('Upgrade')
-    if (upgradeHeader !== 'websocket') {
-      return
-    }
+export const upgradeWebSocket: UpgradeWebSocket<
+  WebSocket,
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  any,
+  Omit<WSEvents<WebSocket>, 'onOpen'>
+> = defineWebSocketHelper(async (c, events) => {
+  const upgradeHeader = c.req.header('Upgrade')
+  if (upgradeHeader !== 'websocket') {
+    return
+  }
 
-    // @ts-expect-error WebSocketPair is not typed
-    const webSocketPair = new WebSocketPair()
-    const client: WebSocket = webSocketPair[0]
-    const server: WebSocket = webSocketPair[1]
+  // @ts-expect-error WebSocketPair is not typed
+  const webSocketPair = new WebSocketPair()
+  const client: WebSocket = webSocketPair[0]
+  const server: WebSocket = webSocketPair[1]
 
-    const wsContext = new WSContext<WebSocket>({
-      close: (code, reason) => server.close(code, reason),
-      get protocol() {
-        return server.protocol
-      },
-      raw: server,
-      get readyState() {
-        return server.readyState as WSReadyState
-      },
-      url: server.url ? new URL(server.url) : null,
-      send: (source) => server.send(source),
-    })
+  const wsContext = new WSContext<WebSocket>({
+    close: (code, reason) => server.close(code, reason),
+    get protocol() {
+      return server.protocol
+    },
+    raw: server,
+    get readyState() {
+      return server.readyState as WSReadyState
+    },
+    url: server.url ? new URL(server.url) : null,
+    send: (source) => server.send(source),
+  })
 
-    if (events.onOpen) {
-      server.addEventListener('open', (evt: Event) => events.onOpen?.(evt, wsContext))
-    }
-    if (events.onClose) {
-      server.addEventListener('close', (evt: CloseEvent) => events.onClose?.(evt, wsContext))
-    }
-    if (events.onMessage) {
-      server.addEventListener('message', (evt: MessageEvent) => events.onMessage?.(evt, wsContext))
-    }
-    if (events.onError) {
-      server.addEventListener('error', (evt: Event) => events.onError?.(evt, wsContext))
-    }
+  // note: cloudflare workers doesn't support 'open' event
 
-    // @ts-expect-error - server.accept is not typed
-    server.accept?.()
-    return new Response(null, {
-      status: 101,
-      // @ts-expect-error - webSocket is not typed
-      webSocket: client,
-    })
+  if (events.onClose) {
+    server.addEventListener('close', (evt: CloseEvent) => events.onClose?.(evt, wsContext))
+  }
+  if (events.onMessage) {
+    server.addEventListener('message', (evt: MessageEvent) => events.onMessage?.(evt, wsContext))
   }
-)
+  if (events.onError) {
+    server.addEventListener('error', (evt: Event) => events.onError?.(evt, wsContext))
+  }
+
+  // @ts-expect-error - server.accept is not typed
+  server.accept?.()
+  return new Response(null, {
+    status: 101,
+    // @ts-expect-error - webSocket is not typed
+    webSocket: client,
+  })
+})
diff --git a/src/helper/websocket/index.ts b/src/helper/websocket/index.ts
index 9976fe186..0c896bed0 100644
--- a/src/helper/websocket/index.ts
+++ b/src/helper/websocket/index.ts
@@ -20,8 +20,8 @@ export interface WSEvents<T = unknown> {
 /**
  * Upgrade WebSocket Type
  */
-export type UpgradeWebSocket<T = unknown, U = any> = (
-  createEvents: (c: Context) => WSEvents<T> | Promise<WSEvents<T>>,
+export type UpgradeWebSocket<T = unknown, U = any, _WSEvents = WSEvents<T>> = (
+  createEvents: (c: Context) => _WSEvents | Promise<_WSEvents>,
   options?: U
 ) => MiddlewareHandler<
   any,