Core: Fix WebSocket connection for StackBlitz/WebContainers

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",

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';

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,
   });

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();
+  });
 });

code/core/src/core-server/utils/get-server-channel.ts

@@ -14,6 +14,7 @@ import { isValidToken } from './validate-token';
 type Server = NonNullable<NonNullable<ConstructorParameters<typeof WebSocketServer>[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) => {

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"