feat(node-ws): Environment variable support

.changeset/eleven-corners-smash.md

@@ -0,0 +1,5 @@
+---
+'@hono/node-ws': minor
+---
+
+load custom env when init

packages/node-ws/src/index.test.ts

@@ -2,12 +2,20 @@
 // eslint-disable-next-line @typescript-eslint/ban-ts-comment
 // @ts-ignore
 import type { ServerType } from '@hono/node-server/dist/types'
-import { Hono } from 'hono'
+import { Env, Hono, Context } from 'hono'
 import { cors } from 'hono/cors'
 import { HTTPException } from 'hono/http-exception'
 import type { WSMessageReceive } from 'hono/ws'
 import { WebSocket } from 'ws'
 import { createNodeWebSocket } from '.'
+import type { IncomingMessage } from 'node:http'
+import type { Http2ServerRequest } from 'node:http2'
+
+interface CustomEnv extends Env {
+  Bindings: {
+    customVar: string
+  }
+}
 
 describe('WebSocket helper', () => {
   let app: Hono
@@ -18,15 +26,18 @@
 
   beforeEach(async () => {
     app = new Hono()
-    ;({ injectWebSocket, upgradeWebSocket, wss } = createNodeWebSocket({ app }))
+    const getEnv = (req: IncomingMessage | Http2ServerRequest) => {
+      return { customVar: req.headers.upgrade }
+    }
+    ;({ injectWebSocket, upgradeWebSocket, wss } = createNodeWebSocket({ app, getEnv }))
 
     server = await new Promise<ServerType>((resolve) => {
       const server = serve({ fetch: app.fetch, port: 3030 }, () => {
         resolve(server)
       })
     })
     injectWebSocket(server)
   })
 
   afterEach(() => {
     server.close()
@@ -329,4 +340,17 @@
       })
     )
   })
+
+  it('Should server can obtain environment variables', async () => {
+    const mainPromise = new Promise<string>((resolve) =>
+      app.get('/', (c: Context<CustomEnv>, next) => {
+        let ws = upgradeWebSocket((_c) => ({}))
+        resolve(c.env.customVar)
+        return ws(c, next)
+      })
+    )
+
+    new WebSocket('ws://localhost:3030/')
+    expect(await mainPromise).toBe('websocket')
+  })
 })

packages/node-ws/src/index.ts

@@ -1,11 +1,11 @@
-import type { Hono } from 'hono'
+import type { Hono, Env } from 'hono'
 import { defineWebSocketHelper } from 'hono/ws'
 import type { UpgradeWebSocket, WSContext } from 'hono/ws'
 import type { WebSocket } from 'ws'
 import { WebSocketServer } from 'ws'
 import { STATUS_CODES } from 'node:http'
 import type { IncomingMessage, Server } from 'node:http'
-import type { Http2SecureServer, Http2Server } from 'node:http2'
+import type { Http2SecureServer, Http2Server, Http2ServerRequest } from 'node:http2'
 import type { Duplex } from 'node:stream'
 import { CloseEvent } from './events'
 
@@ -23,6 +23,7 @@
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   app: Hono<any, any, any>
   baseUrl?: string | URL
+  getEnv?: (request: IncomingMessage | Http2ServerRequest) => Env['Bindings']
 }
 
 const generateConnectionSymbol = () => Symbol('connection')
@@ -60,26 +61,28 @@
     wss,
     injectWebSocket(server) {
       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) {
           const value = request.headers[key]
           if (!value) {
             continue
           }
           headers.append(key, Array.isArray(value) ? value[0] : value)
         }
 
+        const customEnv = init.getEnv?.(request) ?? {}
         const env: {
           incoming: IncomingMessage
           outgoing: undefined
           [CONNECTION_SYMBOL_KEY]?: symbol
         } = {
+          ...customEnv,
           incoming: request,
           outgoing: undefined,
         }
         const response = await init.app.request(url, { headers: headers }, env)
         const waiter = waiterMap.get(request)
 
         if (!waiter || waiter.connectionSymbol !== env[CONNECTION_SYMBOL_KEY]) {
           socket.end(
@@ -88,11 +91,11 @@
               'Content-Length: 0\r\n' +
               '\r\n'
           )
           waiterMap.delete(request)
           return
         }
 
         wss.handleUpgrade(request, socket, head, (ws) => {
           wss.emit('connection', ws, request)
         })
       })
@@ -106,7 +109,7 @@
       const connectionSymbol = generateConnectionSymbol()
       c.env[CONNECTION_SYMBOL_KEY] = connectionSymbol
       ;(async () => {
         const ws = await nodeUpgradeWebSocket(c.env.incoming, connectionSymbol)
 
         // buffer messages to handle messages received before the events are set up
         const messagesReceivedInStarting: [data: WebSocket.RawData, isBinary: boolean][] = []