diff --git a/tunnel/e2e/index.ts b/tunnel/e2e/index.ts index 48f45c903..034870c54 100644 --- a/tunnel/e2e/index.ts +++ b/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 } ] diff --git a/tunnel/e2e/websocket.ts b/tunnel/e2e/websocket.ts new file mode 100644 index 000000000..6f42d9c7d --- /dev/null +++ b/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((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((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 })._capabilities = new Set() + + const visitor = new WebSocket(`ws://localhost:${port}/${TUNNEL_ID}/some/path`) + const closed = await new Promise((resolve) => visitor.addEventListener('close', resolve)) + expect(String(closed.code)).toBe(String(CLOSE_CODES.WS_UNSUPPORTED)) + + tunnelTail.close() + server.close() +} diff --git a/tunnel/package.json b/tunnel/package.json index 57e6ef3d1..61250db19 100644 --- a/tunnel/package.json +++ b/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" -} \ No newline at end of file +} diff --git a/tunnel/src/errors.ts b/tunnel/src/errors.ts index e7314e9fd..8fafcd611 100644 --- a/tunnel/src/errors.ts +++ b/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 { diff --git a/tunnel/src/rooting.ts b/tunnel/src/rooting.ts index 618576f3e..bbde3f6ab 100644 --- a/tunnel/src/rooting.ts +++ b/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 }) } +} diff --git a/tunnel/src/tunnel-client.ts b/tunnel/src/tunnel-client.ts index 39f76b4d6..501a4d5da 100644 --- a/tunnel/src/tunnel-client.ts +++ b/tunnel/src/tunnel-client.ts @@ -3,7 +3,19 @@ import WebSocket from 'isomorphic-ws' import * as errors from './errors' import { EventEmitter } from './event-emitter' import * as rooting from './rooting' -import { Hello, TunnelRequest, TunnelResponse, headSchema, tailSchema } from './types' +import { + Hello, + TunnelRequest, + TunnelResponse, + TunnelWsAccept, + TunnelWsClose, + TunnelWsFrame, + TunnelWsOpen, + TunnelWsReject, + WS_CAPABILITY, + headSchema, + tailSchema +} from './types' export type ClientCloseEvent = WebSocket.CloseEvent export type ClientErrorEvent = WebSocket.Event @@ -18,7 +30,12 @@ export abstract class TunnelClient { open: WebSocket.Event request: TunnelRequest response: TunnelResponse - hello: {} + hello: Hello + ws_open: TunnelWsOpen + ws_accept: TunnelWsAccept + ws_reject: TunnelWsReject + ws_frame: TunnelWsFrame + ws_close: TunnelWsClose }>() public get closed() { @@ -61,12 +78,31 @@ export abstract class TunnelClient { this._ws.close(code ?? errors.CLOSE_CODES.NORMAL_CLOSURE, reason) } - public readonly hello = () => { + public readonly hello = (capabilities?: string[]) => { this._throwIfClosed() - const hello: Hello = { type: 'hello' } + const hello: Hello = { type: 'hello', ...(capabilities?.length && { capabilities }) } this._ws.send(JSON.stringify(hello)) } + /** Relay one WebSocket frame of the bridged connection `id`. */ + public readonly sendWebSocketFrame = (id: string, data: string, binary?: boolean) => { + this._throwIfClosed() + const frame: TunnelWsFrame = { type: 'ws_frame', id, data, ...(binary && { binary }) } + this._ws.send(JSON.stringify(frame)) + } + + /** Close the bridged WebSocket connection `id` on the other end. */ + public readonly closeWebSocket = (id: string, code?: number, reason?: string) => { + this._throwIfClosed() + const close: TunnelWsClose = { + type: 'ws_close', + id, + ...(code !== undefined && { code }), + ...(reason !== undefined && { reason }) + } + this._ws.send(JSON.stringify(close)) + } + protected _throwIfClosed = () => { if (this._closed) { throw new Error('tunnel is closed') @@ -88,7 +124,10 @@ export class TunnelTail extends TunnelClient { error: ClientErrorEvent request: TunnelRequest open: ClientOpenEvent - hello: {} + hello: Hello + ws_open: TunnelWsOpen + ws_frame: TunnelWsFrame + ws_close: TunnelWsClose }> = this._ev public static new(host: string, tunnelId: string): Promise { @@ -103,12 +142,31 @@ export class TunnelTail extends TunnelClient { const url = rooting.formatUrl(host, tunnelId) const headers = { 'User-Agent': 'tunnel-client' } // for firewall + // The bundled `ws` client throws "Unexpected server response: 101" under + // Bun; Bun's native WebSocket handles the upgrade correctly and accepts a + // `{ headers }` option, so prefer it there. Keep `ws` under Node. + type AnyWebSocketConstructor = new (url: string, options?: unknown) => WebSocket + const runtime = globalThis as unknown as { Bun?: unknown; WebSocket?: AnyWebSocketConstructor } + const TunnelWS: AnyWebSocketConstructor = + runtime.Bun !== undefined && runtime.WebSocket + ? runtime.WebSocket + : (WebSocket as unknown as AnyWebSocketConstructor) const socket = isBrowser - ? new WebSocket(url) // headers are not supported in browser, but the browser will add the User-Agent header automatically - : new WebSocket(url, { headers }) + ? new TunnelWS(url) // headers are not supported in browser, but the browser will add the User-Agent header automatically + : new TunnelWS(url, { headers }) super(socket) + // Advertise protocol extensions as soon as the tunnel opens — the head + // only bridges visitor WebSockets to tails that declared support. + this._ev.once('open', () => { + try { + this.hello([WS_CAPABILITY]) + } catch { + // The tunnel closed between open and hello — nothing to advertise. + } + }) + this._ev.on('message', (ev: WebSocket.MessageEvent) => { const message = this._parseMessage(ev) if (!message) { @@ -116,10 +174,14 @@ export class TunnelTail extends TunnelClient { return } if (message.type === 'hello') { - this.events.emit('hello', {}) + this.events.emit('hello', message.hello) + return + } + if (message.type === 'request') { + this.events.emit('request', message.request) return } - this.events.emit('request', message.request) + this.events.emit(message.message.type, message.message as never) }) } @@ -130,9 +192,27 @@ export class TunnelTail extends TunnelClient { this._ws.send(JSON.stringify(res)) } + /** Confirm a `ws_open` — frames may flow for this connection from now on. */ + public readonly acceptWebSocket = (id: string) => { + this._throwIfClosed() + const accept: TunnelWsAccept = { type: 'ws_accept', id } + this._ws.send(JSON.stringify(accept)) + } + + /** Refuse a `ws_open` — the head closes the visitor's socket. */ + public readonly rejectWebSocket = (id: string, reason?: string) => { + this._throwIfClosed() + const reject: TunnelWsReject = { type: 'ws_reject', id, ...(reason !== undefined && { reason }) } + this._ws.send(JSON.stringify(reject)) + } + private _parseMessage = ( ev: WebSocket.MessageEvent - ): { type: 'hello' } | { type: 'request'; request: TunnelRequest } | undefined => { + ): + | { type: 'hello'; hello: Hello } + | { type: 'request'; request: TunnelRequest } + | { type: 'ws'; message: TunnelWsOpen | TunnelWsFrame | TunnelWsClose } + | undefined => { const data = JSON.parse(ev.data.toString()) const parseResult = tailSchema.safeParse(data) @@ -141,7 +221,15 @@ export class TunnelTail extends TunnelClient { } if (parseResult.data.type === 'hello') { - return { type: 'hello' } + return { type: 'hello', hello: parseResult.data } + } + + if ( + parseResult.data.type === 'ws_open' || + parseResult.data.type === 'ws_frame' || + parseResult.data.type === 'ws_close' + ) { + return { type: 'ws', message: parseResult.data } } return { type: 'request', request: parseResult.data } @@ -149,12 +237,18 @@ export class TunnelTail extends TunnelClient { } export class TunnelHead extends TunnelClient { + private _capabilities: Set = new Set() + public readonly events: EventEmitter<{ close: ClientCloseEvent error: ClientErrorEvent response: TunnelResponse open: ClientOpenEvent - hello: {} + hello: Hello + ws_accept: TunnelWsAccept + ws_reject: TunnelWsReject + ws_frame: TunnelWsFrame + ws_close: TunnelWsClose }> = this._ev public constructor(public readonly tunnelId: string, ws: WebSocket) { @@ -167,13 +261,25 @@ export class TunnelHead extends TunnelClient { return } if (message.type === 'hello') { - this.events.emit('hello', {}) + for (const capability of message.hello.capabilities ?? []) { + this._capabilities.add(capability) + } + this.events.emit('hello', message.hello) + return + } + if (message.type === 'response') { + this.events.emit('response', message.response) return } - this.events.emit('response', message.response) + this.events.emit(message.message.type, message.message as never) }) } + /** True once the tail advertised WebSocket bridging (hello capabilities). */ + public get supportsWebSockets(): boolean { + return this._capabilities.has(WS_CAPABILITY) + } + public readonly send = (request: Omit) => { this._throwIfClosed() @@ -181,9 +287,20 @@ export class TunnelHead extends TunnelClient { this._ws.send(JSON.stringify(req)) } + /** Ask the tail to open a bridged WebSocket connection; answered by ws_accept / ws_reject. */ + public readonly openWebSocket = (open: Omit) => { + this._throwIfClosed() + const req: TunnelWsOpen = { type: 'ws_open', ...open } + this._ws.send(JSON.stringify(req)) + } + private _parseMessage = ( ev: WebSocket.MessageEvent - ): { type: 'hello' } | { type: 'response'; response: TunnelResponse } | undefined => { + ): + | { type: 'hello'; hello: Hello } + | { type: 'response'; response: TunnelResponse } + | { type: 'ws'; message: TunnelWsAccept | TunnelWsReject | TunnelWsFrame | TunnelWsClose } + | undefined => { const data = JSON.parse(ev.data.toString()) const parseResult = headSchema.safeParse(data) @@ -192,7 +309,16 @@ export class TunnelHead extends TunnelClient { } if (parseResult.data.type === 'hello') { - return { type: 'hello' } + return { type: 'hello', hello: parseResult.data } + } + + if ( + parseResult.data.type === 'ws_accept' || + parseResult.data.type === 'ws_reject' || + parseResult.data.type === 'ws_frame' || + parseResult.data.type === 'ws_close' + ) { + return { type: 'ws', message: parseResult.data } } return { type: 'response', response: parseResult.data } diff --git a/tunnel/src/tunnel-server.ts b/tunnel/src/tunnel-server.ts index cd6c548e1..4c0cc34da 100644 --- a/tunnel/src/tunnel-server.ts +++ b/tunnel/src/tunnel-server.ts @@ -5,6 +5,7 @@ import * as errors from './errors' import { EventEmitter } from './event-emitter' import * as rooting from './rooting' import { TunnelHead } from './tunnel-client' +import { TunnelHeader } from './types' export type TunnelServerProps = { port?: number @@ -17,9 +18,34 @@ export type ServerListeningEvent = {} export type ServerConnectionEvent = TunnelHead export type ServerDisconnectionEvent = { tunnelId: string } +/** How long a visitor socket waits for the tail's ws_accept before giving up. */ +const VISITOR_ACCEPT_TIMEOUT_MS = 10_000 + +/** + * Frames a visitor sends before the tail accepts are buffered (the visitor's + * upgrade already completed — it cannot know the handshake is in flight). + * The visitor is unauthenticated at this point, so the buffer is hard-capped; + * exceeding it closes the socket instead of holding attacker-controlled bytes. + */ +const MAX_PENDING_FRAMES = 64 +const MAX_PENDING_BYTES = 256 * 1024 + +/** Headers a bridged ws_open forwards to the tail (auth cookies, origin checks). */ +const FORWARDED_VISITOR_HEADERS = ['cookie', 'origin', 'user-agent', 'x-forwarded-for'] + +/** Only codes an endpoint may SEND (RFC 6455): reserved ones (1005, 1006, …) become a normal closure. */ +const sendableCloseCode = (code: number | undefined): number => + code !== undefined && ((code >= 1000 && code <= 1003) || (code >= 1007 && code <= 1011) || (code >= 3000 && code <= 4999)) + ? code + : 1000 + +const visitorConnectionId = (): string => + `wsc_${Date.now().toString(36)}${Math.random().toString(36).slice(2, 12)}` + export class TunnelServer { private _wss: WebSocket.WebSocketServer private _tunnels: Record = {} + private _visitors: Record> = {} private _closed = false public readonly events = new EventEmitter<{ @@ -84,14 +110,23 @@ export class TunnelServer { private _handleConnection = (ws: WebSocket, req: IncomingMessage) => { const parseResult = rooting.parseUrl(req.url) - if (parseResult.status === 'error') { - const code = errors.CLOSE_CODES.INVALID_TUNNEL_ID - const { reason } = parseResult - ws.close(code, reason) + if (parseResult.status === 'success') { + this._handleTailConnection(ws, parseResult.tunnelId) return } - const { tunnelId } = parseResult + // Not a tail registration (`/:tunnelId`): a public visitor upgrading a + // WebSocket against a path served by the tail (`/:tunnelId/:path`). + const publicResult = rooting.parsePublicUrl(req.url) + if (publicResult.status === 'success') { + this._handleVisitorConnection(ws, req, publicResult) + return + } + + ws.close(errors.CLOSE_CODES.INVALID_TUNNEL_ID, parseResult.reason) + } + + private _handleTailConnection = (ws: WebSocket, tunnelId: string) => { if (this._tunnels[tunnelId]) { ws.close(errors.CLOSE_CODES.TUNNEL_ID_CONFLICT, 'tunnel ID already in use') return @@ -104,8 +139,148 @@ export class TunnelServer { this._tunnels[tunnelId] = tunnel } + /** + * Bridge a visitor's WebSocket to the tail: ws_open asks the tail to dial + * its local counterpart, then frames relay in both directions until either + * side closes. Frames the visitor sends before the tail accepts are + * buffered — the visitor has no way to know the handshake is still in + * flight (its upgrade already completed). + */ + private _handleVisitorConnection = ( + ws: WebSocket, + req: IncomingMessage, + target: { tunnelId: string; path: string; query?: string } + ) => { + const tunnel = this._tunnels[target.tunnelId] + if (!tunnel) { + ws.close(errors.CLOSE_CODES.INVALID_TUNNEL_ID, 'tunnel not found') + return + } + if (!tunnel.supportsWebSockets) { + ws.close(errors.CLOSE_CODES.WS_UNSUPPORTED, 'tunnel does not support websocket bridging') + return + } + + const id = visitorConnectionId() + const visitors = (this._visitors[target.tunnelId] ??= new Map()) + visitors.set(id, ws) + + const headers: Record = {} + for (const name of FORWARDED_VISITOR_HEADERS) { + const value = req.headers[name] + if (value !== undefined) headers[name] = value + } + + let accepted = false + let finished = false + let pendingBytes = 0 + const pendingFrames: { data: string; binary: boolean }[] = [] + + const teardown = (opts: { visitorCode?: number; visitorReason?: string; notifyTail: boolean }) => { + if (finished) return + finished = true + clearTimeout(acceptTimeout) + pendingFrames.length = 0 + visitors.delete(id) + tunnel.events.off('ws_accept', onAccept) + tunnel.events.off('ws_reject', onReject) + tunnel.events.off('ws_frame', onFrame) + tunnel.events.off('ws_close', onClose) + tunnel.events.off('close', onTunnelClose) + if (opts.notifyTail && !tunnel.closed) { + try { + tunnel.closeWebSocket(id, opts.visitorCode, opts.visitorReason) + } catch { + // tunnel raced shut — nothing to notify + } + } + try { + ws.close(sendableCloseCode(opts.visitorCode), opts.visitorReason) + } catch { + // already closing + } + } + + const acceptTimeout = setTimeout( + // Notify the tail: a late accept would otherwise stream into the void. + () => teardown({ visitorCode: 1011, visitorReason: 'tunnel websocket accept timed out', notifyTail: true }), + VISITOR_ACCEPT_TIMEOUT_MS + ) + + const onAccept = (msg: { id: string }) => { + if (msg.id !== id || finished) return + accepted = true + clearTimeout(acceptTimeout) + for (const frame of pendingFrames.splice(0)) { + tunnel.sendWebSocketFrame(id, frame.data, frame.binary) + } + pendingBytes = 0 + } + const onReject = (msg: { id: string; reason?: string }) => { + if (msg.id !== id) return + teardown({ visitorCode: 1011, visitorReason: msg.reason ?? 'tunnel websocket rejected', notifyTail: false }) + } + const onFrame = (msg: { id: string; data: string; binary?: boolean }) => { + if (msg.id !== id || finished) return + ws.send(msg.binary ? Buffer.from(msg.data, 'base64') : msg.data) + } + const onClose = (msg: { id: string; code?: number; reason?: string }) => { + if (msg.id !== id) return + teardown({ visitorCode: msg.code, visitorReason: msg.reason, notifyTail: false }) + } + const onTunnelClose = () => teardown({ visitorCode: 1001, visitorReason: 'tunnel closed', notifyTail: false }) + + tunnel.events.on('ws_accept', onAccept) + tunnel.events.on('ws_reject', onReject) + tunnel.events.on('ws_frame', onFrame) + tunnel.events.on('ws_close', onClose) + tunnel.events.on('close', onTunnelClose) + + ws.addEventListener('message', (ev: WebSocket.MessageEvent) => { + if (finished) return + const binary = typeof ev.data !== 'string' + const data = binary ? Buffer.from(ev.data as Buffer).toString('base64') : (ev.data as string) + if (!accepted) { + // Unauthenticated bytes: hard-capped, never held indefinitely. + if (pendingFrames.length >= MAX_PENDING_FRAMES || pendingBytes + data.length > MAX_PENDING_BYTES) { + teardown({ visitorCode: 1009, visitorReason: 'too much data before tunnel accept', notifyTail: true }) + return + } + pendingBytes += data.length + pendingFrames.push({ data, binary }) + return + } + tunnel.sendWebSocketFrame(id, data, binary) + }) + ws.addEventListener('close', (ev: WebSocket.CloseEvent) => { + if (finished) return + finished = true + clearTimeout(acceptTimeout) + pendingFrames.length = 0 + visitors.delete(id) + tunnel.events.off('ws_accept', onAccept) + tunnel.events.off('ws_reject', onReject) + tunnel.events.off('ws_frame', onFrame) + tunnel.events.off('ws_close', onClose) + tunnel.events.off('close', onTunnelClose) + if (!tunnel.closed) { + try { + tunnel.closeWebSocket(id, ev.code, ev.reason?.toString()) + } catch { + // tunnel raced shut — nothing to notify + } + } + }) + ws.addEventListener('error', () => { + // close event follows and owns the cleanup + }) + + tunnel.openWebSocket({ id, path: target.path, ...(target.query !== undefined && { query: target.query }), headers }) + } + private _handleDisconnection = (tunnelId: string) => { delete this._tunnels[tunnelId] + delete this._visitors[tunnelId] this.events.emit('disconnection', { tunnelId }) } diff --git a/tunnel/src/types.ts b/tunnel/src/types.ts index 302dab075..a7b3c83b6 100644 --- a/tunnel/src/types.ts +++ b/tunnel/src/types.ts @@ -23,9 +23,77 @@ export const tunnelResponseSchema = z.object({ body: z.string().optional() }) -// dummy data to keep the connection alive +// dummy data to keep the connection alive; `capabilities` advertises optional +// protocol extensions (a head only sends ws_* frames to a tail that said 'ws', +// so a pre-websocket tail never receives frames it cannot parse) export type Hello = z.infer -export const helloSchema = z.object({ type: z.literal('hello') }) +export const helloSchema = z.object({ + type: z.literal('hello'), + capabilities: z.string().array().optional() +}) + +export const WS_CAPABILITY = 'ws' + +/** + * WebSocket bridging: a public visitor's socket terminated at the tunnel + * server is multiplexed over the tunnel connection frame-by-frame, keyed by a + * per-socket `id`. head → tail: ws_open / ws_frame / ws_close. + * tail → head: ws_accept | ws_reject (answer to ws_open) / ws_frame / ws_close. + */ +export type TunnelWsOpen = z.infer +export const tunnelWsOpenSchema = z.object({ + type: z.literal('ws_open'), + id: z.string(), + path: z.string(), + query: z.string().optional(), + headers: z.record(tunnelHeaderSchema).optional() +}) + +// No subprotocol on ws_accept: the visitor's upgrade completes (without one) +// before the tail answers, so a tail-selected subprotocol could never be +// honored — offering the field would only discard it silently. +export type TunnelWsAccept = z.infer +export const tunnelWsAcceptSchema = z.object({ + type: z.literal('ws_accept'), + id: z.string() +}) + +export type TunnelWsReject = z.infer +export const tunnelWsRejectSchema = z.object({ + type: z.literal('ws_reject'), + id: z.string(), + reason: z.string().optional() +}) + +export type TunnelWsFrame = z.infer +export const tunnelWsFrameSchema = z.object({ + type: z.literal('ws_frame'), + id: z.string(), + /** utf-8 text, or base64 when `binary` is true */ + data: z.string(), + binary: z.boolean().optional() +}) + +export type TunnelWsClose = z.infer +export const tunnelWsCloseSchema = z.object({ + type: z.literal('ws_close'), + id: z.string(), + code: z.number().optional(), + reason: z.string().optional() +}) -export const tailSchema = z.union([tunnelRequestSchema, helloSchema]) -export const headSchema = z.union([tunnelResponseSchema, helloSchema]) +export const tailSchema = z.union([ + tunnelRequestSchema, + helloSchema, + tunnelWsOpenSchema, + tunnelWsFrameSchema, + tunnelWsCloseSchema +]) +export const headSchema = z.union([ + tunnelResponseSchema, + helloSchema, + tunnelWsAcceptSchema, + tunnelWsRejectSchema, + tunnelWsFrameSchema, + tunnelWsCloseSchema +])