Merge pull request #14 from tinyhttp/fix/websocket-upgrade-head-buffer

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

Author: v1rtl

.github/workflows/main.yml

@@ -1,32 +1,22 @@
-# This is a basic workflow to help you get started with Actions
-
 name: CI
 
-# Controls when the action will run. Triggers the workflow on push or pull request
-# events but only for the master branch
 on:
   push:
-    branches: [master]
   pull_request:
-    branches: [master]
 
-# A workflow run is made up of one or more jobs that can run sequentially or in parallel
 jobs:
-  # This workflow contains a single job called "test"
   test:
-    # The type of runner that the job will run on
     runs-on: ubuntu-latest
-
-    # Steps represent a sequence of tasks that will be executed as part of the job
     steps:
-      # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it
       - uses: actions/checkout@v4
       - uses: oven-sh/setup-bun@v2
-      with:
-        bun-version: latest
+        with:
+          bun-version: latest
       - run: bun i
+      - run: bun run lint
       - run: bun test --coverage
       - name: Coveralls
+        if: github.ref == 'refs/heads/master'
         uses: coverallsapp/github-action@master
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/release.yml

@@ -0,0 +1,32 @@
+name: Release
+
+on:
+  push:
+    tags:
+      - 'v*'
+
+permissions:
+  id-token: write
+  contents: read
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-node@v4
+        with:
+          node-version: '22'
+          registry-url: 'https://registry.npmjs.org'
+
+      - uses: oven-sh/setup-bun@v2
+        with:
+          bun-version: latest
+
+      - run: bun i
+      - run: bun run build
+
+      - run: npm publish --provenance --access public
+        env:
+          NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}

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.

biome.json

@@ -1,12 +1,7 @@
 {
   "$schema": "https://biomejs.dev/schemas/1.8.2/schema.json",
   "files": {
-    "ignore": [
-      "node_modules",
-      "dist",
-      "coverage",
-      ".pnpm-store"
-    ]
+    "ignore": ["node_modules", "dist", "coverage", ".pnpm-store"]
   },
   "formatter": {
     "enabled": true,

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

@@ -2,9 +2,7 @@
   "name": "tinyws",
   "version": "0.1.0",
   "description": "Tiny WebSocket middleware for Node.js based on ws.",
-  "files": [
-    "dist"
-  ],
+  "files": ["dist"],
   "engines": {
     "node": ">=12.4"
   },
@@ -13,51 +11,39 @@
   "types": "dist/index.d.ts",
   "scripts": {
     "build": "tsc -p tsconfig.build.json",
-    "test": "uvu -r tsm tests",
-    "test:coverage": "c8 --include=src pnpm test",
-    "test:report": "c8 report --reporter=text-lcov > coverage.lcov",
-    "lint": "eslint \"./**/*.ts\"",
-    "format": "prettier --write \"./**/*.ts\"",
-    "prepublishOnly": "npm run test && npm run lint && npm run build"
+    "test": "bun test",
+    "test:coverage": "bun test --coverage",
+    "lint": "biome check --write .",
+    "prepublishOnly": "bun test && bun run lint && bun run build"
   },
   "repository": {
     "type": "git",
     "url": "git+https://github.com/talentlessguy/tinyws.git"
   },
-  "keywords": [
-    "ws",
-    "express",
-    "tinyhttp",
-    "websocket",
-    "middleware",
-    "polka",
-    "http",
-    "server"
-  ],
+  "keywords": ["ws", "express", "tinyhttp", "websocket", "middleware", "polka", "http", "server"],
   "author": "v1rtl (https://v1rtl.site)",
   "license": "MIT",
   "bugs": {
     "url": "https://github.com/talentlessguy/tinyws/issues"
   },
   "homepage": "https://github.com/talentlessguy/tinyws#readme",
   "devDependencies": {
-    "@biomejs/biome": "^1.8.2",
-    "@commitlint/cli": "^17.6.5",
-    "@commitlint/config-conventional": "^17.6.5",
-    "@tinyhttp/app": "^2.1.0",
-    "@types/bun": "^1.1.5",
-    "@types/express": "^4.17.17",
-    "@types/node": "^18.16.18",
-    "@types/ws": "^8.5.5",
+    "@biomejs/biome": "^1.9.4",
+    "@commitlint/cli": "^17.8.1",
+    "@commitlint/config-conventional": "^17.8.1",
+    "@tinyhttp/app": "^2.5.2",
+    "@types/bun": "^1.3.8",
+    "@types/express": "^4.17.25",
+    "@types/node": "^18.19.130",
+    "@types/ws": "^8.18.1",
     "c8": "7.12.0",
-    "express": "^4.18.2",
+    "express": "^4.22.1",
     "husky": "^8.0.3",
-    "polka": "^1.0.0-next.25",
+    "polka": "^1.0.0-next.28",
     "typescript": "^4.9.5",
-    "ws": "^8.13.0"
+    "ws": "^8.19.0"
   },
   "peerDependencies": {
     "ws": ">=8"
-  },
-  "packageManager": "pnpm@9.3.0+sha256.e1f9e8d1a16607a46dd3c158b5f7a7dc7945501d1c6222d454d63d033d1d918f"
+  }
 }

package.json

@@ -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
+}