fix: use server upgrade event to access head buffer (#12)

Change API from middleware pattern to setup function to fix WebSocket
connections timing out on HTTPS servers. The middleware pattern couldn't
access the upgrade head buffer, causing Node.js to not properly recognize
WebSocket upgrades.

BREAKING CHANGE: API changed from `app.use(tinyws())` to `tinyws(app, server)`

- Add paths option to restrict WebSocket handling to specific routes
- Export TinyWSOptions interface for TypeScript users
- Update examples and documentation

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: v1rtl

README.md

@@ -41,9 +41,7 @@ pnpm i ws tinyws
 import { App, Request } from '@tinyhttp/app'
 import { tinyws, TinyWSRequest } from 'tinyws'
 
-const app = new App<any, Request & TinyWSRequest>()
-
-app.use(tinyws())
+const app = new App<Request & TinyWSRequest>()
 
 app.use('/ws', async (req, res) => {
   if (req.ws) {
@@ -55,7 +53,20 @@ app.use('/ws', async (req, res) => {
   }
 })
 
-app.listen(3000)
+const server = app.listen(3000)
+tinyws(app, server)
+```
+
+### Restricting WebSocket to specific paths
+
+You can restrict WebSocket handling to specific paths using the `paths` option:
+
+```ts
+// Single path
+tinyws(app, server, { paths: '/ws' })
+
+// Multiple paths
+tinyws(app, server, { paths: ['/ws', '/socket'] })
 ```
 
 See [examples](examples) for express and polka integration.

bun.lockb

@@ -4,8 +4,6 @@ import { type TinyWSRequest, tinyws } from '../src/index'
 
 const app = new App<Request & TinyWSRequest>()
 
-app.use(tinyws())
-
 app.use('/hmr', async (req, res) => {
   if (req.ws) {
     const ws = await req.ws()
@@ -15,4 +13,5 @@ app.use('/hmr', async (req, res) => {
   res.send('Hello from HTTP!')
 })
 
-app.listen(3000)
+const server = app.listen(3000)
+tinyws(app, server)

examples/basic.ts

@@ -12,7 +12,7 @@ declare global {
 
 const app = express()
 
-app.use('/hmr', tinyws(), async (req, res) => {
+app.use('/hmr', async (req, res) => {
   if (req.ws) {
     const ws = await req.ws()
 
@@ -21,4 +21,5 @@ app.use('/hmr', tinyws(), async (req, res) => {
   res.send('Hello from HTTP!')
 })
 
-app.listen(3000)
+const server = app.listen(3000)
+tinyws({ handler: app }, server)

examples/express.ts

@@ -4,8 +4,6 @@ import { type TinyWSRequest, tinyws } from '../src/index'
 
 const app = polka<polka.Request & TinyWSRequest>()
 
-app.use(tinyws())
-
 app.use('/hmr', async (req, res) => {
   if (req.ws) {
     const ws = await req.ws()
@@ -15,4 +13,5 @@ app.use('/hmr', async (req, res) => {
   res.end('Hello from HTTP!')
 })
 
-app.listen(3000)
+const server = app.listen(3000)
+tinyws({ handler: app.handler }, server)

examples/polka.ts

@@ -1,30 +1,65 @@
-import type * as http from 'node:http'
+import * as http from 'node:http'
+import type { Socket } from 'node:net'
 import type { ServerOptions, WebSocket } from 'ws'
-import { WebSocketServer as Server } from 'ws'
+import { WebSocketServer } from 'ws'
 
 export interface TinyWSRequest extends http.IncomingMessage {
   ws: () => Promise<WebSocket>
 }
 
+export interface TinyWSOptions extends ServerOptions {
+  paths?: string | string[]
+}
+
 /**
  * tinyws - adds `req.ws` method that resolves when websocket request appears
- * @param wsOptions
+ * @param app - The application instance with a handler function
+ * @param server - The HTTP server instance
+ * @param options - Optional WebSocket server options and paths to restrict WebSocket handling
+ * @param wss - Optional existing WebSocketServer instance
+ * @returns The WebSocketServer instance
  */
-export const tinyws =
-  (wsOptions?: ServerOptions, wss: Server = new Server({ ...wsOptions, noServer: true })) =>
-  async (req: TinyWSRequest, _: unknown, next: () => void | Promise<void>) => {
-    const upgradeHeader = (req.headers.upgrade || '').split(',').map((s) => s.trim())
-
-    // When upgrade header contains "websocket" it's index is 0
-    if (upgradeHeader.indexOf('websocket') === 0) {
-      req.ws = () =>
+export const tinyws = (
+  app: { handler: (req: any, res: any) => void },
+  server: http.Server,
+  options?: TinyWSOptions,
+  wss: WebSocketServer = new WebSocketServer({ ...options, noServer: true })
+) => {
+  const { paths, ...wsOptions } = options || {}
+  const allowedPaths = paths ? (Array.isArray(paths) ? paths : [paths]) : null
+
+  const upgradeHandler = (request: http.IncomingMessage, socket: Socket, head: Buffer) => {
+    const response = new http.ServerResponse(request)
+    response.assignSocket(socket)
+
+    // Copy the head buffer to avoid keeping the entire slab buffer alive
+    const copyOfHead = Buffer.alloc(head.length)
+    head.copy(copyOfHead)
+
+    response.on('finish', () => {
+      if (response.socket !== null) {
+        response.socket.destroy()
+      }
+    })
+
+    const upgradeHeader = (request.headers.upgrade || '').split(',').map((s) => s.trim())
+    const requestPath = request.url?.split('?')[0] || '/'
+
+    const pathMatches = allowedPaths === null || allowedPaths.some((p) => requestPath.startsWith(p))
+
+    if (upgradeHeader.indexOf('websocket') === 0 && pathMatches) {
+      ;(request as TinyWSRequest).ws = () =>
         new Promise((resolve) => {
-          wss.handleUpgrade(req, req.socket, Buffer.alloc(0), (ws) => {
-            wss.emit('connection', ws, req)
+          wss.handleUpgrade(request, socket, copyOfHead, (ws) => {
+            wss.emit('connection', ws, request)
             resolve(ws)
           })
         })
     }
 
-    await next()
+    app.handler(request, response)
   }
+
+  server.on('upgrade', upgradeHandler)
+  return wss
+}

src/index.ts

@@ -2,30 +2,30 @@ import { it } from 'bun:test'
 import * as assert from 'node:assert'
 import { once } from 'node:events'
 import { App, type Request } from '@tinyhttp/app'
-import { type Server, type ServerOptions, WebSocketServer } from 'ws'
+import { type Server, type ServerOptions, WebSocket, WebSocketServer } from 'ws'
 import { type TinyWSRequest, tinyws } from '../src/index'
 
 const s = (handler: (req: TinyWSRequest) => void, opts?: ServerOptions, inst?: Server) => {
   const app = new App<Request & TinyWSRequest>()
 
-  app.use(tinyws(opts, inst))
   app.use('/ws', async (req) => {
     if (typeof req.ws !== 'undefined') {
       handler(req)
     }
   })
 
-  return app
+  return { app, opts, inst }
 }
 
 it('should respond with a message', async () => {
-  const app = s(async (req) => {
+  const { app } = s(async (req) => {
     const ws = await req?.ws()
 
     return ws.send('hello there')
   })
 
   const server = app.listen(4443, undefined, 'localhost')
+  tinyws(app, server)
 
   const ws = new WebSocket('ws://localhost:4443/ws')
 
@@ -37,15 +37,16 @@ it('should respond with a message', async () => {
 })
 
 it('should resolve a `.ws` property', async () => {
-  const app = s(async (req) => {
+  const { app } = s(async (req) => {
     const ws = await req.ws()
 
-    assert.ok(ws instanceof WebSocket)
+    assert.ok(typeof ws.send === 'function')
 
     return ws.send('hello there')
   })
 
   const server = app.listen(4444, undefined, 'localhost')
+  tinyws(app, server)
 
   const ws = new WebSocket('ws://localhost:4444/ws')
 
@@ -56,11 +57,11 @@ it('should resolve a `.ws` property', async () => {
 })
 
 it('should pass ws options', async () => {
-  const app = s(
+  const { app, opts } = s(
     async (req) => {
       const ws = await req.ws()
 
-      assert.ok(ws instanceof WebSocket)
+      assert.ok(typeof ws.send === 'function')
 
       ws.on('error', (err) => {
         assert.match(err.message, /Max payload size exceeded/)
@@ -74,6 +75,7 @@ it('should pass ws options', async () => {
   )
 
   const server = app.listen(4445, undefined, 'localhost')
+  tinyws(app, server, opts)
 
   const ws = new WebSocket('ws://localhost:4445/ws')
 
@@ -86,15 +88,16 @@ it('should pass ws options', async () => {
 })
 
 it('should accept messages', async () => {
-  const app = s(async (req) => {
+  const { app } = s(async (req) => {
     const ws = await req.ws()
 
-    assert.ok(ws instanceof WebSocket)
+    assert.ok(typeof ws.send === 'function')
 
     return ws.on('message', (msg) => ws.send(`You sent: ${msg}`))
   })
 
   const server = app.listen(4446, undefined, 'localhost')
+  tinyws(app, server)
 
   const ws = new WebSocket('ws://localhost:4446/ws')
 
@@ -114,14 +117,14 @@ it('supports passing a server instance', async () => {
   const wss = new WebSocketServer({ noServer: true })
 
   wss.on('connection', (socket) => {
-    assert.ok(socket instanceof WebSocket)
+    assert.ok(typeof socket.send === 'function')
   })
 
-  const app = s(
+  const { app, inst } = s(
     async (req) => {
       const ws = await req.ws()
 
-      assert.ok(ws instanceof WebSocket)
+      assert.ok(typeof ws.send === 'function')
 
       return ws.send('hello there')
     },
@@ -130,6 +133,7 @@ it('supports passing a server instance', async () => {
   )
 
   const server = app.listen(4447, undefined, 'localhost')
+  tinyws(app, server, {}, inst)
 
   const ws = new WebSocket('ws://localhost:4447/ws')
 
@@ -139,6 +143,80 @@ it('supports passing a server instance', async () => {
   ws.close()
 })
 
-it('returns a middleware function', () => {
-  assert.ok(typeof tinyws() === 'function')
+it('returns a WebSocketServer instance', () => {
+  const app = new App()
+  const server = app.listen(4448, undefined, 'localhost')
+  const wss = tinyws(app, server)
+  assert.ok(wss instanceof WebSocketServer)
+  server.close()
+})
+
+it('restricts WebSocket to specified paths', async () => {
+  const app = new App<Request & TinyWSRequest>()
+
+  app.use('/ws', async (req, res) => {
+    if (req.ws) {
+      const ws = await req.ws()
+      return ws.send('allowed')
+    }
+    res.send('no ws')
+  })
+
+  app.use('/other', async (req, res) => {
+    if (req.ws) {
+      const ws = await req.ws()
+      return ws.send('should not happen')
+    }
+    res.send('no ws on other')
+  })
+
+  const server = app.listen(4449, undefined, 'localhost')
+  tinyws(app, server, { paths: '/ws' })
+
+  // Connection to /ws should work
+  const ws1 = new WebSocket('ws://localhost:4449/ws')
+  const [data] = await once(ws1, 'message')
+  assert.equal(data.toString(), 'allowed')
+  ws1.close()
+
+  // Connection to /other should not have req.ws
+  const ws2 = new WebSocket('ws://localhost:4449/other')
+  const [err] = await once(ws2, 'error')
+  assert.ok(err)
+  ws2.close()
+
+  server.close()
+})
+
+it('supports multiple paths', async () => {
+  const app = new App<Request & TinyWSRequest>()
+
+  app.use('/ws1', async (req) => {
+    if (req.ws) {
+      const ws = await req.ws()
+      return ws.send('ws1')
+    }
+  })
+
+  app.use('/ws2', async (req) => {
+    if (req.ws) {
+      const ws = await req.ws()
+      return ws.send('ws2')
+    }
+  })
+
+  const server = app.listen(4450, undefined, 'localhost')
+  tinyws(app, server, { paths: ['/ws1', '/ws2'] })
+
+  const ws1 = new WebSocket('ws://localhost:4450/ws1')
+  const [data1] = await once(ws1, 'message')
+  assert.equal(data1.toString(), 'ws1')
+  ws1.close()
+
+  const ws2 = new WebSocket('ws://localhost:4450/ws2')
+  const [data2] = await once(ws2, 'message')
+  assert.equal(data2.toString(), 'ws2')
+  ws2.close()
+
+  server.close()
 })