diff --git a/backend/src/index.ts b/backend/src/index.ts index 67a0b450c..cb0675a4b 100644 --- a/backend/src/index.ts +++ b/backend/src/index.ts @@ -341,7 +341,7 @@ app.route('/api/health', createHealthRoutes(db, openCodeSupervisor)) app.route('/api/mcp-oauth-proxy', createMcpOauthProxyRoutes(openCodeClient, requireAuth)) app.route('/api/internal', createInternalRoutes(db, scheduleService, notificationService, settingsService, openCodeClient, sessionPermissionModeService, repoWorkspaces, gitAuthService, sessionGoalService, multiRunService)) -app.route('/api/opencode-proxy', createOpenCodeProxyRoutes(db, settingsService)) +app.route('/api/opencode-proxy', createOpenCodeProxyRoutes(db, settingsService, upgradeWebSocket)) const protectedApi = new Hono() protectedApi.use('/*', requireAuth) diff --git a/backend/src/routes/opencode-proxy.ts b/backend/src/routes/opencode-proxy.ts index 04ad88878..fe2ae637a 100644 --- a/backend/src/routes/opencode-proxy.ts +++ b/backend/src/routes/opencode-proxy.ts @@ -1,5 +1,6 @@ import { Hono } from 'hono' import type { Context } from 'hono' +import type { UpgradeWebSocket } from 'hono/ws' import type { Database } from 'bun:sqlite' import { buildOpenCodeBasicAuth } from '@opencode-manager/shared/opencode' import { createInternalTokenMiddleware } from '../auth/internal-token-middleware' @@ -8,12 +9,22 @@ import { opencodeServerManager } from '../services/opencode-single-server' import { getOpenCodeUpstreamBaseUrl, OPENCODE_DIRECTORY_HEADER, + toWebSocketUrl, withDefaultOpenCodeDirectory, } from '../services/opencode/upstream' import { getRepoById } from '../db/queries' -import { getWorkspacePath } from '@opencode-manager/shared/config/env' +import { getTrustedOrigins, getWorkspacePath } from '@opencode-manager/shared/config/env' import { isPathWithinRoot } from '../services/sandbox/command' import { buildProxyResponseHeaders, filterProxyHeaders } from '../utils/proxy-headers' +import { + connectUpstreamWebSocket, + createBridgedSocketEvents, + isAllowedUpgradeOrigin, + type PendingUpstream, +} from '../utils/websocket-bridge' + +const PROXY_PREFIX = /^\/api\/opencode-proxy(?:\/repos\/[^/]+)?/ +const PERSISTENT_PTY_CONNECT_PATH = /^\/api\/experimental\/persistent-pty\/[^/]+\/connect$/ interface ProxyRequestParts { method: string @@ -32,6 +43,26 @@ function isJsonContentType(contentType: string | undefined): boolean { return (contentType ?? '').toLowerCase().includes('application/json') } +function stripProxyPrefix(pathname: string): string { + return pathname.replace(PROXY_PREFIX, '') || '/' +} + +function isWebSocketUpgrade(c: Context): boolean { + return (c.req.header('connection')?.toLowerCase() ?? '').includes('upgrade') + && c.req.header('upgrade')?.toLowerCase() === 'websocket' +} + +/** + * Matches OpenCode's ticketed persistent PTY socket connects. WebSocket clients cannot + * send an Authorization header, so these rely on the single-use ticket minted through the + * authenticated proxy and validated by the upstream server. Only persistent PTY connects + * are ticketed; OpenCode's /api/pty route builds per-directory services before it can + * validate a ticket, so it must never bypass the internal token. + */ +function isTicketedPersistentPtyConnect(pathSuffix: string, searchParams: URLSearchParams): boolean { + return PERSISTENT_PTY_CONNECT_PATH.test(pathSuffix) && !!searchParams.get('ticket') +} + function decodeDirectoryHeader(value: string | undefined): string | undefined { if (!value) return undefined try { @@ -81,10 +112,37 @@ function rewriteRepoLocationBody(bodyText: string, resolveDirectory: RepoDirecto return JSON.stringify(parsed) } -export function createOpenCodeProxyRoutes(db: Database, settingsService: SettingsService) { +export function createOpenCodeProxyRoutes(db: Database, settingsService: SettingsService, upgradeWebSocket: UpgradeWebSocket) { const app = new Hono() + const requireInternalToken = createInternalTokenMiddleware(db) + + app.use('/*', (c, next) => { + const url = new URL(c.req.url) + const isTicketedSocket = isWebSocketUpgrade(c) + && isTicketedPersistentPtyConnect(stripProxyPrefix(url.pathname), url.searchParams) + return isTicketedSocket ? proxyTicketedPtySocket(c, url) : requireInternalToken(c, next) + }) + + async function proxyTicketedPtySocket(c: Context, url: URL): Promise { + if (!opencodeServerManager.isLifecycleInitialized()) { + return c.json({ error: 'OpenCode lifecycle initialization is incomplete; refusing to proxy to an unmanaged server' }, 503) + } + + if (!isAllowedUpgradeOrigin(c.req.header('origin'), getTrustedOrigins())) { + return c.json({ error: 'Origin not allowed' }, 403) + } + + const pathSuffix = stripProxyPrefix(url.pathname) + let pending: PendingUpstream + try { + pending = await connectUpstreamWebSocket(toWebSocketUrl(`${getOpenCodeUpstreamBaseUrl()}${pathSuffix}${url.search}`)) + } catch { + return c.json({ error: 'PTY connection rejected' }, 403) + } - app.use('/*', createInternalTokenMiddleware(db)) + const upgrade = upgradeWebSocket(() => createBridgedSocketEvents((peer) => pending.attach(peer))) + return (await upgrade(c, async () => {})) ?? c.body(null) + } async function forwardToOpenCode( c: Context, @@ -96,10 +154,8 @@ export function createOpenCodeProxyRoutes(db: Database, settingsService: Setting return c.json({ error: 'OpenCode lifecycle initialization is incomplete; refusing to proxy to an unmanaged server' }, 503) } - const connectionHeader = c.req.header('connection')?.toLowerCase() ?? '' - const upgradeHeader = c.req.header('upgrade')?.toLowerCase() ?? '' - if (connectionHeader.includes('upgrade') && upgradeHeader === 'websocket') { - return c.json({ error: 'WebSocket proxying is not supported' }, 501) + if (isWebSocketUpgrade(c)) { + return c.json({ error: 'WebSocket proxying is only supported for ticketed persistent PTY connects' }, 501) } const url = new URL(c.req.url) @@ -164,7 +220,7 @@ export function createOpenCodeProxyRoutes(db: Database, settingsService: Setting } const url = new URL(c.req.url) - const pathSuffix = url.pathname.replace(/^\/api\/opencode-proxy\/repos\/[^/]+/, '') || '/' + const pathSuffix = stripProxyPrefix(url.pathname) const resolveDirectory = createRepoDirectoryResolver(repo.fullPath) const directory = resolveDirectory( url.searchParams.get('location[directory]') || decodeDirectoryHeader(c.req.header(OPENCODE_DIRECTORY_HEADER)), @@ -180,7 +236,7 @@ export function createOpenCodeProxyRoutes(db: Database, settingsService: Setting app.all('/*', async (c) => { const url = new URL(c.req.url) - const pathSuffix = url.pathname.replace(/^\/api\/opencode-proxy/, '') || '/' + const pathSuffix = stripProxyPrefix(url.pathname) return forwardToOpenCode(c, pathSuffix) }) diff --git a/backend/src/routes/repo-terminal-socket.ts b/backend/src/routes/repo-terminal-socket.ts index c2ec12348..cfb27b66d 100644 --- a/backend/src/routes/repo-terminal-socket.ts +++ b/backend/src/routes/repo-terminal-socket.ts @@ -6,8 +6,7 @@ import { getTrustedOrigins } from '@opencode-manager/shared/config/env' import type { GitAuthService } from '../services/git-auth' import type { OpenCodeClient } from '../services/opencode/client' import type { TerminalService } from '../services/terminal' -import { forwardPeerMessage, peerBufferedAmount, type WebSocketBridge } from '../utils/websocket-bridge' -import { getErrorMessage } from '../utils/error-utils' +import { createBridgedSocketEvents, isAllowedUpgradeOrigin } from '../utils/websocket-bridge' import { resolveRepoRequestDirectory } from './repo-directory' interface RepoTerminalSocketVariables { @@ -15,17 +14,6 @@ interface RepoTerminalSocketVariables { terminalCursor: number | undefined } -const MAX_CLOSE_REASON_BYTES = 120 - -function isAllowedUpgradeOrigin( - origin: string | undefined, - trustedOrigins: readonly string[], -): boolean { - if (origin === undefined) return true - if (origin === 'null') return false - return trustedOrigins.includes(origin) -} - function parseCursor(raw: string | undefined): number | undefined | null { if (raw === undefined) return undefined if (!/^-?\d+$/.test(raw)) return null @@ -34,12 +22,6 @@ function parseCursor(raw: string | undefined): number | undefined | null { return parsed } -function truncateCloseReason(message: string): string { - const bytes = Buffer.from(message, 'utf8') - if (bytes.length <= MAX_CLOSE_REASON_BYTES) return message - return bytes.subarray(0, MAX_CLOSE_REASON_BYTES).toString('utf8') -} - export function createRepoTerminalSocketRoutes( database: Database, gitAuthService: GitAuthService, @@ -79,40 +61,7 @@ export function createRepoTerminalSocketRoutes( const cursor = c.get('terminalCursor') const ptyID = c.req.param('ptyID') - let bridge: WebSocketBridge | undefined - let closed = false - let connectPromise: Promise = Promise.resolve() - - return { - onOpen(_event, ws) { - connectPromise = terminalService - .connect(directory, ptyID, cursor, { - send: (data) => ws.send(data), - close: (code, reason) => ws.close(code, reason), - bufferedAmount: () => peerBufferedAmount(ws), - }) - .then((connected) => { - if (closed) { - connected.close() - return - } - bridge = connected - }) - .catch((error: unknown) => { - if (closed) return - ws.close(1011, truncateCloseReason(getErrorMessage(error))) - }) - }, - async onMessage(event) { - await connectPromise - if (!bridge) return - await forwardPeerMessage(bridge, event.data) - }, - onClose() { - closed = true - bridge?.close() - }, - } + return createBridgedSocketEvents((peer) => terminalService.connect(directory, ptyID, cursor, peer)) }), ) diff --git a/backend/src/services/opencode/upstream.ts b/backend/src/services/opencode/upstream.ts index f5a8a2012..0945b0e77 100644 --- a/backend/src/services/opencode/upstream.ts +++ b/backend/src/services/opencode/upstream.ts @@ -17,3 +17,11 @@ export function withDefaultOpenCodeDirectory(headers: Record): R const hasDirectory = Object.keys(headers).some((key) => key.toLowerCase() === OPENCODE_DIRECTORY_HEADER) return hasDirectory ? headers : { ...headers, [OPENCODE_DIRECTORY_HEADER]: encodeURIComponent(getWorkspacePath()) } } + +/** Maps an http(s) URL to its ws(s) equivalent, leaving ws/wss URLs unchanged. */ +export function toWebSocketUrl(httpUrl: string): string { + const url = new URL(httpUrl) + if (url.protocol === 'https:') url.protocol = 'wss:' + else if (url.protocol === 'http:') url.protocol = 'ws:' + return url.toString() +} diff --git a/backend/src/services/preview/gateway.ts b/backend/src/services/preview/gateway.ts index 20387bc79..6115dd73d 100644 --- a/backend/src/services/preview/gateway.ts +++ b/backend/src/services/preview/gateway.ts @@ -5,12 +5,7 @@ import { createNodeWebSocket, type NodeWebSocket } from '@hono/node-ws' import type { PreviewPort } from '@opencode-manager/shared/schemas' import { AUTH_COOKIE_PREFIX } from '../../auth/cookies' import { buildProxyResponseHeaders, filterProxyHeaders } from '../../utils/proxy-headers' -import { - bridgeWebSocket, - forwardPeerMessage, - peerBufferedAmount, - type WebSocketBridge, -} from '../../utils/websocket-bridge' +import { bridgeWebSocket, createBridgedSocketEvents } from '../../utils/websocket-bridge' export const PREVIEW_COOKIE = 'ocm_preview' @@ -229,39 +224,20 @@ export function createPreviewGatewayApp(store: PreviewSessionStore, fetchFn: typ upgradeWebSocket((c) => { const requestUrl = new URL(c.req.url) const protocols = parseProtocols(c.req.header('sec-websocket-protocol')) - let bridge: WebSocketBridge | undefined - - return { - onOpen(_event, ws) { - const sessionId = readPreviewCookie(c.req.header('cookie')) - const session = sessionId ? store.getSession(sessionId) : undefined - if (!session) { - ws.close(1011) - return - } - - let upstream: WebSocket - try { - upstream = new WebSocket(upstreamWebSocketUrl(session, requestUrl), protocols.length > 0 ? protocols : undefined) - } catch { - ws.close(1011) - return - } - - bridge = bridgeWebSocket(upstream, { - send: (data) => ws.send(data), - close: (code, reason) => ws.close(code, reason), - bufferedAmount: () => peerBufferedAmount(ws), - }) - }, - async onMessage(event) { - if (!bridge) return - await forwardPeerMessage(bridge, event.data) - }, - onClose() { - bridge?.close() - }, - } + + return createBridgedSocketEvents((peer) => { + const sessionId = readPreviewCookie(c.req.header('cookie')) + const session = sessionId ? store.getSession(sessionId) : undefined + if (!session) { + throw new Error('Preview session expired. Reopen it from OpenCode Manager.') + } + + const upstream = new WebSocket( + upstreamWebSocketUrl(session, requestUrl), + protocols.length > 0 ? protocols : undefined, + ) + return bridgeWebSocket(upstream, peer) + }) }), ) diff --git a/backend/src/services/terminal.ts b/backend/src/services/terminal.ts index cba45276f..d32f42fdd 100644 --- a/backend/src/services/terminal.ts +++ b/backend/src/services/terminal.ts @@ -3,9 +3,11 @@ import { formatTerminalTitle, parseTerminalTitle } from '@opencode-manager/share import type { TerminalInfo, TerminalKind } from '@opencode-manager/shared/types' import type { OpenCodeClient } from './opencode/client' import type { CredentialProvider } from './credential-provider' +import { toWebSocketUrl } from './opencode/upstream' import { bridgeWebSocket, type WebSocketBridge, type WebSocketPeer } from '../utils/websocket-bridge' type OpenCodePty = Awaited>['data'][number] +type OpenCodePersistentPty = Awaited>[number] export interface CreateTerminalInput { kind: TerminalKind @@ -79,6 +81,7 @@ export class TerminalService { throw error } } + await this.removePersistentTerminals(directory) } async connect( @@ -101,9 +104,38 @@ export class TerminalService { socketUrl.searchParams.set('ticket', result.data.ticket) socketUrl.searchParams.set('location[directory]', directory) if (cursor !== undefined) socketUrl.searchParams.set('cursor', String(cursor)) - socketUrl.protocol = socketUrl.protocol === 'https:' ? 'wss:' : 'ws:' - return bridgeWebSocket(this.openSocket(socketUrl.toString()), peer) + return bridgeWebSocket(this.openSocket(toWebSocketUrl(socketUrl.toString())), peer) + } + + private async removePersistentTerminals(directory: string): Promise { + let cursor: string | undefined + do { + const page = await this.openCodeClient.api.session.list({ directory, cursor }) + for (const session of page.data) { + await this.removeSessionPersistentTerminals(session.id) + } + cursor = page.cursor.next ?? undefined + } while (cursor !== undefined) + } + + private async removeSessionPersistentTerminals(sessionID: string): Promise { + let terminals: OpenCodePersistentPty[] + try { + terminals = await this.openCodeClient.api.experimental.persistentPty.list({ sessionID }) + } catch (error: unknown) { + if (openCodeErrorStatus(error) === 404) return + throw error + } + + for (const terminal of terminals) { + try { + await this.openCodeClient.api.experimental.persistentPty.remove({ ptyID: terminal.id }) + } catch (error: unknown) { + if (openCodeErrorStatus(error) === 404) continue + throw error + } + } } private toTerminalInfo(pty: OpenCodePty): TerminalInfo { diff --git a/backend/src/utils/websocket-bridge.ts b/backend/src/utils/websocket-bridge.ts index 3e1941575..21945c623 100644 --- a/backend/src/utils/websocket-bridge.ts +++ b/backend/src/utils/websocket-bridge.ts @@ -1,3 +1,6 @@ +import type { WSContext, WSEvents } from 'hono/ws' +import { getErrorMessage } from './error-utils' + export interface WebSocketPeer { send(data: string | Uint8Array): void close(code?: number, reason?: string): void @@ -9,13 +12,57 @@ export interface WebSocketBridge { close(): void } +export interface WebSocketBridgeOptions { + initialMessages?: Array + connectTimeoutMs?: number +} + +export interface PendingUpstream { + attach(peer: WebSocketPeer): WebSocketBridge + abandon(): void +} + +export interface ConnectUpstreamOptions { + protocols?: string[] + timeoutMs?: number +} + export const MAX_PEER_BUFFERED_BYTES = 4 * 1024 * 1024 +export const MAX_UPSTREAM_BUFFERED_BYTES = 4 * 1024 * 1024 + +export const UPSTREAM_CONNECT_TIMEOUT_MS = 10_000 + +const MAX_CLOSE_REASON_BYTES = 120 + export function peerBufferedAmount(peer: { raw?: unknown }): number { const raw = peer.raw as { bufferedAmount?: number } | undefined return typeof raw?.bufferedAmount === 'number' ? raw.bufferedAmount : 0 } +export function toWebSocketPeer(ws: WSContext): WebSocketPeer { + return { + send: (data) => ws.send(data), + close: (code, reason) => ws.close(code, reason), + bufferedAmount: () => peerBufferedAmount(ws), + } +} + +export function isAllowedUpgradeOrigin( + origin: string | undefined, + trustedOrigins: readonly string[], +): boolean { + if (origin === undefined) return true + if (origin === 'null') return false + return trustedOrigins.includes(origin) +} + +function truncateCloseReason(message: string): string { + const bytes = Buffer.from(message, 'utf8') + if (bytes.length <= MAX_CLOSE_REASON_BYTES) return message + return bytes.subarray(0, MAX_CLOSE_REASON_BYTES).toString('utf8') +} + export async function forwardPeerMessage( bridge: WebSocketBridge, data: string | Blob | ArrayBufferLike, @@ -43,37 +90,73 @@ function toBytes(data: ArrayBuffer | Uint8Array): Uint8Array { return data instanceof Uint8Array ? (data as Uint8Array) : new Uint8Array(data) } -export function bridgeWebSocket(upstream: WebSocket, peer: WebSocketPeer): WebSocketBridge { +function toArrayBuffer(data: Uint8Array): ArrayBuffer { + return new Uint8Array(data).buffer +} + +export function bridgeWebSocket( + upstream: WebSocket, + peer: WebSocketPeer, + options: WebSocketBridgeOptions = {}, +): WebSocketBridge { upstream.binaryType = 'arraybuffer' const queue: Array = [] + let queuedBytes = 0 let upstreamOpen = false let finished = false + let connectTimer: ReturnType | undefined + + const clearConnectTimer = (): void => { + if (connectTimer !== undefined) { + clearTimeout(connectTimer) + connectTimer = undefined + } + } const finish = (): void => { finished = true + clearConnectTimer() } - upstream.addEventListener('open', () => { - upstreamOpen = true - while (queue.length > 0) { - upstream.send(queue.shift()!) - } - }) - - upstream.addEventListener('message', (event) => { - if (finished) return - const data = (event as MessageEvent).data + const deliverToPeer = (data: string | ArrayBuffer | Uint8Array): void => { if (typeof data === 'string') { peer.send(data) } else if (data instanceof ArrayBuffer || data instanceof Uint8Array) { peer.send(toBytes(data)) + } else { + return } if (peer.bufferedAmount() > MAX_PEER_BUFFERED_BYTES) { finish() upstream.close() peer.close(1013, 'Client too slow') } + } + + const failUpstream = (code: number, reason: string): void => { + if (finished) return + finish() + upstream.close() + peer.close(code, reason) + } + + upstream.addEventListener('open', () => { + if (finished) return + upstreamOpen = true + clearConnectTimer() + while (queue.length > 0) { + upstream.send(queue.shift()!) + if (upstream.bufferedAmount > MAX_UPSTREAM_BUFFERED_BYTES) { + failUpstream(1013, 'Upstream too slow') + return + } + } + }) + + upstream.addEventListener('message', (event) => { + if (finished) return + deliverToPeer((event as MessageEvent).data) }) upstream.addEventListener('close', (event) => { @@ -88,14 +171,40 @@ export function bridgeWebSocket(upstream: WebSocket, peer: WebSocketPeer): WebSo peer.close(1011) }) + const readyState = upstream.readyState + if (readyState === WebSocket.OPEN) { + upstreamOpen = true + } else if (readyState === WebSocket.CLOSING || readyState === WebSocket.CLOSED) { + finish() + peer.close(1011) + } else { + connectTimer = setTimeout(() => { + failUpstream(1011, 'Upstream connect timeout') + }, options.connectTimeoutMs ?? UPSTREAM_CONNECT_TIMEOUT_MS) + } + + if (!finished && options.initialMessages) { + for (const message of options.initialMessages) { + if (finished) break + deliverToPeer(message) + } + } + return { send(data) { if (finished) return const payload = typeof data === 'string' ? data : toBytes(data) if (upstreamOpen) { upstream.send(payload) - } else { - queue.push(payload) + if (upstream.bufferedAmount > MAX_UPSTREAM_BUFFERED_BYTES) { + failUpstream(1013, 'Upstream too slow') + } + return + } + queue.push(payload) + queuedBytes += typeof payload === 'string' ? Buffer.byteLength(payload) : payload.byteLength + if (queuedBytes > MAX_UPSTREAM_BUFFERED_BYTES) { + failUpstream(1013, 'Upstream too slow') } }, close() { @@ -105,3 +214,139 @@ export function bridgeWebSocket(upstream: WebSocket, peer: WebSocketPeer): WebSo }, } } + +export function connectUpstreamWebSocket( + url: string, + options: ConnectUpstreamOptions = {}, +): Promise { + return new Promise((resolve, reject) => { + let socket: WebSocket + try { + socket = new WebSocket(url, options.protocols) + } catch (error) { + reject(error) + return + } + socket.binaryType = 'arraybuffer' + + const timeoutMs = options.timeoutMs ?? UPSTREAM_CONNECT_TIMEOUT_MS + const buffered: Array = [] + let settled = false + let attached = false + let timer: ReturnType | undefined + + const clearTimer = (): void => { + if (timer !== undefined) { + clearTimeout(timer) + timer = undefined + } + } + + const detachListeners = (): void => { + socket.removeEventListener('open', handleOpen) + socket.removeEventListener('message', handleMessage) + socket.removeEventListener('error', handleError) + socket.removeEventListener('close', handleClose) + } + + const closeSocket = (): void => { + socket.close() + } + + const rejectBeforeOpen = (error: unknown): void => { + if (settled) return + settled = true + clearTimer() + detachListeners() + closeSocket() + reject(error) + } + + const handleOpen = (): void => { + if (settled) return + settled = true + clearTimer() + timer = setTimeout(() => { + if (!attached) closeSocket() + }, timeoutMs) + resolve({ + attach(peer) { + if (attached) throw new Error('Upstream WebSocket is already attached') + attached = true + clearTimer() + detachListeners() + return bridgeWebSocket(socket, peer, { initialMessages: buffered.splice(0, buffered.length) }) + }, + abandon() { + if (attached) return + attached = true + clearTimer() + detachListeners() + closeSocket() + }, + }) + } + + const handleMessage = (event: Event): void => { + if (attached) return + const data = (event as MessageEvent).data + if (typeof data === 'string') { + buffered.push(data) + } else if (data instanceof ArrayBuffer) { + buffered.push(data) + } else if (data instanceof Uint8Array) { + buffered.push(toArrayBuffer(data)) + } + } + + const handleError = (): void => { + rejectBeforeOpen(new Error(`Failed to connect to upstream WebSocket: ${url}`)) + } + + const handleClose = (): void => { + rejectBeforeOpen(new Error(`Upstream WebSocket closed before connect: ${url}`)) + } + + timer = setTimeout(() => { + rejectBeforeOpen(new Error(`Upstream WebSocket connect timed out: ${url}`)) + }, timeoutMs) + + socket.addEventListener('open', handleOpen) + socket.addEventListener('message', handleMessage) + socket.addEventListener('error', handleError) + socket.addEventListener('close', handleClose) + }) +} + +export function createBridgedSocketEvents( + openBridge: (peer: WebSocketPeer) => WebSocketBridge | Promise, +): WSEvents { + let bridge: WebSocketBridge | undefined + let closed = false + let connectPromise: Promise = Promise.resolve() + + return { + onOpen(_event, ws) { + connectPromise = (async () => { + const connected = await openBridge(toWebSocketPeer(ws)) + if (closed) { + connected.close() + return + } + bridge = connected + })().catch((error: unknown) => { + if (closed) return + ws.close(1011, truncateCloseReason(getErrorMessage(error))) + }) + }, + async onMessage(event) { + await connectPromise + if (!bridge) return + await forwardPeerMessage(bridge, event.data) + }, + onClose() { + closed = true + bridge?.close() + }, + } +} diff --git a/backend/test/helpers/websocket-server.ts b/backend/test/helpers/websocket-server.ts new file mode 100644 index 000000000..73122517c --- /dev/null +++ b/backend/test/helpers/websocket-server.ts @@ -0,0 +1,43 @@ +import type { Server } from 'node:http' +import type { AddressInfo } from 'node:net' +import type { Hono } from 'hono' +import type { UpgradeWebSocket } from 'hono/ws' +import { serve } from '@hono/node-server' +import { createNodeWebSocket } from '@hono/node-ws' + +export interface StartWebSocketServerOptions { + app: Hono + register?: (app: Hono, upgradeWebSocket: UpgradeWebSocket) => void + injectWebSocket?: (server: Server) => void +} + +export interface WebSocketTestServer { + server: Server + port: number + close: () => Promise +} + +/** + * Starts a Hono app on an ephemeral loopback port with WebSocket upgrades wired in. + * Creates its own node WebSocket binding unless the caller passes an existing `injectWebSocket`. + */ +export async function startWebSocketServer(options: StartWebSocketServerOptions): Promise { + const owned = options.injectWebSocket ? undefined : createNodeWebSocket({ app: options.app }) + if (owned && options.register) { + options.register(options.app, owned.upgradeWebSocket) + } + const server = await new Promise((resolve) => { + const started = serve({ fetch: options.app.fetch, port: 0, hostname: '127.0.0.1' }, () => resolve(started)) as unknown as Server + }) + const injectWebSocket = options.injectWebSocket ?? owned?.injectWebSocket + if (!injectWebSocket) throw new Error('startWebSocketServer could not resolve an injectWebSocket') + injectWebSocket(server) + return { + server, + port: (server.address() as AddressInfo).port, + close: () => { + server.closeAllConnections() + return new Promise((resolve) => server.close(() => resolve())) + }, + } +} diff --git a/backend/test/routes/opencode-proxy.test.ts b/backend/test/routes/opencode-proxy.test.ts index 9e9d2d201..cb491644f 100644 --- a/backend/test/routes/opencode-proxy.test.ts +++ b/backend/test/routes/opencode-proxy.test.ts @@ -1,19 +1,21 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest' -import { spawn } from 'child_process' -import { createServer } from 'http' +import { describe, it, expect, vi, beforeEach, afterEach, beforeAll, afterAll } from 'vitest' +import { execFileSync, spawn } from 'child_process' +import { createServer, request } from 'http' +import type { Server } from 'http' import type { AddressInfo } from 'net' import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from 'fs' import { tmpdir } from 'os' import { join } from 'path' import { Hono } from 'hono' -import { serve } from '@hono/node-server' +import { createNodeWebSocket } from '@hono/node-ws' import type { Database } from 'bun:sqlite' import { buildOpenCodeBasicAuth } from '@opencode-manager/shared/opencode' import { getWorkspacePath } from '@opencode-manager/shared/config/env' import { createOpenCodeProxyRoutes } from '../../src/routes/opencode-proxy' import type { SettingsService } from '../../src/services/settings' import { OpenCodeSupervisor } from '../../src/services/opencode-supervisor' -import { resolveOpenCode2Binary, startOpenCodeServe } from '../helpers/opencode-binary' +import { resolveOpenCode2Binary, startOpenCodeServe, type OpenCodeServe } from '../helpers/opencode-binary' +import { startWebSocketServer } from '../helpers/websocket-server' vi.mock('bun:sqlite', () => ({ Database: vi.fn(), @@ -52,6 +54,12 @@ const mockSettingsService = { const mockDb = {} as Database +function createProxyApp(): Hono { + const proxyApp = new Hono() + proxyApp.route('/api/opencode-proxy', createOpenCodeProxyRoutes(mockDb, mockSettingsService, createNodeWebSocket({ app: proxyApp }).upgradeWebSocket)) + return proxyApp +} + function upstreamOk() { const upstreamFetch = vi.fn().mockResolvedValue( new Response('ok', { status: 200, headers: { 'content-type': 'text/plain' } }) @@ -68,8 +76,7 @@ describe('opencode-proxy routes', () => { vi.clearAllMocks() isLifecycleInitializedMock.mockReturnValue(true) originalFetch = globalThis.fetch - app = new Hono() - app.route('/api/opencode-proxy', createOpenCodeProxyRoutes(mockDb, mockSettingsService)) + app = createProxyApp() }) afterEach(() => { @@ -230,6 +237,45 @@ describe('opencode-proxy routes', () => { expect(body.error).toContain('WebSocket') }) + it('still requires authorization for WebSocket upgrades without a PTY ticket', async () => { + const res = await app.request('/api/opencode-proxy/api/experimental/persistent-pty/pty_1/connect', { + headers: { Connection: 'Upgrade', Upgrade: 'websocket' }, + }) + expect(res.status).toBe(401) + }) + + it('requires authorization for a ticketed /api/pty connect', async () => { + const res = await app.request('/api/opencode-proxy/api/pty/pty_1/connect?ticket=abc', { + headers: { Connection: 'Upgrade', Upgrade: 'websocket' }, + }) + expect(res.status).toBe(401) + }) + + it('returns 403 for a ticketed persistent PTY connect from an untrusted origin', async () => { + const res = await app.request('/api/opencode-proxy/api/experimental/persistent-pty/pty_1/connect?ticket=abc', { + headers: { Connection: 'Upgrade', Upgrade: 'websocket', Origin: 'http://evil.example' }, + }) + expect(res.status).toBe(403) + const body = await res.json() as { error: string } + expect(body.error).toBe('Origin not allowed') + }) + + it('returns 503 for a ticketed persistent PTY connect before lifecycle initialization', async () => { + isLifecycleInitializedMock.mockReturnValue(false) + const res = await app.request('/api/opencode-proxy/api/experimental/persistent-pty/pty_1/connect?ticket=abc', { + headers: { Connection: 'Upgrade', Upgrade: 'websocket' }, + }) + expect(res.status).toBe(503) + }) + + it('does not resolve a repo for a ticketed persistent PTY connect through the repo mount', async () => { + const res = await app.request('/api/opencode-proxy/repos/999/api/experimental/persistent-pty/pty_1/connect?ticket=abc', { + headers: { Connection: 'Upgrade', Upgrade: 'websocket', Origin: 'http://evil.example' }, + }) + expect(res.status).toBe(403) + expect(getRepoByIdMock).not.toHaveBeenCalled() + }) + it('preserves SSE content-type header from upstream', async () => { const upstreamFetch = vi.fn().mockResolvedValue( new Response('event: message\ndata: hello\n\n', { @@ -770,8 +816,7 @@ describe('opencode-proxy repo-scoped mount', () => { vi.clearAllMocks() isLifecycleInitializedMock.mockReturnValue(true) originalFetch = globalThis.fetch - app = new Hono() - app.route('/api/opencode-proxy', createOpenCodeProxyRoutes(mockDb, mockSettingsService)) + app = createProxyApp() }) afterEach(() => { @@ -1082,17 +1127,21 @@ function runOpenCodeAgainstProxy( } describe.skipIf(SHIPPED_OPENCODE_BIN === null)('opencode-proxy repo-scoped mount against the shipped OpenCode 2 binary', () => { - it('pins the repo directory for a client started from an unrelated cwd', async () => { - const root = mkdtempSync(join(tmpdir(), 'ocm-proxy-e2e-')) - const configHome = join(root, 'config') - const homeDirectory = join(root, 'home') - const clientCwd = join(root, 'client-cwd') - const repoPath = join(root, 'repo') - for (const directory of [join(configHome, 'opencode'), homeDirectory, clientCwd, repoPath, join(root, 'data'), join(root, 'state'), join(root, 'cache')]) { + let instanceRoot: string + let configHome: string + let homeDirectory: string + let instance: OpenCodeServe + let llm: { port: number; close: () => void } + + beforeAll(async () => { + instanceRoot = mkdtempSync(join(tmpdir(), 'ocm-proxy-e2e-')) + configHome = join(instanceRoot, 'config') + homeDirectory = join(instanceRoot, 'home') + for (const directory of [join(configHome, 'opencode'), homeDirectory, join(instanceRoot, 'data'), join(instanceRoot, 'state'), join(instanceRoot, 'cache')]) { mkdirSync(directory, { recursive: true }) } - const llm = await startMockLlm() + llm = await startMockLlm() writeFileSync( join(configHome, 'opencode', 'opencode.json'), JSON.stringify({ @@ -1109,17 +1158,33 @@ describe.skipIf(SHIPPED_OPENCODE_BIN === null)('opencode-proxy repo-scoped mount }), ) - const instance = await startOpenCodeServe({ env: { HOME: homeDirectory, XDG_CONFIG_HOME: configHome } }) + instance = await startOpenCodeServe({ env: { HOME: homeDirectory, XDG_CONFIG_HOME: configHome } }) + }) + + afterEach(() => { + getOpenCodeServerPasswordMock.mockReturnValue('test-password') + upstreamBaseUrl.value = 'http://127.0.0.1:5551' + }) + + afterAll(async () => { + await instance.stop() + llm.close() + rmSync(instanceRoot, { recursive: true, force: true }) + }) + + it('pins the repo directory for a client started from an unrelated cwd', async () => { + const root = mkdtempSync(join(tmpdir(), 'ocm-proxy-e2e-')) + const clientCwd = join(root, 'client-cwd') + const repoPath = join(root, 'repo') + for (const directory of [clientCwd, repoPath]) { + mkdirSync(directory, { recursive: true }) + } + upstreamBaseUrl.value = instance.baseUrl getOpenCodeServerPasswordMock.mockReturnValue(instance.password) getRepoByIdMock.mockReturnValue({ id: 1, fullPath: repoPath, cloneStatus: 'ready' }) - const proxyApp = new Hono() - proxyApp.route('/api/opencode-proxy', createOpenCodeProxyRoutes(mockDb, mockSettingsService)) - const proxyServer = await new Promise>((resolve) => { - const server = serve({ fetch: proxyApp.fetch, port: 0, hostname: '127.0.0.1' }, () => resolve(server)) - }) - const proxyPort = (proxyServer.address() as AddressInfo).port + const { proxyServer, proxyPort } = await startProxyServer() try { const result = await runOpenCodeAgainstProxy( @@ -1137,11 +1202,164 @@ describe.skipIf(SHIPPED_OPENCODE_BIN === null)('opencode-proxy repo-scoped mount expect(body.data.some((session) => session.location.directory === repoPath)).toBe(true) } finally { await new Promise((resolve) => proxyServer.close(() => resolve())) - await instance.stop() - llm.close() - getOpenCodeServerPasswordMock.mockReturnValue('test-password') - upstreamBaseUrl.value = 'http://127.0.0.1:5551' + rmSync(root, { recursive: true, force: true }) + } + }, 120000) + + it('bridges a ticketed persistent PTY WebSocket for a remote client without an Authorization header', async () => { + const ptyDaemonsBefore = listPtyDaemonPids() + const root = mkdtempSync(join(tmpdir(), 'ocm-proxy-pty-')) + const repoPath = join(root, 'repo') + mkdirSync(repoPath, { recursive: true }) + + upstreamBaseUrl.value = instance.baseUrl + getOpenCodeServerPasswordMock.mockReturnValue(instance.password) + getRepoByIdMock.mockReturnValue({ id: 1, fullPath: repoPath, cloneStatus: 'ready' }) + + const { proxyServer, proxyPort } = await startProxyServer() + const repoProxy = `127.0.0.1:${proxyPort}/api/opencode-proxy/repos/1` + const clientAuth = 'Basic ' + Buffer.from('opencode:test-internal-token').toString('base64') + let leakedPtyDaemons: number[] = [] + + try { + const created = await fetch(`http://${repoProxy}/api/session`, { + method: 'POST', + headers: { Authorization: clientAuth, 'content-type': 'application/json' }, + body: JSON.stringify({ title: 'pty-bridge' }), + }) + expect(created.status, await created.clone().text()).toBe(200) + const sessionID = ((await created.json()) as { data: { id: string } }).data.id + + const terminal = await fetch(`http://${repoProxy}/api/experimental/session/${sessionID}/terminal`, { + method: 'POST', + headers: { Authorization: clientAuth, 'content-type': 'application/json' }, + body: JSON.stringify({ command: '/bin/sh', args: [], title: 'pty-bridge', env: {}, cwd: repoPath }), + }) + expect(terminal.status, await terminal.clone().text()).toBe(200) + const ptyID = ((await terminal.json()) as { data: { id: string } }).data.id + + const token = await fetch(`http://${repoProxy}/api/experimental/persistent-pty/${ptyID}/connect-token`, { + method: 'POST', + headers: { Authorization: clientAuth, 'x-opencode-ticket': '1' }, + }) + expect(token.status, await token.clone().text()).toBe(200) + const ticket = ((await token.json()) as { data: { ticket: string } }).data.ticket + + const socket = new WebSocket(`ws://${repoProxy}/api/experimental/persistent-pty/${ptyID}/connect?ticket=${encodeURIComponent(ticket)}`) + socket.binaryType = 'arraybuffer' + const output = await new Promise((resolve, reject) => { + let received = '' + let sent = false + const timer = setTimeout(() => reject(new Error(`timed out; received: ${received}`)), 20000) + socket.addEventListener('message', (event) => { + received += typeof event.data === 'string' ? event.data : new TextDecoder().decode(event.data as ArrayBuffer) + if (!sent && received.includes('replay_complete')) { + sent = true + socket.send('echo proxied-$((40+2))\n') + } + if (received.includes('proxied-42')) { + clearTimeout(timer) + resolve(received) + } + }) + socket.addEventListener('close', (event) => { + clearTimeout(timer) + reject(new Error(`socket closed ${event.code}; received: ${received}`)) + }) + }) + socket.close() + + expect(output).toContain('proxied-42') + } finally { + await shutdownPersistentPty(repoProxy, clientAuth) + await new Promise((resolve) => proxyServer.close(() => resolve())) + leakedPtyDaemons = await waitForNoNewPtyDaemons(ptyDaemonsBefore) + rmSync(root, { recursive: true, force: true }) + } + + expect(leakedPtyDaemons).toEqual([]) + }, 120000) + + it('rejects a persistent PTY connect with a bogus ticket as HTTP 403', async () => { + const root = mkdtempSync(join(tmpdir(), 'ocm-proxy-pty-')) + const repoPath = join(root, 'repo') + mkdirSync(repoPath, { recursive: true }) + + upstreamBaseUrl.value = instance.baseUrl + getOpenCodeServerPasswordMock.mockReturnValue(instance.password) + getRepoByIdMock.mockReturnValue({ id: 1, fullPath: repoPath, cloneStatus: 'ready' }) + + const { proxyServer, proxyPort } = await startProxyServer() + + try { + const status = await readUpgradeRejectionStatus( + proxyPort, + '/api/opencode-proxy/repos/1/api/experimental/persistent-pty/pty_bogus/connect?ticket=bogus', + ) + expect(status).toBe(403) + } finally { + await new Promise((resolve) => proxyServer.close(() => resolve())) rmSync(root, { recursive: true, force: true }) } }, 120000) }) + +function listPtyDaemonPids(): number[] { + try { + return execFileSync('pgrep', ['-f', 'opencode-pty'], { encoding: 'utf8' }) + .split('\n') + .map((line) => Number(line.trim())) + .filter((pid) => Number.isInteger(pid) && pid > 0) + } catch { + return [] + } +} + +async function waitForNoNewPtyDaemons(baseline: number[]): Promise { + const deadline = Date.now() + 10000 + let leaked: number[] = [] + while (Date.now() < deadline) { + leaked = listPtyDaemonPids().filter((pid) => !baseline.includes(pid)) + if (leaked.length === 0) return [] + await new Promise((resolve) => setTimeout(resolve, 250)) + } + return leaked +} + +async function shutdownPersistentPty(repoProxy: string, clientAuth: string): Promise { + await fetch(`http://${repoProxy}/api/experimental/persistent-pty/shutdown`, { + method: 'POST', + headers: { Authorization: clientAuth }, + }).catch(() => undefined) +} + +function readUpgradeRejectionStatus(port: number, path: string): Promise { + return new Promise((resolve, reject) => { + const req = request({ + host: '127.0.0.1', + port, + path, + headers: { Connection: 'Upgrade', Upgrade: 'websocket' }, + }) + req.on('response', (res) => { + res.resume() + resolve(res.statusCode ?? 0) + }) + req.on('upgrade', (res) => { + res.destroy() + resolve(101) + }) + req.on('error', reject) + req.end() + }) +} + +async function startProxyServer(): Promise<{ proxyServer: Server; proxyPort: number }> { + const { server, port } = await startWebSocketServer({ + app: new Hono(), + register: (proxyApp, upgradeWebSocket) => { + proxyApp.route('/api/opencode-proxy', createOpenCodeProxyRoutes(mockDb, mockSettingsService, upgradeWebSocket)) + }, + }) + return { proxyServer: server, proxyPort: port } +} diff --git a/backend/test/routes/repo-terminal-socket.test.ts b/backend/test/routes/repo-terminal-socket.test.ts index f82c9f697..4a2685dcc 100644 --- a/backend/test/routes/repo-terminal-socket.test.ts +++ b/backend/test/routes/repo-terminal-socket.test.ts @@ -3,9 +3,7 @@ import { randomBytes } from 'node:crypto' import { request as httpRequest } from 'node:http' import { Hono } from 'hono' import { cors } from 'hono/cors' -import { serve } from '@hono/node-server' import { createNodeWebSocket } from '@hono/node-ws' -import type { Server } from 'node:http' import type { Database } from 'bun:sqlite' import type { Repo } from '@opencode-manager/shared/types' import { createRepoTerminalSocketRoutes } from '../../src/routes/repo-terminal-socket' @@ -13,6 +11,7 @@ import { TerminalService } from '../../src/services/terminal' import type { GitAuthService } from '../../src/services/git-auth' import type { OpenCodeClient } from '../../src/services/opencode/client' import type { CredentialProvider } from '../../src/services/credential-provider' +import { startWebSocketServer, type WebSocketTestServer } from '../helpers/websocket-server' vi.mock('../../src/services/repo', () => ({ resolveRepoOrAssistant: vi.fn(), @@ -49,8 +48,9 @@ const openCodeClient = { const credentialProvider = { getGhCliEnv: vi.fn(() => ({})) } as unknown as CredentialProvider -let upstreamServer: Server -let managerServer: Server +let upstream: WebSocketTestServer +let manager: WebSocketTestServer +let authenticated: WebSocketTestServer let managerPort = 0 let upstreamBase = '' let capturedUrl: URL | undefined @@ -117,7 +117,6 @@ authenticatedApp.route( ), ) -let authenticatedServer: Server let authenticatedPort = 0 interface RawHandshakeResult { @@ -155,27 +154,6 @@ function rawHandshake(options: { port: number; path: string; origin?: string; co }) } -function listen(server: Server): Promise { - return new Promise((resolve) => { - if (server.listening) { - resolve(readPort(server)) - return - } - server.once('listening', () => resolve(readPort(server))) - }) -} - -function readPort(server: Server): number { - const address = server.address() - if (address && typeof address === 'object') return address.port - throw new Error('Server is not listening on a TCP port') -} - -function closeServer(server: Server): Promise { - server.closeAllConnections() - return new Promise((resolve) => server.close(() => resolve())) -} - function clientUrl(ptyID: string, directory = '/tmp/repo'): string { const url = new URL(`ws://127.0.0.1:${managerPort}/1/terminals/${ptyID}/connect`) url.searchParams.set('directory', directory) @@ -210,17 +188,14 @@ async function waitFor(predicate: () => boolean, timeoutMs = 5000): Promise { - upstreamServer = serve({ fetch: upstreamApp.fetch, port: 0, hostname: '127.0.0.1' }) as unknown as Server - upstreamWs.injectWebSocket(upstreamServer) - upstreamBase = `http://127.0.0.1:${await listen(upstreamServer)}` + upstream = await startWebSocketServer({ app: upstreamApp, injectWebSocket: upstreamWs.injectWebSocket }) + upstreamBase = `http://127.0.0.1:${upstream.port}` - managerServer = serve({ fetch: managerApp.fetch, port: 0, hostname: '127.0.0.1' }) as unknown as Server - managerWs.injectWebSocket(managerServer) - managerPort = await listen(managerServer) + manager = await startWebSocketServer({ app: managerApp, injectWebSocket: managerWs.injectWebSocket }) + managerPort = manager.port - authenticatedServer = serve({ fetch: authenticatedApp.fetch, port: 0, hostname: '127.0.0.1' }) as unknown as Server - authenticatedWs.injectWebSocket(authenticatedServer) - authenticatedPort = await listen(authenticatedServer) + authenticated = await startWebSocketServer({ app: authenticatedApp, injectWebSocket: authenticatedWs.injectWebSocket }) + authenticatedPort = authenticated.port vi.mocked(resolveRepoOrAssistant).mockImplementation((_db, id) => (id === 1 ? readyRepo : null)) vi.mocked(resolveRepoWorkingDirectory).mockImplementation(async (_repo, directory) => { @@ -234,9 +209,9 @@ afterAll(async () => { upstreamWs.wss.close() managerWs.wss.close() authenticatedWs.wss.close() - await closeServer(upstreamServer) - await closeServer(managerServer) - await closeServer(authenticatedServer) + await upstream.close() + await manager.close() + await authenticated.close() }) describe('Repo Terminal Socket Route', () => { diff --git a/backend/test/services/opencode/upstream.test.ts b/backend/test/services/opencode/upstream.test.ts index 83e032bdc..1661ba8ea 100644 --- a/backend/test/services/opencode/upstream.test.ts +++ b/backend/test/services/opencode/upstream.test.ts @@ -1,6 +1,6 @@ import { describe, it, expect, afterEach } from 'vitest' import { ENV, getWorkspacePath } from '@opencode-manager/shared/config/env' -import { getOpenCodeUpstreamBaseUrl, withDefaultOpenCodeDirectory } from '../../../src/services/opencode/upstream' +import { getOpenCodeUpstreamBaseUrl, toWebSocketUrl, withDefaultOpenCodeDirectory } from '../../../src/services/opencode/upstream' const originalHost = ENV.OPENCODE.HOST @@ -58,3 +58,25 @@ describe('withDefaultOpenCodeDirectory', () => { expect(withDefaultOpenCodeDirectory(headers)).toBe(headers) }) }) + +describe('toWebSocketUrl', () => { + it('maps an http URL to ws and preserves the path, port and query', () => { + expect(toWebSocketUrl('http://127.0.0.1:5551/api/pty/pty-1/connect?ticket=abc')).toBe( + 'ws://127.0.0.1:5551/api/pty/pty-1/connect?ticket=abc', + ) + }) + + it('maps an https URL to wss', () => { + expect(toWebSocketUrl('https://opencode.internal/api/pty/pty-1/connect')).toBe( + 'wss://opencode.internal/api/pty/pty-1/connect', + ) + }) + + it('leaves a ws URL unchanged', () => { + expect(toWebSocketUrl('ws://127.0.0.1:5551/api/pty/pty-1/connect')).toBe('ws://127.0.0.1:5551/api/pty/pty-1/connect') + }) + + it('leaves a wss URL unchanged', () => { + expect(toWebSocketUrl('wss://opencode.internal/api')).toBe('wss://opencode.internal/api') + }) +}) diff --git a/backend/test/services/preview/gateway-ws.test.ts b/backend/test/services/preview/gateway-ws.test.ts index 00ce9ce25..b4f4b7b6b 100644 --- a/backend/test/services/preview/gateway-ws.test.ts +++ b/backend/test/services/preview/gateway-ws.test.ts @@ -1,16 +1,14 @@ import { describe, it, expect, beforeAll, afterAll } from 'vitest' import { randomBytes } from 'node:crypto' import { request as httpRequest } from 'node:http' -import type { Server } from 'node:http' -import type { AddressInfo } from 'node:net' import type { Duplex } from 'node:stream' import { Hono } from 'hono' -import { serve } from '@hono/node-server' import { createNodeWebSocket } from '@hono/node-ws' import { createPreviewGatewayApp, PreviewSessionStore, } from '../../../src/services/preview/gateway' +import { startWebSocketServer, type WebSocketTestServer } from '../../helpers/websocket-server' class HandshakeError extends Error { constructor(readonly statusCode: number) { @@ -156,8 +154,8 @@ class RawWebSocketClient { } } -let upstreamServer: Server -let gatewayServer: Server +let upstream: WebSocketTestServer +let gatewayServer: WebSocketTestServer let upstreamPort = 0 let gatewayPort = 0 let capturedUrl: URL | undefined @@ -184,21 +182,6 @@ upstreamApp.get( upstreamApp.all('/*', (c) => c.json({ path: c.req.path, search: new URL(c.req.url).search })) -function listen(server: Server): Promise { - return new Promise((resolve) => { - if (server.listening) { - resolve((server.address() as AddressInfo).port) - return - } - server.once('listening', () => resolve((server.address() as AddressInfo).port)) - }) -} - -function closeServer(server: Server): Promise { - server.closeAllConnections() - return new Promise((resolve) => server.close(() => resolve())) -} - async function openSessionCookie(): Promise { const token = store.issueStartToken({ port: upstreamPort, host: '127.0.0.1' }) const res = await gateway.app.request( @@ -211,19 +194,17 @@ async function openSessionCookie(): Promise { } beforeAll(async () => { - upstreamServer = serve({ fetch: upstreamApp.fetch, port: 0, hostname: '127.0.0.1' }) as unknown as Server - upstreamWs.injectWebSocket(upstreamServer) - upstreamPort = await listen(upstreamServer) + upstream = await startWebSocketServer({ app: upstreamApp, injectWebSocket: upstreamWs.injectWebSocket }) + upstreamPort = upstream.port - gatewayServer = serve({ fetch: gateway.app.fetch, port: 0, hostname: '127.0.0.1' }) as unknown as Server - gateway.injectWebSocket(gatewayServer) - gatewayPort = await listen(gatewayServer) + gatewayServer = await startWebSocketServer({ app: gateway.app, injectWebSocket: gateway.injectWebSocket }) + gatewayPort = gatewayServer.port }) afterAll(async () => { upstreamWs.wss.close() - await closeServer(upstreamServer) - await closeServer(gatewayServer) + await upstream.close() + await gatewayServer.close() }) describe('Preview Gateway WebSocket', () => { diff --git a/backend/test/services/terminal.test.ts b/backend/test/services/terminal.test.ts index 75dee51b2..d9a1d71b1 100644 --- a/backend/test/services/terminal.test.ts +++ b/backend/test/services/terminal.test.ts @@ -18,8 +18,8 @@ function createPty(overrides: Partial = {}): PtyRecord { return { id: 'pty-1', title: 'Terminal', command: 'sh', args: [], cwd: '/repo', status: 'running', pid: 1, ...overrides } } -function createClient(pty: Record): OpenCodeClient { - return { api: { pty } } as unknown as OpenCodeClient +function createClient(pty: Record, api: Record = {}): OpenCodeClient { + return { api: { pty, ...api } } as unknown as OpenCodeClient } function createCredentialProvider(shellEnv: Record = {}): CredentialProvider { @@ -30,8 +30,9 @@ function createTerminalService( client: OpenCodeClient, provider: CredentialProvider, openSocket?: (url: string) => WebSocket, + baseUrl = 'http://127.0.0.1:5551', ): TerminalService { - return new TerminalService(client, provider, () => 'http://127.0.0.1:5551', openSocket) + return new TerminalService(client, provider, () => baseUrl, openSocket) } class FakeUpstreamSocket { @@ -211,6 +212,16 @@ describe('TerminalService', () => { }) describe('removeAll', () => { + function createSessionApi(pages: Array<{ data: Array<{ id: string }>; next: string | null }>) { + let index = 0 + const list = vi.fn(async () => { + const page = pages[index] ?? { data: [], next: null } + index += 1 + return { data: page.data, cursor: { next: page.next } } + }) + return { list } + } + it('removes every terminal and tolerates a 404', async () => { const list = vi.fn(async () => ({ location: { directory: '/repo' }, @@ -220,7 +231,9 @@ describe('TerminalService', () => { const remove = vi.fn(async ({ ptyID }: { ptyID: string }) => { if (ptyID === 'a') throw notFound }) - const service = createTerminalService(createClient({ list, remove }), createCredentialProvider()) + const { list: sessionList } = createSessionApi([]) + const client = createClient({ list, remove }, { session: { list: sessionList } }) + const service = createTerminalService(client, createCredentialProvider()) await expect(service.removeAll('/repo')).resolves.toBeUndefined() expect(remove).toHaveBeenCalledTimes(2) @@ -236,7 +249,129 @@ describe('TerminalService', () => { const remove = vi.fn(async () => { throw new Error('boom') }) - const service = createTerminalService(createClient({ list, remove }), createCredentialProvider()) + const { list: sessionList } = createSessionApi([]) + const client = createClient({ list, remove }, { session: { list: sessionList } }) + const service = createTerminalService(client, createCredentialProvider()) + + await expect(service.removeAll('/repo')).rejects.toThrow('boom') + expect(sessionList).not.toHaveBeenCalled() + }) + + it('removes every persistent PTY across paginated sessions', async () => { + const list = vi.fn(async () => ({ location: { directory: '/repo' }, data: [] })) + const remove = vi.fn(async () => undefined) + const { list: sessionList } = createSessionApi([ + { data: [{ id: 'ses-1' }], next: 'page-2' }, + { data: [{ id: 'ses-2' }], next: null }, + ]) + const persistentList = vi.fn(async ({ sessionID }: { sessionID: string }) => [ + { id: `${sessionID}-a` }, + { id: `${sessionID}-b` }, + ]) + const persistentRemove = vi.fn(async () => undefined) + const client = createClient( + { list, remove }, + { + session: { list: sessionList }, + experimental: { persistentPty: { list: persistentList, remove: persistentRemove } }, + }, + ) + const service = createTerminalService(client, createCredentialProvider()) + + await expect(service.removeAll('/repo')).resolves.toBeUndefined() + + expect(sessionList).toHaveBeenNthCalledWith(1, { directory: '/repo' }) + expect(sessionList).toHaveBeenNthCalledWith(2, { directory: '/repo', cursor: 'page-2' }) + expect(persistentList).toHaveBeenCalledWith({ sessionID: 'ses-1' }) + expect(persistentList).toHaveBeenCalledWith({ sessionID: 'ses-2' }) + expect(persistentRemove).toHaveBeenCalledTimes(4) + expect(persistentRemove).toHaveBeenCalledWith({ ptyID: 'ses-1-a' }) + expect(persistentRemove).toHaveBeenCalledWith({ ptyID: 'ses-1-b' }) + expect(persistentRemove).toHaveBeenCalledWith({ ptyID: 'ses-2-a' }) + expect(persistentRemove).toHaveBeenCalledWith({ ptyID: 'ses-2-b' }) + }) + + it('treats a 404 from listing persistent PTYs as none', async () => { + const list = vi.fn(async () => ({ location: { directory: '/repo' }, data: [] })) + const remove = vi.fn(async () => undefined) + const { list: sessionList } = createSessionApi([{ data: [{ id: 'ses-1' }], next: null }]) + const notFound = Object.assign(new Error('not found'), { _tag: 'PtyNotFoundError' }) + const persistentList = vi.fn(async () => { + throw notFound + }) + const persistentRemove = vi.fn(async () => undefined) + const client = createClient( + { list, remove }, + { + session: { list: sessionList }, + experimental: { persistentPty: { list: persistentList, remove: persistentRemove } }, + }, + ) + const service = createTerminalService(client, createCredentialProvider()) + + await expect(service.removeAll('/repo')).resolves.toBeUndefined() + expect(persistentRemove).not.toHaveBeenCalled() + }) + + it('tolerates a 404 while removing a persistent PTY', async () => { + const list = vi.fn(async () => ({ location: { directory: '/repo' }, data: [] })) + const remove = vi.fn(async () => undefined) + const { list: sessionList } = createSessionApi([{ data: [{ id: 'ses-1' }], next: null }]) + const notFound = Object.assign(new Error('not found'), { _tag: 'PtyNotFoundError' }) + const persistentList = vi.fn(async () => [{ id: 'pp-1' }, { id: 'pp-2' }]) + const persistentRemove = vi.fn(async ({ ptyID }: { ptyID: string }) => { + if (ptyID === 'pp-1') throw notFound + }) + const client = createClient( + { list, remove }, + { + session: { list: sessionList }, + experimental: { persistentPty: { list: persistentList, remove: persistentRemove } }, + }, + ) + const service = createTerminalService(client, createCredentialProvider()) + + await expect(service.removeAll('/repo')).resolves.toBeUndefined() + expect(persistentRemove).toHaveBeenCalledTimes(2) + expect(persistentRemove).toHaveBeenCalledWith({ ptyID: 'pp-1' }) + expect(persistentRemove).toHaveBeenCalledWith({ ptyID: 'pp-2' }) + }) + + it('propagates a non-404 failure from listing persistent PTYs', async () => { + const list = vi.fn(async () => ({ location: { directory: '/repo' }, data: [] })) + const remove = vi.fn(async () => undefined) + const { list: sessionList } = createSessionApi([{ data: [{ id: 'ses-1' }], next: null }]) + const persistentList = vi.fn(async () => { + throw new Error('boom') + }) + const client = createClient( + { list, remove }, + { + session: { list: sessionList }, + experimental: { persistentPty: { list: persistentList, remove: vi.fn(async () => undefined) } }, + }, + ) + const service = createTerminalService(client, createCredentialProvider()) + + await expect(service.removeAll('/repo')).rejects.toThrow('boom') + }) + + it('propagates a non-404 failure from removing a persistent PTY', async () => { + const list = vi.fn(async () => ({ location: { directory: '/repo' }, data: [] })) + const remove = vi.fn(async () => undefined) + const { list: sessionList } = createSessionApi([{ data: [{ id: 'ses-1' }], next: null }]) + const persistentList = vi.fn(async () => [{ id: 'pp-1' }]) + const persistentRemove = vi.fn(async () => { + throw new Error('boom') + }) + const client = createClient( + { list, remove }, + { + session: { list: sessionList }, + experimental: { persistentPty: { list: persistentList, remove: persistentRemove } }, + }, + ) + const service = createTerminalService(client, createCredentialProvider()) await expect(service.removeAll('/repo')).rejects.toThrow('boom') }) @@ -303,6 +438,26 @@ describe('TerminalService', () => { expect(url.searchParams.has('cursor')).toBe(false) }) + it('uses wss when the upstream base is https', async () => { + const { client } = createConnectClient() + let opened: FakeUpstreamSocket | undefined + const service = createTerminalService( + client, + createCredentialProvider(), + (url) => { + opened = new FakeUpstreamSocket(url) + return opened as unknown as WebSocket + }, + 'https://opencode.internal', + ) + + await service.connect('/repo', 'pty-1', undefined, createPeer()) + + const url = new URL(opened!.url) + expect(url.protocol).toBe('wss:') + expect(url.pathname).toBe('/api/pty/pty-1/connect') + }) + it('rejects an unknown terminal before requesting a ticket', async () => { const { client, token } = createConnectClient({ terminals: [] }) const service = createTerminalService(client, createCredentialProvider()) diff --git a/backend/test/utils/websocket-bridge.test.ts b/backend/test/utils/websocket-bridge.test.ts index cdf00630b..9967ca5a7 100644 --- a/backend/test/utils/websocket-bridge.test.ts +++ b/backend/test/utils/websocket-bridge.test.ts @@ -1,8 +1,14 @@ -import { describe, it, expect, vi } from 'vitest' +import { describe, it, expect, vi, afterEach } from 'vitest' +import type { WSContext } from 'hono/ws' import { bridgeWebSocket, + connectUpstreamWebSocket, + createBridgedSocketEvents, forwardPeerMessage, + isAllowedUpgradeOrigin, + toWebSocketPeer, MAX_PEER_BUFFERED_BYTES, + MAX_UPSTREAM_BUFFERED_BYTES, type WebSocketBridge, type WebSocketPeer, } from '../../src/utils/websocket-bridge' @@ -11,6 +17,8 @@ type Listener = (event: unknown) => void class FakeUpstream { binaryType: 'blob' | 'arraybuffer' = 'blob' + readyState = 0 + bufferedAmount = 0 sent: Array = [] private readonly listeners = new Map>() @@ -59,11 +67,127 @@ function asUpstream(upstream: FakeUpstream): WebSocket { return upstream as unknown as WebSocket } +interface FakeWs { + raw: { bufferedAmount: number } + sent: Array + closeCalls: Array<{ code?: number; reason?: string }> + send(data: string | ArrayBuffer | Uint8Array): void + close(code?: number, reason?: string): void +} + +function createWs(): FakeWs { + const ws: FakeWs = { + raw: { bufferedAmount: 0 }, + sent: [], + closeCalls: [], + send(data) { + ws.sent.push(data) + }, + close(code, reason) { + ws.closeCalls.push({ code, reason }) + }, + } + return ws +} + +function asWsContext(ws: FakeWs): WSContext { + return ws as unknown as WSContext +} + +interface FakeBridge extends WebSocketBridge { + sent: Array + closeCalls: number +} + +function createBridge(): FakeBridge { + const bridge: FakeBridge = { + sent: [], + closeCalls: 0, + send(data) { + bridge.sent.push(data) + }, + close() { + bridge.closeCalls += 1 + }, + } + return bridge +} + +function flushAsync(): Promise { + return new Promise((resolve) => setTimeout(resolve, 0)) +} + +class ControllableUpstream { + static readonly CONNECTING = 0 + static readonly OPEN = 1 + static readonly CLOSING = 2 + static readonly CLOSED = 3 + static instances: ControllableUpstream[] = [] + + binaryType: 'blob' | 'arraybuffer' = 'blob' + readyState = ControllableUpstream.CONNECTING + bufferedAmount = 0 + sent: Array = [] + closeCalls = 0 + private readonly listeners = new Map>() + + constructor( + readonly url: string, + readonly protocols?: string | string[], + ) { + ControllableUpstream.instances.push(this) + } + + static last(): ControllableUpstream { + const instance = ControllableUpstream.instances.at(-1) + if (!instance) throw new Error('No ControllableUpstream instance') + return instance + } + + static reset(): void { + ControllableUpstream.instances = [] + } + + addEventListener(type: string, listener: Listener): void { + const set = this.listeners.get(type) ?? new Set() + set.add(listener) + this.listeners.set(type, set) + } + + removeEventListener(type: string, listener: Listener): void { + this.listeners.get(type)?.delete(listener) + } + + send(data: string | Uint8Array): void { + this.sent.push(data) + } + + close(): void { + this.closeCalls += 1 + this.readyState = ControllableUpstream.CLOSED + } + + open(): void { + this.readyState = ControllableUpstream.OPEN + this.emit('open', {}) + } + + emit(type: string, event: unknown): void { + for (const listener of [...(this.listeners.get(type) ?? [])]) listener(event) + } +} + +afterEach(() => { + vi.useRealTimers() + vi.unstubAllGlobals() +}) + describe('bridgeWebSocket', () => { it('requests arraybuffer frames from the upstream', () => { const upstream = new FakeUpstream() - bridgeWebSocket(asUpstream(upstream), createPeer()) + const bridge = bridgeWebSocket(asUpstream(upstream), createPeer()) expect(upstream.binaryType).toBe('arraybuffer') + bridge.close() }) it('queues client sends until the upstream opens and preserves order', () => { @@ -82,13 +206,14 @@ describe('bridgeWebSocket', () => { it('forwards text as text and binary as bytes to the peer', () => { const upstream = new FakeUpstream() const peer = createPeer() - bridgeWebSocket(asUpstream(upstream), peer) + const bridge = bridgeWebSocket(asUpstream(upstream), peer) upstream.emit('message', { data: 'plain text' }) upstream.emit('message', { data: new Uint8Array([0, 1, 2]).buffer }) expect(peer.sent[0]).toBe('plain text') expect(peer.sent[1]).toEqual(new Uint8Array([0, 1, 2])) + bridge.close() }) it('passes a normal close through and maps abnormal codes to 1011', () => { @@ -141,12 +266,13 @@ describe('bridgeWebSocket', () => { const peer = createPeer() peer.buffered = MAX_PEER_BUFFERED_BYTES - bridgeWebSocket(asUpstream(upstream), peer) + const bridge = bridgeWebSocket(asUpstream(upstream), peer) upstream.emit('message', { data: 'within limit' }) expect(peer.sent).toEqual(['within limit']) expect(peer.closeCalls).toEqual([]) expect(upstream.close).not.toHaveBeenCalled() + bridge.close() }) it('closes the upstream only once and ignores later sends', () => { @@ -160,10 +286,335 @@ describe('bridgeWebSocket', () => { expect(upstream.close).toHaveBeenCalledTimes(1) expect(upstream.sent).toEqual([]) }) + + it('treats an already-open upstream as open immediately', () => { + const upstream = new FakeUpstream() + upstream.readyState = WebSocket.OPEN + const bridge = bridgeWebSocket(asUpstream(upstream), createPeer()) + + bridge.send('immediate') + + expect(upstream.sent).toEqual(['immediate']) + bridge.close() + }) + + it('closes the peer with 1011 when the upstream is already closed', () => { + const upstream = new FakeUpstream() + upstream.readyState = WebSocket.CLOSED + const peer = createPeer() + const bridge = bridgeWebSocket(asUpstream(upstream), peer) + + expect(peer.closeCalls).toEqual([{ code: 1011, reason: undefined }]) + + bridge.send('ignored') + expect(upstream.sent).toEqual([]) + }) + + it('closes both sides with 1013 when queued client bytes exceed the cap', () => { + const upstream = new FakeUpstream() + const peer = createPeer() + const bridge = bridgeWebSocket(asUpstream(upstream), peer) + + bridge.send('a'.repeat(MAX_UPSTREAM_BUFFERED_BYTES + 1)) + + expect(peer.closeCalls).toEqual([{ code: 1013, reason: 'Upstream too slow' }]) + expect(upstream.close).toHaveBeenCalledTimes(1) + }) + + it('closes both sides with 1013 when the upstream buffer exceeds the cap after a send', () => { + const upstream = new FakeUpstream() + upstream.readyState = WebSocket.OPEN + upstream.bufferedAmount = MAX_UPSTREAM_BUFFERED_BYTES + 1 + const peer = createPeer() + const bridge = bridgeWebSocket(asUpstream(upstream), peer) + + bridge.send('payload') + + expect(upstream.sent).toEqual(['payload']) + expect(peer.closeCalls).toEqual([{ code: 1013, reason: 'Upstream too slow' }]) + }) + + it('closes both sides with 1013 when the upstream buffer overflows during a queued flush', () => { + const upstream = new FakeUpstream() + const peer = createPeer() + const bridge = bridgeWebSocket(asUpstream(upstream), peer) + + bridge.send('queued') + upstream.bufferedAmount = MAX_UPSTREAM_BUFFERED_BYTES + 1 + upstream.emit('open', {}) + + expect(peer.closeCalls).toEqual([{ code: 1013, reason: 'Upstream too slow' }]) + expect(upstream.close).toHaveBeenCalledTimes(1) + }) + + it('closes both sides with 1011 when the upstream does not open in time', () => { + vi.useFakeTimers() + const upstream = new FakeUpstream() + const peer = createPeer() + bridgeWebSocket(asUpstream(upstream), peer, { connectTimeoutMs: 50 }) + + vi.advanceTimersByTime(50) + + expect(upstream.close).toHaveBeenCalledTimes(1) + expect(peer.closeCalls).toEqual([{ code: 1011, reason: 'Upstream connect timeout' }]) + }) + + it('clears the connect deadline once the upstream opens', () => { + vi.useFakeTimers() + const upstream = new FakeUpstream() + const peer = createPeer() + bridgeWebSocket(asUpstream(upstream), peer, { connectTimeoutMs: 50 }) + + upstream.emit('open', {}) + vi.advanceTimersByTime(1000) + + expect(peer.closeCalls).toEqual([]) + expect(upstream.close).not.toHaveBeenCalled() + }) + + it('delivers initialMessages before later upstream messages', () => { + const upstream = new FakeUpstream() + const peer = createPeer() + const bridge = bridgeWebSocket(asUpstream(upstream), peer, { + initialMessages: ['early text', new Uint8Array([9]).buffer], + }) + + upstream.emit('message', { data: 'later' }) + + expect(peer.sent).toEqual(['early text', new Uint8Array([9]), 'later']) + bridge.close() + }) + + it('applies the peer buffer cap while delivering initialMessages', () => { + const upstream = new FakeUpstream() + const peer = createPeer() + peer.buffered = MAX_PEER_BUFFERED_BYTES + 1 + + bridgeWebSocket(asUpstream(upstream), peer, { initialMessages: ['overflow'] }) + + expect(peer.closeCalls).toEqual([{ code: 1013, reason: 'Client too slow' }]) + expect(upstream.close).toHaveBeenCalledTimes(1) + }) +}) + +describe('connectUpstreamWebSocket', () => { + it('resolves on open, forwards protocols and buffers messages until attach', async () => { + vi.stubGlobal('WebSocket', ControllableUpstream) + ControllableUpstream.reset() + + const promise = connectUpstreamWebSocket('ws://upstream', { protocols: ['vite-hmr'] }) + const socket = ControllableUpstream.last() + expect(socket.protocols).toEqual(['vite-hmr']) + expect(socket.binaryType).toBe('arraybuffer') + + socket.open() + const pending = await promise + + socket.emit('message', { data: 'early' }) + const peer = createPeer() + const bridge = pending.attach(peer) + + expect(peer.sent).toEqual(['early']) + + socket.emit('message', { data: new Uint8Array([1, 2]).buffer }) + expect(peer.sent[1]).toEqual(new Uint8Array([1, 2])) + + bridge.close() + }) + + it('rejects and closes the socket when the upstream errors before open', async () => { + vi.stubGlobal('WebSocket', ControllableUpstream) + ControllableUpstream.reset() + + const promise = connectUpstreamWebSocket('ws://upstream') + const socket = ControllableUpstream.last() + socket.emit('error', {}) + + await expect(promise).rejects.toThrow('Failed to connect') + expect(socket.closeCalls).toBe(1) + }) + + it('rejects and closes the socket when the upstream closes before open', async () => { + vi.stubGlobal('WebSocket', ControllableUpstream) + ControllableUpstream.reset() + + const promise = connectUpstreamWebSocket('ws://upstream') + const socket = ControllableUpstream.last() + socket.emit('close', {}) + + await expect(promise).rejects.toThrow('closed before connect') + expect(socket.closeCalls).toBe(1) + }) + + it('rejects and closes the socket when the connect times out', async () => { + vi.useFakeTimers() + vi.stubGlobal('WebSocket', ControllableUpstream) + ControllableUpstream.reset() + + const promise = connectUpstreamWebSocket('ws://upstream', { timeoutMs: 50 }) + const socket = ControllableUpstream.last() + const rejection = expect(promise).rejects.toThrow('timed out') + + await vi.advanceTimersByTimeAsync(50) + await rejection + + expect(socket.closeCalls).toBe(1) + }) + + it('closes the socket when attach is not called within the timeout', async () => { + vi.useFakeTimers() + vi.stubGlobal('WebSocket', ControllableUpstream) + ControllableUpstream.reset() + + const promise = connectUpstreamWebSocket('ws://upstream', { timeoutMs: 50 }) + const socket = ControllableUpstream.last() + socket.open() + const pending = await promise + + expect(pending).toBeDefined() + await vi.advanceTimersByTimeAsync(50) + + expect(socket.closeCalls).toBe(1) + }) + + it('abandon closes the socket', async () => { + vi.stubGlobal('WebSocket', ControllableUpstream) + ControllableUpstream.reset() + + const promise = connectUpstreamWebSocket('ws://upstream') + const socket = ControllableUpstream.last() + socket.open() + const pending = await promise + + pending.abandon() + + expect(socket.closeCalls).toBe(1) + }) + + it('rejects when the WebSocket constructor throws', async () => { + vi.stubGlobal('WebSocket', class { + constructor() { + throw new Error('constructor failed') + } + }) + + await expect(connectUpstreamWebSocket('ws://upstream')).rejects.toThrow('constructor failed') + }) +}) + +describe('createBridgedSocketEvents', () => { + it('opens the bridge and forwards messages after the connect resolves', async () => { + const bridge = createBridge() + const events = createBridgedSocketEvents(() => Promise.resolve(bridge)) + const ws = createWs() + + events.onOpen!({} as Event, asWsContext(ws)) + await events.onMessage!({ data: 'hello' } as MessageEvent, asWsContext(ws)) + + expect(bridge.sent).toEqual(['hello']) + }) + + it('closes the client with 1011 when the bridge factory throws synchronously', async () => { + const events = createBridgedSocketEvents(() => { + throw new Error('boom') + }) + const ws = createWs() + + events.onOpen!({} as Event, asWsContext(ws)) + await flushAsync() + + expect(ws.closeCalls).toEqual([{ code: 1011, reason: 'boom' }]) + }) + + it('closes the client with 1011 when the bridge factory rejects', async () => { + const events = createBridgedSocketEvents(() => Promise.reject(new Error('nope'))) + const ws = createWs() + + events.onOpen!({} as Event, asWsContext(ws)) + await flushAsync() + + expect(ws.closeCalls).toEqual([{ code: 1011, reason: 'nope' }]) + }) + + it('does not close the client when it already closed', async () => { + const events = createBridgedSocketEvents(() => { + throw new Error('late') + }) + const ws = createWs() + + events.onOpen!({} as Event, asWsContext(ws)) + events.onClose!({} as CloseEvent, asWsContext(ws)) + await flushAsync() + + expect(ws.closeCalls).toEqual([]) + }) + + it('closes a bridge that resolves after the client closed', async () => { + let resolveBridge!: (bridge: WebSocketBridge) => void + const events = createBridgedSocketEvents( + () => new Promise((resolve) => { + resolveBridge = resolve + }), + ) + const ws = createWs() + + events.onOpen!({} as Event, asWsContext(ws)) + events.onClose!({} as CloseEvent, asWsContext(ws)) + + const bridge = createBridge() + resolveBridge(bridge) + await flushAsync() + + expect(bridge.closeCalls).toBe(1) + }) + + it('waits for the connect before forwarding messages', async () => { + let resolveBridge!: (bridge: WebSocketBridge) => void + const events = createBridgedSocketEvents( + () => new Promise((resolve) => { + resolveBridge = resolve + }), + ) + const ws = createWs() + + events.onOpen!({} as Event, asWsContext(ws)) + const messagePromise = events.onMessage!({ data: 'queued' } as MessageEvent, asWsContext(ws)) + + const bridge = createBridge() + resolveBridge(bridge) + await messagePromise + + expect(bridge.sent).toEqual(['queued']) + }) +}) + +describe('toWebSocketPeer', () => { + it('forwards send, close and bufferedAmount to the context', () => { + const ws = createWs() + ws.raw.bufferedAmount = 42 + const peer = toWebSocketPeer(asWsContext(ws)) + + peer.send('data') + peer.send(new Uint8Array([1])) + peer.close(1000, 'done') + + expect(ws.sent).toEqual(['data', new Uint8Array([1])]) + expect(ws.closeCalls).toEqual([{ code: 1000, reason: 'done' }]) + expect(peer.bufferedAmount()).toBe(42) + }) +}) + +describe('isAllowedUpgradeOrigin', () => { + it('allows a missing origin, rejects null and otherwise checks the list', () => { + expect(isAllowedUpgradeOrigin(undefined, ['http://localhost:5003'])).toBe(true) + expect(isAllowedUpgradeOrigin('null', ['http://localhost:5003'])).toBe(false) + expect(isAllowedUpgradeOrigin('http://localhost:5003', ['http://localhost:5003'])).toBe(true) + expect(isAllowedUpgradeOrigin('http://localhost:5004', ['http://localhost:5003'])).toBe(false) + }) }) describe('forwardPeerMessage', () => { - function createBridge(sent: Array): WebSocketBridge { + function createSentBridge(sent: Array): WebSocketBridge { return { send(data) { sent.push(data) @@ -174,7 +625,7 @@ describe('forwardPeerMessage', () => { it('forwards strings, ArrayBuffers and Blobs to the bridge', async () => { const sent: Array = [] - const bridge = createBridge(sent) + const bridge = createSentBridge(sent) await forwardPeerMessage(bridge, 'plain text') await forwardPeerMessage(bridge, new Uint8Array([1, 2]).buffer)