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 upstreams = new WeakMap<WSContext, WebSocket>()
-
+const pendingMessages = new WeakMap<WSContext, (string | ArrayBuffer | Uint8Array)[]>()
-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>
-}
 
 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) {
+  if (!app || app.state !== 'running' || !app.port) return null
-    return new Response(`App "${subdomain}" not found or not running`, { status: 502 })
-  }
 
   const url = new URL(req.url)
   const path = url.pathname + url.search
@@ -105,82 +94,77 @@ 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> = {}
+  const port = app.port
-  if (protocolHeader) upgradeHeaders['sec-websocket-protocol'] = protocolHeader
 
-  const ok = server.upgrade(req, { data: { port: app.port, path, protocols, headers: forwardHeaders } as WsData, headers: upgradeHeaders })
+  return {
-  if (ok) return undefined
+    onOpen(_evt: Event, ws: WSContext) {
-  return new Response('WebSocket upgrade failed', { status: 500 })
+      const upstream = new WebSocket(`ws://localhost:${port}${path}`, {
-}
+        headers: { ...forwardHeaders, host: `localhost:${port}` },
+        protocols,
+      })
 
-export const websocket = {
+      upstream.binaryType = 'arraybuffer'
-  open(ws: ServerWebSocket<WsData>) {
+      upstreams.set(ws, upstream)
-    const { port, path } = ws.data
+      pendingMessages.set(ws, [])
-    const upstream = new WebSocket(`ws://localhost:${port}${path}`, {
-      headers: { ...ws.data.headers, host: `localhost:${port}` },
-      protocols: ws.data.protocols,
-    })
 
-    upstream.binaryType = 'arraybuffer'
+      const timeout = setTimeout(() => {
-    upstreams.set(ws, upstream)
+        if (upstream.readyState !== WebSocket.OPEN) {
-    pendingMessages.set(ws, [])
+          upstream.close()
+          ws.close()
+        }
+      }, 10_000)
 
-    const timeout = setTimeout(() => {
+      upstream.addEventListener('open', () => {
-      if (upstream.readyState !== WebSocket.OPEN) {
+        clearTimeout(timeout)
-        upstream.close()
+        const buffered = pendingMessages.get(ws)
-        ws.close()
+        if (buffered) {
-      }
+          for (const msg of buffered) upstream.send(msg)
-    }, 10_000)
+          pendingMessages.delete(ws)
+        }
+      })
 
-    upstream.addEventListener('open', () => {
+      upstream.addEventListener('message', e => {
-      clearTimeout(timeout)
+        ws.send(e.data as string | ArrayBuffer)
-      const buffered = pendingMessages.get(ws)
+      })
-      if (buffered) {
+
-        for (const msg of buffered) upstream.send(msg)
+      upstream.addEventListener('close', () => {
+        clearTimeout(timeout)
         pendingMessages.delete(ws)
+        upstreams.delete(ws)
+        ws.close()
+      })
+
+      upstream.addEventListener('error', () => {
+        clearTimeout(timeout)
+        pendingMessages.delete(ws)
+        upstreams.delete(ws)
+        ws.close()
+      })
+    },
+
+    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(typeof evt.data === 'string' ? evt.data : evt.data as ArrayBuffer)
+    },
 
-    upstream.addEventListener('message', e => {
+    onClose(_evt: CloseEvent, ws: WSContext) {
-      // binaryType is 'arraybuffer' so data is always string | ArrayBuffer
+      const upstream = upstreams.get(ws)
-      ws.send(e.data as string | ArrayBuffer)
+      if (upstream) {
-    })
+        upstream.close()
-
+        upstreams.delete(ws)
-    upstream.addEventListener('close', () => {
+      }
-      clearTimeout(timeout)
       pendingMessages.delete(ws)
-      upstreams.delete(ws)
+    },
-      ws.close()
+  }
-    })
-
-    upstream.addEventListener('error', () => {
-      clearTimeout(timeout)
-      pendingMessages.delete(ws)
-      upstreams.delete(ws)
-      ws.close()
-    })
-  },
-
-  message(ws: ServerWebSocket<WsData>, msg: string | ArrayBuffer | Uint8Array) {
-    const upstream = upstreams.get(ws)
-    if (!upstream) return
-    if (upstream.readyState !== WebSocket.OPEN) {
-      pendingMessages.get(ws)?.push(msg)
-      return
-    }
-    upstream.send(msg)
-  },
-
-  close(ws: ServerWebSocket<WsData>) {
-    const upstream = upstreams.get(ws)
-    if (upstream) {
-      upstream.close()
-      upstreams.delete(ws)
-    }
-    pendingMessages.delete(ws)
-  },
 }