From d910f6e67ffd21794e1209977953458a8dbee3af Mon Sep 17 00:00:00 2001 From: Gert Hengeveld Date: Mon, 23 Mar 2026 16:20:31 +0100 Subject: [PATCH] Fix WebSocket connection for StackBlitz/WebContainers --- code/core/package.json | 1 + code/core/src/common/utils/envs.ts | 2 + code/core/src/core-server/build-dev.ts | 2 + .../utils/__tests__/server-channel.test.ts | 61 +++++++++++++++++++ .../core-server/utils/get-server-channel.ts | 17 +++--- yarn.lock | 8 +++ 6 files changed, 84 insertions(+), 7 deletions(-) diff --git a/code/core/package.json b/code/core/package.json index bb195e3cbf40..ed2ed441dc68 100644 --- a/code/core/package.json +++ b/code/core/package.json @@ -234,6 +234,7 @@ "@testing-library/user-event": "^14.6.1", "@vitest/expect": "3.2.4", "@vitest/spy": "3.2.4", + "@webcontainer/env": "^1.1.1", "esbuild": "^0.18.0 || ^0.19.0 || ^0.20.0 || ^0.21.0 || ^0.22.0 || ^0.23.0 || ^0.24.0 || ^0.25.0 || ^0.26.0 || ^0.27.0", "open": "^10.2.0", "recast": "^0.23.5", diff --git a/code/core/src/common/utils/envs.ts b/code/core/src/common/utils/envs.ts index d62eeb6048be..c87005101d5c 100644 --- a/code/core/src/common/utils/envs.ts +++ b/code/core/src/common/utils/envs.ts @@ -1,3 +1,5 @@ +export { isWebContainer } from '@webcontainer/env'; + // eslint-disable-next-line @typescript-eslint/ban-ts-comment // @ts-ignore - Needed for Angular sandbox running without --no-link option. Do NOT convert to @ts-expect-error! import { nodePathsToArray } from './paths'; diff --git a/code/core/src/core-server/build-dev.ts b/code/core/src/core-server/build-dev.ts index bba61a8f6ddd..93233fb497f5 100644 --- a/code/core/src/core-server/build-dev.ts +++ b/code/core/src/core-server/build-dev.ts @@ -5,6 +5,7 @@ import { getConfigInfo, getInterpretedFile, getProjectRoot, + isWebContainer, loadAllPresets, loadMainConfig, resolveAddonName, @@ -236,6 +237,7 @@ export async function buildDevStandalone( token: getWsToken(), host: options.host, allowedHosts, + skipValidation: isWebContainer(), localAddress, networkAddress, }); diff --git a/code/core/src/core-server/utils/__tests__/server-channel.test.ts b/code/core/src/core-server/utils/__tests__/server-channel.test.ts index 9128df924339..3b1e29276615 100644 --- a/code/core/src/core-server/utils/__tests__/server-channel.test.ts +++ b/code/core/src/core-server/utils/__tests__/server-channel.test.ts @@ -16,6 +16,11 @@ const options = { token: mockToken, } as any; +const webContainerOptions = { + ...options, + skipValidation: true, +} as any; + describe('getServerChannel', () => { it('should return a channel', () => { const server = { on: vi.fn() } as any as Server; @@ -303,4 +308,60 @@ describe('ServerChannelTransport', () => { // Socket should not be destroyed for wrong path (just ignored) expect(destroySpy).not.toHaveBeenCalled(); }); + + it('accepts connections without token when validation is disabled', () => { + const server = new EventEmitter() as any as Server; + const socket = new EventEmitter() as any; + socket.write = vi.fn(); + socket.destroy = vi.fn(); + const destroySpy = vi.spyOn(socket, 'destroy'); + const handleUpgradeSpy = vi.fn(); + const transport = new ServerChannelTransport(server, webContainerOptions); + + // Mock handleUpgrade to track if it's called + // @ts-expect-error (accessing private property) + transport.socket.handleUpgrade = handleUpgradeSpy; + + const request = { + url: '/storybook-server-channel', + headers: { + origin: 'http://localhost:6006', + }, + } as any; + const head = Buffer.from(''); + + server.listeners('upgrade')[0](request, socket, head); + + expect(socket.write).not.toHaveBeenCalled(); + expect(destroySpy).not.toHaveBeenCalled(); + expect(handleUpgradeSpy).toHaveBeenCalled(); + }); + + it('accepts connections with invalid origin when validation is disabled', () => { + const server = new EventEmitter() as any as Server; + const socket = new EventEmitter() as any; + socket.write = vi.fn(); + socket.destroy = vi.fn(); + const destroySpy = vi.spyOn(socket, 'destroy'); + const handleUpgradeSpy = vi.fn(); + const transport = new ServerChannelTransport(server, webContainerOptions); + + // Mock handleUpgrade to track if it's called + // @ts-expect-error (accessing private property) + transport.socket.handleUpgrade = handleUpgradeSpy; + + const request = { + url: '/storybook-server-channel?token=wrong-token', + headers: { + origin: 'http://malicious-site.com', + }, + } as any; + const head = Buffer.from(''); + + server.listeners('upgrade')[0](request, socket, head); + + expect(socket.write).not.toHaveBeenCalled(); + expect(destroySpy).not.toHaveBeenCalled(); + expect(handleUpgradeSpy).toHaveBeenCalled(); + }); }); diff --git a/code/core/src/core-server/utils/get-server-channel.ts b/code/core/src/core-server/utils/get-server-channel.ts index 2050bf9d71e7..a338f5e2e298 100644 --- a/code/core/src/core-server/utils/get-server-channel.ts +++ b/code/core/src/core-server/utils/get-server-channel.ts @@ -14,6 +14,7 @@ import { isValidToken } from './validate-token'; type Server = NonNullable[0]>['server']>; type ServerChannelTransportOptions = HostValidationOptions & { + skipValidation?: boolean; token: string; }; @@ -38,14 +39,16 @@ export class ServerChannelTransport { return; } - const originHost = request.headers.origin && new URL(request.headers.origin).host; - if (!isValidHost(originHost, options)) { - throw new Error('Invalid websocket origin'); - } + if (!options.skipValidation) { + const originHost = request.headers.origin && new URL(request.headers.origin).host; + if (!isValidHost(originHost, options)) { + throw new Error('Invalid websocket origin'); + } - const requestToken = url.searchParams.get('token'); - if (!isValidToken(requestToken, options.token)) { - throw new Error('Invalid websocket token'); + const requestToken = url.searchParams.get('token'); + if (!isValidToken(requestToken, options.token)) { + throw new Error('Invalid websocket token'); + } } this.socket.handleUpgrade(request, socket, head, (ws) => { diff --git a/yarn.lock b/yarn.lock index f341af810e8d..14ad9807f4b9 100644 --- a/yarn.lock +++ b/yarn.lock @@ -11475,6 +11475,13 @@ __metadata: languageName: node linkType: hard +"@webcontainer/env@npm:^1.1.1": + version: 1.1.1 + resolution: "@webcontainer/env@npm:1.1.1" + checksum: 10c0/bc64114ffa7ee92f4985cc2bdd5e27f6f31d892b9aa5cde68eaf93df02d13ee6edf13faeebdd701464183b6f8f9c47c14975958cdd6fc20e7356ad32f6ee39e7 + languageName: node + linkType: hard + "@xtuc/ieee754@npm:^1.2.0": version: 1.2.0 resolution: "@xtuc/ieee754@npm:1.2.0" @@ -28581,6 +28588,7 @@ __metadata: "@vitest/mocker": "npm:3.2.4" "@vitest/spy": "npm:3.2.4" "@vitest/utils": "npm:^3.2.4" + "@webcontainer/env": "npm:^1.1.1" "@yarnpkg/fslib": "npm:2.10.3" "@yarnpkg/libzip": "npm:2.3.0" ansi-to-html: "npm:^0.7.2"