Refactor WebSocket proxy to use Hono middleware

src/server/index.tsx

@@ -6,13 +6,26 @@ import syncRouter from './api/sync'
 import systemRouter from './api/system'
 import { Hype } from '@because/hype'
 import { cleanupStalePublishers } from './mdns'
-import { extractSubdomain, proxySubdomain, proxyWebSocket, websocket } from './proxy'
+import { upgradeWebSocket, websocket } from 'hono/bun'
+import { extractSubdomain, proxySubdomain, wsProxyEvents } from './proxy'
 import { Shell } from './shell'
-import type { Server } from 'bun'
-import type { WsData } from './proxy'
 
 const app = new Hype({ layout: false, logging: !!process.env.DEBUG })
 
+// Subdomain proxy — runs before all Hono routes
+app.use('*', async (c, next) => {
+  const subdomain = extractSubdomain(c.req.header('host') ?? '')
+  if (!subdomain) return next()
+
+  if (c.req.header('upgrade')?.toLowerCase() === 'websocket') {
+    const events = wsProxyEvents(subdomain, c.req.raw)
+    if (!events) return c.text(`App "${subdomain}" not found or not running`, 502)
+    return upgradeWebSocket(c, events)
+  }
+
+  return proxySubdomain(subdomain, c.req.raw)
+})
+
 app.route('/api/apps', appsRouter)
 app.route('/api/events', eventsRouter)
 app.route('/api/sync', syncRouter)
@@ -127,15 +140,5 @@ const defaults = app.defaults
 export default {
   ...defaults,
   maxRequestBodySize: 1024 * 1024 * 50, // 50MB
-  fetch(req: Request, server: Server<WsData>) {
-    const subdomain = extractSubdomain(req.headers.get('host') ?? '')
-    if (subdomain) {
-      if (req.headers.get('upgrade')?.toLowerCase() === 'websocket') {
-        return proxyWebSocket(subdomain, req, server)
-      }
-      return proxySubdomain(subdomain, req)
-    }
-    return defaults.fetch.call(app, req, server)
-  },
   websocket,
 }

src/server/proxy.ts

@@ -1,20 +1,11 @@
-import type { Server, ServerWebSocket } from 'bun'
+import type { WSContext, WSMessageReceive } from 'hono/ws'
 import { getAppBySubdomain } from '$apps'
 import { serveStatic } from '$static'
 
 export const perf = { timing: false }
 
-export type { WsData }
-
-const pendingMessages = new Map<ServerWebSocket<WsData>, (string | ArrayBuffer | Uint8Array)[]>()
-const upstreams = new Map<ServerWebSocket<WsData>, WebSocket>()
-
-interface WsData {
-  port: number
-  path: string
-  protocols: string[]
-  headers: Record<string, string>
-}
+const upstreams = new WeakMap<WSContext, WebSocket>()
+const pendingMessages = new WeakMap<WSContext, (string | ArrayBuffer | Uint8Array)[]>()
 
 export function extractSubdomain(host: string): string | null {
   // Strip port
@@ -87,12 +78,10 @@ export async function proxySubdomain(subdomain: string, req: Request): Promise<R
   }
 }
 
-export function proxyWebSocket(subdomain: string, req: Request, server: Server<WsData>): Response | undefined {
+export function wsProxyEvents(subdomain: string, req: Request) {
   const app = getAppBySubdomain(subdomain)
 
-  if (!app || app.state !== 'running' || !app.port) {
-    return new Response(`App "${subdomain}" not found or not running`, { status: 502 })
-  }
+  if (!app || app.state !== 'running' || !app.port) return null
 
   const url = new URL(req.url)
   const path = url.pathname + url.search
@@ -105,24 +94,18 @@ export function proxyWebSocket(subdomain: string, req: Request, server: Server<W
     const value = req.headers.get(name)
     if (value) forwardHeaders[name] = value
   }
+
   if (!forwardHeaders['x-app-url']) {
     forwardHeaders['x-app-url'] = app.tunnelUrl ?? `${url.protocol}//${subdomain}.${url.hostname}`
   }
 
-  const upgradeHeaders: Record<string, string> = {}
-  if (protocolHeader) upgradeHeaders['sec-websocket-protocol'] = protocolHeader
+  const port = app.port
 
-  const ok = server.upgrade(req, { data: { port: app.port, path, protocols, headers: forwardHeaders } as WsData, headers: upgradeHeaders })
-  if (ok) return undefined
-  return new Response('WebSocket upgrade failed', { status: 500 })
-}
-
-export const websocket = {
-  open(ws: ServerWebSocket<WsData>) {
-    const { port, path } = ws.data
+  return {
+    onOpen(_evt: Event, ws: WSContext) {
       const upstream = new WebSocket(`ws://localhost:${port}${path}`, {
-      headers: { ...ws.data.headers, host: `localhost:${port}` },
-      protocols: ws.data.protocols,
+        headers: { ...forwardHeaders, host: `localhost:${port}` },
+        protocols,
       })
 
       upstream.binaryType = 'arraybuffer'
@@ -146,7 +129,6 @@ export const websocket = {
       })
 
       upstream.addEventListener('message', e => {
-      // binaryType is 'arraybuffer' so data is always string | ArrayBuffer
         ws.send(e.data as string | ArrayBuffer)
       })
 
@@ -165,17 +147,18 @@ export const websocket = {
       })
     },
 
-  message(ws: ServerWebSocket<WsData>, msg: string | ArrayBuffer | Uint8Array) {
+    onMessage(evt: MessageEvent<WSMessageReceive>, ws: WSContext) {
       const upstream = upstreams.get(ws)
       if (!upstream) return
       if (upstream.readyState !== WebSocket.OPEN) {
+        const msg = typeof evt.data === 'string' ? evt.data : evt.data as ArrayBuffer
         pendingMessages.get(ws)?.push(msg)
         return
       }
-    upstream.send(msg)
+      upstream.send(typeof evt.data === 'string' ? evt.data : evt.data as ArrayBuffer)
     },
 
-  close(ws: ServerWebSocket<WsData>) {
+    onClose(_evt: CloseEvent, ws: WSContext) {
       const upstream = upstreams.get(ws)
       if (upstream) {
         upstream.close()
@@ -184,3 +167,4 @@ export const websocket = {
       pendingMessages.delete(ws)
     },
   }
+}