Core: Require token for websocket connections

code/builders/builder-webpack5/src/preview/iframe-webpack.config.ts

@@ -1,4 +1,4 @@
-import { dirname, join, resolve } from 'node:path';
+import { join, resolve } from 'node:path';
 import { fileURLToPath } from 'node:url';
 
 import {

code/core/src/builder-manager/utils/framework.ts

@@ -7,9 +7,9 @@ import {
 import { type Options, SupportedBuilder } from 'storybook/internal/types';
 
 export const buildFrameworkGlobalsFromOptions = async (options: Options) => {
-  const globals: Record<string, string | undefined> = {};
+  const globals: Record<string, any> = {};
 
-  const builderConfig = (await options.presets.apply('core')).builder;
+  const { builder: builderConfig, channelOptions } = await options.presets.apply('core');
   const builderName = typeof builderConfig === 'string' ? builderConfig : builderConfig?.name;
   const builder = Object.values(SupportedBuilder).find((builder) => builderName?.includes(builder));
 
@@ -18,6 +18,10 @@ export const buildFrameworkGlobalsFromOptions = async (options: Options) => {
   const framework = frameworkPackages[frameworkPackageName];
   const renderer = frameworkToRenderer[framework];
 
+  if (options.configType === 'DEVELOPMENT') {
+    // Manager only needs the token currently, so we don't pass any other channel options.
+    globals.CHANNEL_OPTIONS = { wsToken: channelOptions?.wsToken };
+  }
   globals.STORYBOOK_BUILDER = builder;
   globals.STORYBOOK_FRAMEWORK = framework;
   globals.STORYBOOK_RENDERER = renderer;

code/core/src/channels/index.ts

@@ -7,7 +7,7 @@ import { PostMessageTransport } from './postmessage';
 import type { ChannelTransport, Config } from './types';
 import { WebsocketTransport } from './websocket';
 
-const { CONFIG_TYPE } = global;
+const { CHANNEL_OPTIONS, CONFIG_TYPE } = global;
 
 export * from './main';
 
@@ -35,7 +35,8 @@ export function createBrowserChannel({ page, extraTransports = [] }: Options): C
   if (CONFIG_TYPE === 'DEVELOPMENT') {
     const protocol = window.location.protocol === 'http:' ? 'ws' : 'wss';
     const { hostname, port } = window.location;
-    const channelUrl = `${protocol}://${hostname}:${port}/storybook-server-channel`;
+    const { wsToken } = CHANNEL_OPTIONS || {};
+    const channelUrl = `${protocol}://${hostname}:${port}/storybook-server-channel?token=${wsToken}`;
 
     transports.push(new WebsocketTransport({ url: channelUrl, onError: () => {}, page }));
   }

code/core/src/core-server/dev-server.ts

@@ -4,6 +4,7 @@ import { MissingBuilderError } from 'storybook/internal/server-errors';
 import type { Options } from 'storybook/internal/types';
 
 import compression from '@polka/compression';
+import assert from 'assert';
 import polka from 'polka';
 import invariant from 'tiny-invariant';
 
@@ -28,9 +29,11 @@ export async function storybookDevServer(options: Options) {
   const [server, core] = await Promise.all([getServer(options), options.presets.apply('core')]);
   const app = polka({ server });
 
+  assert(core?.channelOptions?.wsToken, 'wsToken is required for securing the server channel');
+
   const serverChannel = await options.presets.apply(
     'experimental_serverChannel',
-    getServerChannel(server)
+    getServerChannel(server, core.channelOptions.wsToken)
   );
 
   const workingDir = process.cwd();
@@ -52,8 +55,6 @@ export async function storybookDevServer(options: Options) {
     options.extendServer(server);
   }
 
-  // CORS middleware must be registered BEFORE route handlers to ensure all routes
-  // (including /index.json) receive proper CORS headers for Storybook Composition
   app.use(getAccessControlMiddleware(core?.crossOriginIsolated ?? false));
   app.use(getCachingMiddleware());
 

code/core/src/core-server/presets/common-preset.ts

@@ -1,3 +1,4 @@
+import { randomUUID } from 'node:crypto';
 import { existsSync } from 'node:fs';
 import { readFile } from 'node:fs/promises';
 
@@ -190,8 +191,13 @@ export const experimental_serverAPI = (extension: Record<string, Function>, opti
  * ...existing, someConfig })`, just overwriting everything and not merging with the existing
  * values.
  */
+const wsToken = randomUUID();
 export const core = async (existing: CoreConfig, options: Options): Promise<CoreConfig> => ({
   ...existing,
+  channelOptions: {
+    ...(existing?.channelOptions ?? {}),
+    ...(options.configType === 'DEVELOPMENT' ? { wsToken } : {}),
+  },
   disableTelemetry: options.disableTelemetry === true,
   enableCrashReports:
     options.enableCrashReports || optionalEnvToBoolean(process.env.STORYBOOK_ENABLE_CRASH_REPORTS),
@@ -257,6 +263,10 @@ export const managerHead = async (_: any, options: Options) => {
   return '';
 };
 
+export const channelToken = async (value: string | undefined) => {
+  return value;
+};
+
 export const experimental_serverChannel = async (
   channel: Channel,
   options: OptionsWithRequiredCache

code/core/src/core-server/utils/__tests__/server-channel.test.ts

@@ -11,22 +11,24 @@ import { ServerChannelTransport, getServerChannel } from '../get-server-channel'
 describe('getServerChannel', () => {
   it('should return a channel', () => {
     const server = { on: vi.fn() } as any as Server;
-    const result = getServerChannel(server);
+    const result = getServerChannel(server, 'test-token-123');
     expect(result).toBeInstanceOf(Channel);
   });
 
   it('should attach to the http server', () => {
     const server = { on: vi.fn() } as any as Server;
-    getServerChannel(server);
+    getServerChannel(server, 'test-token-123');
     expect(server.on).toHaveBeenCalledWith('upgrade', expect.any(Function));
   });
 });
 
 describe('ServerChannelTransport', () => {
+  const mockToken = 'test-token-123';
+
   it('parses simple JSON', () => {
     const server = new EventEmitter() as any as Server;
     const socket = new EventEmitter();
-    const transport = new ServerChannelTransport(server);
+    const transport = new ServerChannelTransport(server, mockToken);
     const handler = vi.fn();
     transport.setHandler(handler);
 
@@ -36,10 +38,11 @@ describe('ServerChannelTransport', () => {
 
     expect(handler).toHaveBeenCalledWith('hello');
   });
+
   it('parses object JSON', () => {
     const server = new EventEmitter() as any as Server;
     const socket = new EventEmitter();
-    const transport = new ServerChannelTransport(server);
+    const transport = new ServerChannelTransport(server, mockToken);
     const handler = vi.fn();
     transport.setHandler(handler);
 
@@ -49,10 +52,11 @@ describe('ServerChannelTransport', () => {
 
     expect(handler).toHaveBeenCalledWith({ type: 'hello' });
   });
+
   it('supports telejson cyclical data', () => {
     const server = new EventEmitter() as any as Server;
     const socket = new EventEmitter();
-    const transport = new ServerChannelTransport(server);
+    const transport = new ServerChannelTransport(server, mockToken);
     const handler = vi.fn();
     transport.setHandler(handler);
 
@@ -70,4 +74,52 @@ describe('ServerChannelTransport', () => {
       }
     `);
   });
+
+  it('rejects connections with invalid token', () => {
+    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');
+    new ServerChannelTransport(server, mockToken);
+
+    // Simulate upgrade request with wrong token
+    const request = {
+      url: '/storybook-server-channel?token=wrong-token',
+    } as any;
+    const head = Buffer.from('');
+
+    server.listeners('upgrade')[0](request, socket, head);
+
+    expect(socket.write).toHaveBeenCalledWith(
+      'HTTP/1.1 403 Forbidden\r\nConnection: close\r\n\r\n'
+    );
+    expect(destroySpy).toHaveBeenCalled();
+  });
+
+  it('accepts connections with valid token', () => {
+    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, mockToken);
+
+    // Mock handleUpgrade to track if it's called
+    // @ts-expect-error (accessing private property)
+    transport.socket.handleUpgrade = handleUpgradeSpy;
+
+    // Simulate upgrade request with correct token
+    const request = {
+      url: `/storybook-server-channel?token=${mockToken}`,
+    } 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

@@ -1,10 +1,13 @@
+import type { IncomingMessage } from 'node:http';
+
 import type { ChannelHandler } from 'storybook/internal/channels';
 import { Channel, HEARTBEAT_INTERVAL } from 'storybook/internal/channels';
 
 import { isJSON, parse, stringify } from 'telejson';
 import WebSocket, { WebSocketServer } from 'ws';
 
 import { UniversalStore } from '../../shared/universal-store';
+import { isValidToken } from './validate-websocket-token';
 
 type Server = NonNullable<NonNullable<ConstructorParameters<typeof WebSocketServer>[0]>['server']>;
 
@@ -19,14 +22,27 @@ export class ServerChannelTransport {
 
   private handler?: ChannelHandler;
 
-  constructor(server: Server) {
+  private token: string;
+
+  constructor(server: Server, token: string) {
+    this.token = token;
     this.socket = new WebSocketServer({ noServer: true });
 
-    server.on('upgrade', (request, socket, head) => {
-      if (request.url === '/storybook-server-channel') {
-        this.socket.handleUpgrade(request, socket, head, (ws) => {
-          this.socket.emit('connection', ws, request);
-        });
+    server.on('upgrade', (request: IncomingMessage, socket, head) => {
+      if (request.url) {
+        const url = new URL(request.url, 'http://localhost');
+        if (url.pathname === '/storybook-server-channel') {
+          const requestToken = url.searchParams.get('token');
+          if (!isValidToken(requestToken, this.token)) {
+            socket.write('HTTP/1.1 403 Forbidden\r\nConnection: close\r\n\r\n');
+            socket.destroy();
+            return;
+          }
+
+          this.socket.handleUpgrade(request, socket, head, (ws) => {
+            this.socket.emit('connection', ws, request);
+          });
+        }
       }
     });
     this.socket.on('connection', (wss) => {
@@ -68,8 +84,8 @@ export class ServerChannelTransport {
   }
 }
 
-export function getServerChannel(server: Server) {
-  const transports = [new ServerChannelTransport(server)];
+export function getServerChannel(server: Server, token: string) {
+  const transports = [new ServerChannelTransport(server, token)];
 
   const channel = new Channel({ transports, async: true });
 

code/core/src/core-server/utils/getAccessControlMiddleware.ts

@@ -2,13 +2,9 @@ import type { Middleware } from '../../types';
 
 export function getAccessControlMiddleware(crossOriginIsolated: boolean): Middleware {
   return (req, res, next) => {
-    res.setHeader('Access-Control-Allow-Origin', '*');
-    res.setHeader('Access-Control-Allow-Headers', 'Origin, X-Requested-With, Content-Type, Accept');
     // These headers are required to enable SharedArrayBuffer
     // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/SharedArrayBuffer
     if (crossOriginIsolated) {
-      // These headers are required to enable SharedArrayBuffer
-      // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/SharedArrayBuffer
       res.setHeader('Cross-Origin-Opener-Policy', 'same-origin');
       res.setHeader('Cross-Origin-Embedder-Policy', 'require-corp');
     }

code/core/src/core-server/utils/index-json.ts

@@ -58,6 +58,11 @@ export function registerIndexJsonRoute({
     try {
       const index = await (await storyIndexGeneratorPromise).getIndex();
       res.setHeader('Content-Type', 'application/json');
+      res.setHeader('Access-Control-Allow-Origin', '*');
+      res.setHeader(
+        'Access-Control-Allow-Headers',
+        'Origin, X-Requested-With, Content-Type, Accept'
+      );
       res.end(JSON.stringify(index));
     } catch (err) {
       res.statusCode = 500;

code/core/src/core-server/utils/validate-websocket-token.ts

@@ -0,0 +1,21 @@
+import { timingSafeEqual } from 'node:crypto';
+
+/**
+ * Validates a secret token using constant-time comparison to prevent timing attacks.
+ *
+ * @returns `true` if tokens match, `false` otherwise
+ */
+export function isValidToken(token: string | null, expectedToken: string): boolean {
+  if (!token || !expectedToken) {
+    return false;
+  }
+
+  const a = Buffer.from(token, 'utf8');
+  const b = Buffer.from(expectedToken, 'utf8');
+  try {
+    // TODO: Remove any types as soon as @types/node is updated
+    return a.length === b.length && timingSafeEqual(a as any, b as any);
+  } catch {
+    return false;
+  }
+}

code/core/src/types/modules/core-common.ts

@@ -32,7 +32,7 @@ export interface CoreConfig {
       };
   renderer?: RendererName;
   disableWebpackDefaults?: boolean;
-  channelOptions?: Partial<TelejsonOptions>;
+  channelOptions?: Partial<TelejsonOptions> & { wsToken?: string };
   /** Disables the generation of project.json, a file containing Storybook metadata */
   disableProjectJson?: boolean;
   /**

scripts/tasks/compile.ts

@@ -7,7 +7,7 @@ import { exec } from '../utils/exec';
 import { maxConcurrentTasks } from '../utils/maxConcurrentTasks';
 
 // The amount of VCPUs for the check task on CI is 4 (large resource)
-const amountOfVCPUs = 3;
+const amountOfVCPUs = 2;
 
 const parallel = `--parallel=${process.env.CI ? amountOfVCPUs - 1 : maxConcurrentTasks}`;