diff --git a/.changeset/fix-websocket-routing.md b/.changeset/fix-websocket-routing.md new file mode 100644 index 00000000..7ed53571 --- /dev/null +++ b/.changeset/fix-websocket-routing.md @@ -0,0 +1,5 @@ +--- +"@cloudflare/sandbox": patch +--- + +Fix WebSocket upgrade requests through exposed ports diff --git a/package-lock.json b/package-lock.json index cfde2721..b65ecb03 100644 --- a/package-lock.json +++ b/package-lock.json @@ -25,6 +25,7 @@ "@types/node": "^24.1.0", "@types/react": "^19.1.8", "@types/react-dom": "^19.1.6", + "@types/ws": "^8.18.1", "@vitejs/plugin-react": "^4.7.0", "@vitest/ui": "^3.2.4", "fast-glob": "^3.3.3", @@ -38,7 +39,8 @@ "typescript": "^5.8.3", "vite": "^7.1.11", "vitest": "^3.2.4", - "wrangler": "^4.42.2" + "wrangler": "^4.42.2", + "ws": "^8.18.3" } }, "examples/basic": { @@ -1458,6 +1460,28 @@ "wrangler": "^4.42.2" } }, + "node_modules/@cloudflare/vite-plugin/node_modules/ws": { + "version": "8.18.0", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.18.0.tgz", + "integrity": "sha512-8VbfWfHLbbwu3+N6OKsOMpBdT4kXPDDB9cJk2bJ6mh9ucxdlnNvH1e+roYkKmN9Nxw2yjz7VzeO9oOz2zJ04Pw==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=10.0.0" + }, + "peerDependencies": { + "bufferutil": "^4.0.1", + "utf-8-validate": ">=5.0.2" + }, + "peerDependenciesMeta": { + "bufferutil": { + "optional": true + }, + "utf-8-validate": { + "optional": true + } + } + }, "node_modules/@cloudflare/vitest-pool-workers": { "version": "0.9.12", "resolved": "https://registry.npmjs.org/@cloudflare/vitest-pool-workers/-/vitest-pool-workers-0.9.12.tgz", @@ -3503,6 +3527,16 @@ "dev": true, "license": "MIT" }, + "node_modules/@types/ws": { + "version": "8.18.1", + "resolved": "https://registry.npmjs.org/@types/ws/-/ws-8.18.1.tgz", + "integrity": "sha512-ThVF6DCVhA8kUGy+aazFQ4kXQ7E1Ty7A3ypFOe0IcJV8O/M511G99AW24irKrW56Wt44yG9+ij8FaqoBGkuBXg==", + "dev": true, + "license": "MIT", + "dependencies": { + "@types/node": "*" + } + }, "node_modules/@ungap/structured-clone": { "version": "1.3.0", "resolved": "https://registry.npmjs.org/@ungap/structured-clone/-/structured-clone-1.3.0.tgz", @@ -6541,6 +6575,28 @@ "node": ">=18.0.0" } }, + "node_modules/miniflare/node_modules/ws": { + "version": "8.18.0", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.18.0.tgz", + "integrity": "sha512-8VbfWfHLbbwu3+N6OKsOMpBdT4kXPDDB9cJk2bJ6mh9ucxdlnNvH1e+roYkKmN9Nxw2yjz7VzeO9oOz2zJ04Pw==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=10.0.0" + }, + "peerDependencies": { + "bufferutil": "^4.0.1", + "utf-8-validate": ">=5.0.2" + }, + "peerDependenciesMeta": { + "bufferutil": { + "optional": true + }, + "utf-8-validate": { + "optional": true + } + } + }, "node_modules/miniflare/node_modules/zod": { "version": "3.22.3", "resolved": "https://registry.npmjs.org/zod/-/zod-3.22.3.tgz", @@ -9850,9 +9906,9 @@ "license": "ISC" }, "node_modules/ws": { - "version": "8.18.0", - "resolved": "https://registry.npmjs.org/ws/-/ws-8.18.0.tgz", - "integrity": "sha512-8VbfWfHLbbwu3+N6OKsOMpBdT4kXPDDB9cJk2bJ6mh9ucxdlnNvH1e+roYkKmN9Nxw2yjz7VzeO9oOz2zJ04Pw==", + "version": "8.18.3", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.18.3.tgz", + "integrity": "sha512-PEIGCY5tSlUt50cqyMXfCzX+oOPqN0vuGqWzbcJ2xvnkzkq46oOpz7dQaTDBdfICb4N14+GARUDw2XV2N4tvzg==", "devOptional": true, "license": "MIT", "engines": { diff --git a/package.json b/package.json index 091ad904..fb42707f 100644 --- a/package.json +++ b/package.json @@ -33,20 +33,22 @@ "@types/node": "^24.1.0", "@types/react": "^19.1.8", "@types/react-dom": "^19.1.6", + "@types/ws": "^8.18.1", "@vitejs/plugin-react": "^4.7.0", "@vitest/ui": "^3.2.4", "fast-glob": "^3.3.3", "happy-dom": "^20.0.0", + "pkg-pr-new": "^0.0.60", "react": "^19.1.0", "react-dom": "^19.1.0", "tsup": "^8.5.0", "tsx": "^4.20.3", "turbo": "^2.5.8", "typescript": "^5.8.3", - "pkg-pr-new": "^0.0.60", "vite": "^7.1.11", "vitest": "^3.2.4", - "wrangler": "^4.42.2" + "wrangler": "^4.42.2", + "ws": "^8.18.3" }, "private": true, "packageManager": "npm@11.5.1" diff --git a/packages/sandbox/package.json b/packages/sandbox/package.json index 8f0574a8..e6f4a08f 100644 --- a/packages/sandbox/package.json +++ b/packages/sandbox/package.json @@ -27,7 +27,7 @@ "docker:local": "cd ../.. && docker build -f packages/sandbox/Dockerfile --build-arg SANDBOX_VERSION=$npm_package_version -t cloudflare/sandbox-test:$npm_package_version .", "docker:publish": "cd ../.. && docker buildx build --platform linux/amd64,linux/arm64 -f packages/sandbox/Dockerfile --build-arg SANDBOX_VERSION=$npm_package_version -t cloudflare/sandbox:$npm_package_version --push .", "docker:publish:beta": "cd ../.. && docker buildx build --platform linux/amd64,linux/arm64 -f packages/sandbox/Dockerfile --build-arg SANDBOX_VERSION=$npm_package_version -t cloudflare/sandbox:$npm_package_version-beta --push .", - "test": "vitest run --config vitest.config.ts", + "test": "vitest run --config vitest.config.ts \"$@\"", "test:e2e": "cd ../.. && vitest run --config vitest.e2e.config.ts \"$@\"" }, "exports": { diff --git a/packages/sandbox/src/request-handler.ts b/packages/sandbox/src/request-handler.ts index ad2a4501..9b7a1765 100644 --- a/packages/sandbox/src/request-handler.ts +++ b/packages/sandbox/src/request-handler.ts @@ -1,4 +1,5 @@ import { createLogger, type LogContext, TraceContext } from "@repo/shared"; +import { switchPort } from "@cloudflare/containers"; import { getSandbox, type Sandbox } from "./sandbox"; import { sanitizeSandboxId, @@ -70,6 +71,14 @@ export async function proxyToSandbox( } } + // Detect WebSocket upgrade request + const upgradeHeader = request.headers.get('Upgrade'); + if (upgradeHeader?.toLowerCase() === 'websocket') { + // WebSocket path: Must use fetch() not containerFetch() + // This bypasses JSRPC serialization boundary which cannot handle WebSocket upgrades + return await sandbox.fetch(switchPort(request, port)); + } + // Build proxy request with proper headers let proxyUrl: string; @@ -96,7 +105,7 @@ export async function proxyToSandbox( duplex: 'half', }); - return sandbox.containerFetch(proxyRequest, port); + return await sandbox.containerFetch(proxyRequest, port); } catch (error) { logger.error('Proxy routing error', error instanceof Error ? error : new Error(String(error))); return new Response('Proxy routing error', { status: 500 }); diff --git a/packages/sandbox/src/sandbox.ts b/packages/sandbox/src/sandbox.ts index 22b3c432..8ea44296 100644 --- a/packages/sandbox/src/sandbox.ts +++ b/packages/sandbox/src/sandbox.ts @@ -238,7 +238,17 @@ export class Sandbox extends Container implements ISandbox { await this.ctx.storage.put('sandboxName', name); } - // Determine which port to route to + // Detect WebSocket upgrade request + const upgradeHeader = request.headers.get('Upgrade'); + const isWebSocket = upgradeHeader?.toLowerCase() === 'websocket'; + + if (isWebSocket) { + // WebSocket path: Let parent Container class handle WebSocket proxying + // This bypasses containerFetch() which uses JSRPC and cannot handle WebSocket upgrades + return await super.fetch(request); + } + + // Non-WebSocket: Use existing port determination and HTTP routing logic const port = this.determinePort(url); // Route to the appropriate port diff --git a/packages/sandbox/tests/request-handler.test.ts b/packages/sandbox/tests/request-handler.test.ts new file mode 100644 index 00000000..f172c613 --- /dev/null +++ b/packages/sandbox/tests/request-handler.test.ts @@ -0,0 +1,240 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { proxyToSandbox, type SandboxEnv } from '../src/request-handler'; +import type { Sandbox } from '../src/sandbox'; + +// Mock getSandbox from sandbox.ts +vi.mock('../src/sandbox', () => { + const mockFn = vi.fn(); + return { + getSandbox: mockFn, + Sandbox: vi.fn(), + }; +}); + +// Import the mock after vi.mock is set up +import { getSandbox } from '../src/sandbox'; + +describe('proxyToSandbox - WebSocket Support', () => { + let mockSandbox: Partial; + let mockEnv: SandboxEnv; + + beforeEach(() => { + vi.clearAllMocks(); + + // Mock Sandbox with necessary methods + mockSandbox = { + validatePortToken: vi.fn().mockResolvedValue(true), + fetch: vi.fn().mockResolvedValue(new Response('WebSocket response')), + containerFetch: vi.fn().mockResolvedValue(new Response('HTTP response')), + }; + + mockEnv = { + Sandbox: {} as any, + }; + + vi.mocked(getSandbox).mockReturnValue(mockSandbox as Sandbox); + }); + + describe('WebSocket detection and routing', () => { + it('should detect WebSocket upgrade header (case-insensitive)', async () => { + const request = new Request('https://8080-sandbox-token12345678901.example.com/ws', { + headers: { + 'Upgrade': 'websocket', + 'Connection': 'Upgrade', + }, + }); + + await proxyToSandbox(request, mockEnv); + + // Should route through fetch() for WebSocket + expect(mockSandbox.fetch).toHaveBeenCalledTimes(1); + expect(mockSandbox.containerFetch).not.toHaveBeenCalled(); + }); + + it('should set cf-container-target-port header for WebSocket', async () => { + const request = new Request('https://8080-sandbox-token12345678901.example.com/ws', { + headers: { + 'Upgrade': 'websocket', + }, + }); + + await proxyToSandbox(request, mockEnv); + + expect(mockSandbox.fetch).toHaveBeenCalledTimes(1); + const fetchCall = vi.mocked(mockSandbox.fetch as any).mock.calls[0][0] as Request; + expect(fetchCall.headers.get('cf-container-target-port')).toBe('8080'); + }); + + it('should preserve original headers for WebSocket', async () => { + const request = new Request('https://8080-sandbox-token12345678901.example.com/ws', { + headers: { + 'Upgrade': 'websocket', + 'Sec-WebSocket-Key': 'test-key-123', + 'Sec-WebSocket-Version': '13', + 'User-Agent': 'test-client', + }, + }); + + await proxyToSandbox(request, mockEnv); + + const fetchCall = vi.mocked(mockSandbox.fetch as any).mock.calls[0][0] as Request; + expect(fetchCall.headers.get('Upgrade')).toBe('websocket'); + expect(fetchCall.headers.get('Sec-WebSocket-Key')).toBe('test-key-123'); + expect(fetchCall.headers.get('Sec-WebSocket-Version')).toBe('13'); + expect(fetchCall.headers.get('User-Agent')).toBe('test-client'); + }); + }); + + describe('HTTP routing (existing behavior)', () => { + it('should route HTTP requests through containerFetch', async () => { + const request = new Request('https://8080-sandbox-token12345678901.example.com/api/data', { + method: 'GET', + }); + + await proxyToSandbox(request, mockEnv); + + // Should route through containerFetch() for HTTP + expect(mockSandbox.containerFetch).toHaveBeenCalledTimes(1); + expect(mockSandbox.fetch).not.toHaveBeenCalled(); + }); + + it('should route POST requests through containerFetch', async () => { + const request = new Request('https://8080-sandbox-token12345678901.example.com/api/data', { + method: 'POST', + body: JSON.stringify({ data: 'test' }), + headers: { + 'Content-Type': 'application/json', + }, + }); + + await proxyToSandbox(request, mockEnv); + + expect(mockSandbox.containerFetch).toHaveBeenCalledTimes(1); + expect(mockSandbox.fetch).not.toHaveBeenCalled(); + }); + + it('should not detect SSE as WebSocket', async () => { + const request = new Request('https://8080-sandbox-token12345678901.example.com/events', { + headers: { + 'Accept': 'text/event-stream', + }, + }); + + await proxyToSandbox(request, mockEnv); + + // SSE should use HTTP path, not WebSocket path + expect(mockSandbox.containerFetch).toHaveBeenCalledTimes(1); + expect(mockSandbox.fetch).not.toHaveBeenCalled(); + }); + }); + + describe('Token validation', () => { + it('should validate token for both WebSocket and HTTP requests', async () => { + const wsRequest = new Request('https://8080-sandbox-token12345678901.example.com/ws', { + headers: { 'Upgrade': 'websocket' }, + }); + + await proxyToSandbox(wsRequest, mockEnv); + expect(mockSandbox.validatePortToken).toHaveBeenCalledWith(8080, 'token12345678901'); + + vi.clearAllMocks(); + + const httpRequest = new Request('https://8080-sandbox-token12345678901.example.com/api'); + await proxyToSandbox(httpRequest, mockEnv); + expect(mockSandbox.validatePortToken).toHaveBeenCalledWith(8080, 'token12345678901'); + }); + + it('should reject requests with invalid token', async () => { + vi.mocked(mockSandbox.validatePortToken as any).mockResolvedValue(false); + + const request = new Request('https://8080-sandbox-invalidtoken1234.example.com/ws', { + headers: { 'Upgrade': 'websocket' }, + }); + + const response = await proxyToSandbox(request, mockEnv); + + expect(response?.status).toBe(404); + expect(mockSandbox.fetch).not.toHaveBeenCalled(); + + const body = await response?.json(); + expect(body).toMatchObject({ + error: 'Access denied: Invalid token or port not exposed', + code: 'INVALID_TOKEN', + }); + }); + + it('should reject reserved port 3000', async () => { + // Port 3000 is reserved as control plane port and rejected by validatePort() + const request = new Request('https://3000-sandbox-anytoken12345678.example.com/status', { + method: 'GET', + }); + + const response = await proxyToSandbox(request, mockEnv); + + // Port 3000 is reserved and should be rejected (extractSandboxRoute returns null) + expect(response).toBeNull(); + expect(mockSandbox.validatePortToken).not.toHaveBeenCalled(); + expect(mockSandbox.containerFetch).not.toHaveBeenCalled(); + }); + }); + + describe('Port routing', () => { + it('should route to correct port from subdomain', async () => { + const request = new Request('https://9000-sandbox-token12345678901.example.com/api', { + method: 'GET', + }); + + await proxyToSandbox(request, mockEnv); + + expect(mockSandbox.validatePortToken).toHaveBeenCalledWith(9000, 'token12345678901'); + }); + }); + + describe('Non-sandbox requests', () => { + it('should return null for non-sandbox URLs', async () => { + const request = new Request('https://example.com/some-path'); + + const response = await proxyToSandbox(request, mockEnv); + + expect(response).toBeNull(); + expect(mockSandbox.fetch).not.toHaveBeenCalled(); + expect(mockSandbox.containerFetch).not.toHaveBeenCalled(); + }); + + it('should return null for invalid subdomain patterns', async () => { + const request = new Request('https://invalid-pattern.example.com'); + + const response = await proxyToSandbox(request, mockEnv); + + expect(response).toBeNull(); + }); + }); + + describe('Error handling', () => { + it('should handle errors during WebSocket routing', async () => { + (mockSandbox.fetch as any).mockImplementation(() => Promise.reject(new Error('Connection failed'))); + + const request = new Request('https://8080-sandbox-token12345678901.example.com/ws', { + headers: { + 'Upgrade': 'websocket', + }, + }); + + const response = await proxyToSandbox(request, mockEnv); + + expect(response?.status).toBe(500); + const text = await response?.text(); + expect(text).toBe('Proxy routing error'); + }); + + it('should handle errors during HTTP routing', async () => { + (mockSandbox.containerFetch as any).mockImplementation(() => Promise.reject(new Error('Service error'))); + + const request = new Request('https://8080-sandbox-token12345678901.example.com/api'); + + const response = await proxyToSandbox(request, mockEnv); + + expect(response?.status).toBe(500); + }); + }); +}); diff --git a/packages/sandbox/tests/sandbox.test.ts b/packages/sandbox/tests/sandbox.test.ts index 479a40ec..c6881cb0 100644 --- a/packages/sandbox/tests/sandbox.test.ts +++ b/packages/sandbox/tests/sandbox.test.ts @@ -1,23 +1,36 @@ import type { DurableObjectState } from '@cloudflare/workers-types'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { Sandbox } from '../src/sandbox'; +import { Container } from '@cloudflare/containers'; // Mock dependencies before imports vi.mock('./interpreter', () => ({ CodeInterpreter: vi.fn().mockImplementation(() => ({})), })); -vi.mock('@cloudflare/containers', () => ({ - Container: class Container { +vi.mock('@cloudflare/containers', () => { + const MockContainer = class Container { ctx: any; env: any; constructor(ctx: any, env: any) { this.ctx = ctx; this.env = env; } - }, - getContainer: vi.fn(), -})); + async fetch(request: Request): Promise { + // Mock implementation - will be spied on in tests + return new Response('Mock Container fetch'); + } + async containerFetch(request: Request, port: number): Promise { + // Mock implementation for HTTP path + return new Response('Mock Container HTTP fetch'); + } + }; + + return { + Container: MockContainer, + getContainer: vi.fn(), + }; +}); describe('Sandbox - Automatic Session Management', () => { let sandbox: Sandbox; @@ -462,4 +475,80 @@ describe('Sandbox - Automatic Session Management', () => { expect(sandbox.client.ports.exposePort).toHaveBeenCalled(); }); }); + + describe('fetch() override - WebSocket detection', () => { + let superFetchSpy: any; + + beforeEach(async () => { + await sandbox.setSandboxName('test-sandbox'); + + // Spy on Container.prototype.fetch to verify WebSocket routing + superFetchSpy = vi.spyOn(Container.prototype, 'fetch') + .mockResolvedValue(new Response('WebSocket response')); + }); + + afterEach(() => { + superFetchSpy?.mockRestore(); + }); + + it('should detect WebSocket upgrade header and route to super.fetch', async () => { + const request = new Request('https://example.com/ws', { + headers: { + 'Upgrade': 'websocket', + 'Connection': 'Upgrade', + }, + }); + + const response = await sandbox.fetch(request); + + // Should route through super.fetch() for WebSocket + expect(superFetchSpy).toHaveBeenCalledTimes(1); + expect(await response.text()).toBe('WebSocket response'); + }); + + it('should route non-WebSocket requests through containerFetch', async () => { + // GET request + const getRequest = new Request('https://example.com/api/data'); + await sandbox.fetch(getRequest); + expect(superFetchSpy).not.toHaveBeenCalled(); + + vi.clearAllMocks(); + + // POST request + const postRequest = new Request('https://example.com/api/data', { + method: 'POST', + body: JSON.stringify({ data: 'test' }), + headers: { 'Content-Type': 'application/json' }, + }); + await sandbox.fetch(postRequest); + expect(superFetchSpy).not.toHaveBeenCalled(); + + vi.clearAllMocks(); + + // SSE request (should not be detected as WebSocket) + const sseRequest = new Request('https://example.com/events', { + headers: { 'Accept': 'text/event-stream' }, + }); + await sandbox.fetch(sseRequest); + expect(superFetchSpy).not.toHaveBeenCalled(); + }); + + it('should preserve WebSocket request unchanged when calling super.fetch()', async () => { + const request = new Request('https://example.com/ws', { + headers: { + 'Upgrade': 'websocket', + 'Sec-WebSocket-Key': 'test-key-123', + 'Sec-WebSocket-Version': '13', + }, + }); + + await sandbox.fetch(request); + + expect(superFetchSpy).toHaveBeenCalledTimes(1); + const passedRequest = superFetchSpy.mock.calls[0][0] as Request; + expect(passedRequest.headers.get('Upgrade')).toBe('websocket'); + expect(passedRequest.headers.get('Sec-WebSocket-Key')).toBe('test-key-123'); + expect(passedRequest.headers.get('Sec-WebSocket-Version')).toBe('13'); + }); + }); }); diff --git a/tests/e2e/fixtures/websocket-echo-server.ts b/tests/e2e/fixtures/websocket-echo-server.ts new file mode 100644 index 00000000..12bc5bd9 --- /dev/null +++ b/tests/e2e/fixtures/websocket-echo-server.ts @@ -0,0 +1,35 @@ +/** + * Simple WebSocket Echo Server for E2E Testing + * + * This server echoes back any messages it receives. + * Used to validate WebSocket routing through the sandbox infrastructure. + * + * Usage: bun run websocket-echo-server.ts + */ + +const port = parseInt(process.argv[2] || '8080', 10); + +Bun.serve({ + port, + fetch(req, server) { + // Upgrade HTTP request to WebSocket + if (server.upgrade(req)) { + return; // Successfully upgraded + } + return new Response('Expected WebSocket', { status: 400 }); + }, + websocket: { + message(ws, message) { + // Echo the message back + ws.send(message); + }, + open(ws) { + console.log('WebSocket client connected'); + }, + close(ws) { + console.log('WebSocket client disconnected'); + }, + }, +}); + +console.log(`WebSocket echo server listening on port ${port}`); diff --git a/tests/e2e/websocket-workflow.test.ts b/tests/e2e/websocket-workflow.test.ts new file mode 100644 index 00000000..5796f7f4 --- /dev/null +++ b/tests/e2e/websocket-workflow.test.ts @@ -0,0 +1,163 @@ +import { describe, test, expect, beforeAll, afterAll, afterEach, vi } from 'vitest'; +import { readFileSync } from 'node:fs'; +import { join } from 'node:path'; +import WebSocket from 'ws'; +import { getTestWorkerUrl, WranglerDevRunner } from './helpers/wrangler-runner'; +import { createSandboxId, createTestHeaders, fetchWithStartup, cleanupSandbox } from './helpers/test-fixtures'; + +// Port exposure tests require custom domain with wildcard DNS routing +// Skip these tests when running against workers.dev deployment (no wildcard support) +const skipWebSocketTests = process.env.TEST_WORKER_URL?.endsWith('.workers.dev') ?? false; + +/** + * WebSocket Workflow Integration Tests + * + * Tests WebSocket support for exposed sandbox ports. + * + * SCOPE: Phase 1 - WebSocket Routing Validation + * This test validates that WebSocket upgrade requests are correctly routed + * through the sandbox infrastructure (Worker → proxyToSandbox → Sandbox.fetch → Container). + * + * NOT TESTED HERE: + * - Long-running connection timeout management (Phase 2) + * - Keep-alive strategies (Phase 2) + * - Error recovery and reconnection (Phase 3) + * - Concurrent connections or load testing (Phase 4) + * + * KNOWN LIMITATION: + * Current runtime has 30-second CPU timeout without keep-alive mechanism. + * This test keeps connections brief (< 5s) to validate routing correctness only. + * Timeout management will be addressed in Phase 2. + */ +describe('WebSocket Workflow', () => { + describe.skipIf(skipWebSocketTests)('local', () => { + let runner: WranglerDevRunner | null = null; + let workerUrl: string; + let currentSandboxId: string | null = null; + + beforeAll(async () => { + const result = await getTestWorkerUrl(); + workerUrl = result.url; + runner = result.runner; + }); + + afterEach(async () => { + // Cleanup sandbox container after each test + if (currentSandboxId) { + await cleanupSandbox(workerUrl, currentSandboxId); + currentSandboxId = null; + } + }); + + afterAll(async () => { + if (runner) { + await runner.stop(); + } + }); + + test('should connect to WebSocket server via exposed port and echo messages', async () => { + currentSandboxId = createSandboxId(); + const headers = createTestHeaders(currentSandboxId); + + // Read the WebSocket echo server fixture + const serverCode = readFileSync( + join(__dirname, 'fixtures', 'websocket-echo-server.ts'), + 'utf-8' + ); + + // Step 1: Write the WebSocket echo server to the container + await vi.waitFor( + async () => fetchWithStartup(`${workerUrl}/api/file/write`, { + method: 'POST', + headers, + body: JSON.stringify({ + path: '/workspace/ws-server.ts', + content: serverCode, + }), + }), + { timeout: 90000, interval: 2000 } + ); + + // Step 2: Start the WebSocket server as a background process + const port = 8080; + const startResponse = await fetch(`${workerUrl}/api/process/start`, { + method: 'POST', + headers, + body: JSON.stringify({ + command: `bun run /workspace/ws-server.ts ${port}`, + }), + }); + + expect(startResponse.status).toBe(200); + const processData = await startResponse.json(); + const processId = processData.id; + expect(processData.status).toBe('running'); + + // Wait for server to be ready (generous timeout for first startup) + await new Promise(resolve => setTimeout(resolve, 2000)); + + // Step 3: Expose the port to get preview URL + const exposeResponse = await fetch(`${workerUrl}/api/port/expose`, { + method: 'POST', + headers, + body: JSON.stringify({ + port, + name: 'websocket-test', + }), + }); + + expect(exposeResponse.status).toBe(200); + const exposeData = await exposeResponse.json(); + expect(exposeData.url).toBeTruthy(); + console.log('[DEBUG] Preview URL:', exposeData.url); + + // Step 4: Connect to WebSocket via preview URL + // Convert http:// to ws:// for WebSocket protocol + const wsUrl = exposeData.url.replace(/^http/, 'ws'); + + const ws = new WebSocket(wsUrl); + + // Wait for connection to open + await new Promise((resolve, reject) => { + ws.on('open', () => resolve()); + ws.on('error', (error) => reject(error)); + setTimeout(() => reject(new Error('WebSocket connection timeout')), 10000); + }); + + console.log('[DEBUG] WebSocket connected'); + + // Step 5: Send a message and verify echo + const testMessage = 'Hello WebSocket!'; + const messagePromise = new Promise((resolve, reject) => { + ws.on('message', (data) => { + resolve(data.toString()); + }); + setTimeout(() => reject(new Error('Echo timeout')), 5000); + }); + + ws.send(testMessage); + const echoedMessage = await messagePromise; + + expect(echoedMessage).toBe(testMessage); + console.log('[DEBUG] Message echoed successfully:', echoedMessage); + + // Step 6: Close WebSocket connection gracefully + ws.close(); + await new Promise((resolve) => { + ws.on('close', () => resolve()); + setTimeout(() => resolve(), 1000); // Fallback timeout + }); + + // Step 7: Cleanup - kill process and unexpose port + await fetch(`${workerUrl}/api/process/${processId}`, { + method: 'DELETE', + headers, + }); + + await fetch(`${workerUrl}/api/exposed-ports/${port}`, { + method: 'DELETE', + headers, + }); + }, 90000); + }); +});