Core: Harden websocket connection

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 9.1.19
+
+- Harden websocket connection
+
 ## 9.1.18
 
 - No-op release. No changes.

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

@@ -33,11 +33,16 @@ export const pluckThirdPartyPackageFromPath = (packagePath: string) =>
 export const buildFrameworkGlobalsFromOptions = async (options: Options) => {
   const globals: Record<string, any> = {};
 
-  const { builder } = await options.presets.apply('core');
+  const { builder, channelOptions } = await options.presets.apply('core');
 
   const frameworkName = await getFrameworkName(options);
   const rendererName = await extractProperRendererNameFromFramework(frameworkName);
 
+  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 };
+  }
+
   if (rendererName) {
     globals.STORYBOOK_RENDERER =
       (await extractProperRendererNameFromFramework(frameworkName)) ?? undefined;

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';
 
@@ -26,9 +27,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)
   );
 
   let indexError: Error | undefined;

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';
 import { dirname, isAbsolute, join } from 'node:path';
@@ -181,7 +182,7 @@ export const experimental_serverAPI = (extension: Record<string, Function>, opti
   }
   return { ...extension, removeAddon };
 };
-
+const wsToken = randomUUID();
 /**
  * If for some reason this config is not applied, the reason is that likely there is an addon that
  * does `export core = () => ({ someConfig })`, instead of `export core = (existing) => ({
@@ -191,6 +192,10 @@ export const experimental_serverAPI = (extension: Record<string, Function>, opti
 export const core = async (existing: CoreConfig, options: Options): Promise<CoreConfig> => ({
   ...existing,
   disableTelemetry: options.disableTelemetry === true,
+  channelOptions: {
+    ...(existing?.channelOptions ?? {}),
+    ...(options.configType === 'DEVELOPMENT' ? { wsToken } : {}),
+  },
   enableCrashReports:
     options.enableCrashReports || optionalEnvToBoolean(process.env.STORYBOOK_ENABLE_CRASH_REPORTS),
 });
@@ -247,6 +252,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,71 @@ describe('ServerChannelTransport', () => {
       }
     `);
   });
+
+  it('skips telejson classes and functions in data', () => {
+    const server = new EventEmitter() as any as Server;
+    const socket = new EventEmitter();
+    const transport = new ServerChannelTransport(server, mockToken);
+    const handler = vi.fn();
+    transport.setHandler(handler);
+
+    // @ts-expect-error (an internal API)
+    transport.socket.emit('connection', socket);
+
+    const input = { a() {}, b: class {}, c: true, d: 3 };
+    socket.emit('message', stringify(input));
+
+    console.log(handler.mock.calls);
+
+    expect(handler.mock.calls[0][0].a).toBeUndefined();
+    expect(handler.mock.calls[0][0].b).toBeUndefined();
+  });
+
+  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/stories-json.ts

@@ -62,6 +62,11 @@ export function useStoriesJson({
       const generator = await initializedStoryIndexGenerator;
       const index = await generator.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,20 @@
+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 {
+    return a.length === b.length && timingSafeEqual(a, b);
+  } catch {
+    return false;
+  }
+}

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

@@ -26,7 +26,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;
   /**

code/package.json

@@ -283,5 +283,6 @@
         "Dependency Upgrades"
       ]
     ]
-  }
+  },
+  "deferredNextVersion": "9.1.19"
 }