feat(tunnel): bridge public WebSockets over the tunnel

tunnel/e2e/index.ts

@@ -2,6 +2,7 @@ import { Logger } from '@bpinternal/log4bot'
 import yargs, { YargsArgv, YargsSchema } from '@bpinternal/yargs-extra'
 import * as browser from './browser'
 import * as nodejs from './nodejs'
+import * as websocket from './websocket'
 import { sleep } from './utils'
 
 const tests = [
@@ -19,6 +20,21 @@ const tests = [
     name: 'nodejs-invalid-request',
     port: 9082,
     test: nodejs.testInvalidRequest
+  },
+  {
+    name: 'websocket-bridge',
+    port: 9083,
+    test: websocket.testWebSocketBridge
+  },
+  {
+    name: 'websocket-unsupported',
+    port: 9084,
+    test: websocket.testWebSocketUnsupported
+  },
+  {
+    name: 'websocket-pending-cap',
+    port: 9085,
+    test: websocket.testWebSocketPendingCap
   }
 ]
 

tunnel/e2e/websocket.ts

@@ -0,0 +1,114 @@
+import { Logger } from '@bpinternal/log4bot'
+import WebSocket from 'isomorphic-ws'
+import { TunnelTail, TunnelServer } from '../src'
+import { CLOSE_CODES } from '../src/errors'
+import { expect } from './utils'
+
+const TUNNEL_ID = 'ws-tunnel-id'
+
+/**
+ * A visitor WebSocket dialed against `/:tunnelId/:path` bridges through the
+ * tunnel: the tail sees ws_open with the visitor's path/query, frames relay
+ * in both directions, and a visitor close reaches the tail.
+ */
+export const testWebSocketBridge = async (port: number, logger: Logger) => {
+  const server = await TunnelServer.new({ port })
+  const tunnelTail = await TunnelTail.new(`ws://localhost:${port}`, TUNNEL_ID)
+
+  // The tail acts as the local endpoint: accept every ws_open, echo every
+  // frame back with a prefix, and record the close.
+  const opened = new Promise<{ path: string; query?: string }>((resolve) => {
+    tunnelTail.events.on('ws_open', (msg) => {
+      logger.debug(`tail ws_open: ${JSON.stringify(msg)}`)
+      tunnelTail.acceptWebSocket(msg.id)
+      resolve({ path: msg.path, ...(msg.query !== undefined && { query: msg.query }) })
+    })
+  })
+  tunnelTail.events.on('ws_frame', (msg) => {
+    logger.debug(`tail ws_frame: ${JSON.stringify(msg)}`)
+    tunnelTail.sendWebSocketFrame(msg.id, `echo:${msg.data}`, msg.binary)
+  })
+  const tailSawClose = new Promise<{ code?: number }>((resolve) => {
+    tunnelTail.events.on('ws_close', (msg) => resolve({ code: msg.code }))
+  })
+
+  // The tail must advertise the ws capability before the visitor connects
+  // (its hello is sent on open; wait for the head to have processed it).
+  const head = server.getTunnel(TUNNEL_ID)
+  if (!head) throw new Error(`Tunnel ${TUNNEL_ID} not found`)
+  for (let i = 0; i < 50 && !head.supportsWebSockets; i++) {
+    await new Promise((r) => setTimeout(r, 20))
+  }
+  expect(String(head.supportsWebSockets)).toBe('true')
+
+  const visitor = new WebSocket(`ws://localhost:${port}/${TUNNEL_ID}/edge/bots/b1/conversations/c1/ws?token=t1`)
+  const received: string[] = []
+  const gotEcho = new Promise<string>((resolve) => {
+    visitor.addEventListener('message', (ev: WebSocket.MessageEvent) => {
+      received.push(ev.data.toString())
+      resolve(ev.data.toString())
+    })
+  })
+  // Send before the tail accepts on purpose — the server must buffer it.
+  visitor.addEventListener('open', () => visitor.send('hello-through-tunnel'))
+
+  const { path, query } = await opened
+  expect(path).toBe('/edge/bots/b1/conversations/c1/ws')
+  expect(query ?? '').toBe('token=t1')
+
+  const echoed = await gotEcho
+  expect(echoed).toBe('echo:hello-through-tunnel')
+
+  visitor.close(1000, 'done')
+  await tailSawClose
+
+  tunnelTail.close()
+  server.close()
+}
+
+/** Pre-accept frames are hard-capped: a visitor flooding before the tail accepts is closed with 1009. */
+export const testWebSocketPendingCap = async (port: number, _logger: Logger) => {
+  const server = await TunnelServer.new({ port })
+  const tunnelTail = await TunnelTail.new(`ws://localhost:${port}`, TUNNEL_ID)
+  // Never accept — every visitor frame lands in the pre-accept buffer.
+  tunnelTail.events.on('ws_open', () => {})
+
+  const head = server.getTunnel(TUNNEL_ID)
+  if (!head) throw new Error(`Tunnel ${TUNNEL_ID} not found`)
+  for (let i = 0; i < 50 && !head.supportsWebSockets; i++) {
+    await new Promise((r) => setTimeout(r, 20))
+  }
+
+  const visitor = new WebSocket(`ws://localhost:${port}/${TUNNEL_ID}/flood`)
+  const closed = new Promise<WebSocket.CloseEvent>((resolve) => visitor.addEventListener('close', resolve))
+  visitor.addEventListener('open', () => {
+    for (let i = 0; i < 100; i++) visitor.send(`frame-${i}`)
+  })
+  const { code } = await closed
+  expect(String(code)).toBe('1009')
+
+  tunnelTail.close()
+  server.close()
+}
+
+/** A visitor upgrading against a tail that never advertised `ws` is refused cleanly. */
+export const testWebSocketUnsupported = async (port: number, _logger: Logger) => {
+  const server = await TunnelServer.new({ port })
+  const tunnelTail = await TunnelTail.new(`ws://localhost:${port}`, TUNNEL_ID)
+
+  const head = server.getTunnel(TUNNEL_ID)
+  if (!head) throw new Error(`Tunnel ${TUNNEL_ID} not found`)
+  // Simulate a pre-websocket tail: let the on-open hello land, then wipe the
+  // advertised capability (wiping first would race the hello repopulating it).
+  for (let i = 0; i < 50 && !head.supportsWebSockets; i++) {
+    await new Promise((r) => setTimeout(r, 20))
+  }
+  ;(head as unknown as { _capabilities: Set<string> })._capabilities = new Set()
+
+  const visitor = new WebSocket(`ws://localhost:${port}/${TUNNEL_ID}/some/path`)
+  const closed = await new Promise<WebSocket.CloseEvent>((resolve) => visitor.addEventListener('close', resolve))
+  expect(String(closed.code)).toBe(String(CLOSE_CODES.WS_UNSUPPORTED))
+
+  tunnelTail.close()
+  server.close()
+}

tunnel/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@bpinternal/tunnel",
-  "version": "0.1.25",
+  "version": "0.2.0",
   "description": "Tunneling logic for client and server",
   "main": "./dist/index.cjs",
   "browser": "./dist/index.mjs",
@@ -40,4 +40,4 @@
     "pnpm": "8.6.2"
   },
   "packageManager": "pnpm@8.6.2"
-}
+}

tunnel/src/errors.ts

@@ -6,6 +6,7 @@ export const CLOSE_CODES = {
   INVALID_REQUEST_PAYLOAD: 4003,
   INTERNAL_TAIL_ERROR: 4004,
   INTERNAL_HEAD_ERROR: 4005,
+  WS_UNSUPPORTED: 4006,
 } as const
 
 export class TunnelError extends Error {

tunnel/src/rooting.ts

@@ -1,4 +1,5 @@
 const URL_REGEX = /^\/([\w|-]+)$/ // /:tunnelId
+const PUBLIC_URL_REGEX = /^\/([\w|-]+)(\/[^?]*)(?:\?(.*))?$/ // /:tunnelId/:path[?query]
 
 export const formatUrl = (host: string, tunnelId: string): string => {
   return `${host}/${tunnelId}`
@@ -27,3 +28,36 @@ export const parseUrl = (url: string | undefined): ParseUrlResult => {
   const tunnelId = match[1] as string
   return { status: 'success', tunnelId }
 }
+
+type ParsePublicUrlResult =
+  | {
+      status: 'error'
+      reason: string
+    }
+  | {
+      status: 'success'
+      tunnelId: string
+      path: string
+      query?: string
+    }
+
+/**
+ * A public visitor URL: the tunnel id followed by the path the visitor
+ * targets on the tail (`/:tunnelId/:path[?query]`). Distinct from the
+ * bare `/:tunnelId` a tail registers with.
+ */
+export const parsePublicUrl = (url: string | undefined): ParsePublicUrlResult => {
+  if (!url) {
+    return { status: 'error', reason: 'url is empty' }
+  }
+
+  const match = url.match(PUBLIC_URL_REGEX)
+  if (!match) {
+    return { status: 'error', reason: 'invalid url' }
+  }
+
+  const tunnelId = match[1] as string
+  const path = match[2] as string
+  const query = match[3] as string | undefined
+  return { status: 'success', tunnelId, path, ...(query !== undefined && { query }) }
+}