feat: direct ws connection fallback (#4474)

packages/core/src/client/hmr.ts

@@ -2,6 +2,7 @@ import type { NormalizedClientConfig } from '../types';
 
 const compilationId = RSBUILD_COMPILATION_NAME;
 const config: NormalizedClientConfig = RSBUILD_CLIENT_CONFIG;
+const resolvedConfig: NormalizedClientConfig = RSBUILD_RESOLVED_CLIENT_CONFIG;
 
 function formatURL({
   port,
@@ -214,10 +215,27 @@ function onClose() {
   setTimeout(connect, 1000 * 1.5 ** reconnectCount);
 }
 
+function onError() {
+  if (!config.port) {
+    console.error(
+      '[HMR] WebSocket connection error, attempting direct fallback',
+    );
+    removeListeners();
+    connection = null;
+    connect(true);
+  }
+}
+
 // Establishing a WebSocket connection with the server.
-function connect() {
+function connect(fallback = false) {
+  let cfg = config;
+  if (fallback) {
+    cfg = resolvedConfig;
+    console.info('[HMR] Using direct websocket fallback');
+  }
+
   const { location } = self;
-  const { host, port, path, protocol } = config;
+  const { host, port, path, protocol } = cfg;
   const socketUrl = formatURL({
     protocol: protocol || (location.protocol === 'https:' ? 'wss' : 'ws'),
     hostname: host || location.hostname,
@@ -231,13 +249,16 @@ function connect() {
   connection.addEventListener('close', onClose);
   // Handle messages from the server.
   connection.addEventListener('message', onMessage);
+  // Handle errors
+  connection.addEventListener('error', onError);
 }
 
 function removeListeners() {
   if (connection) {
     connection.removeEventListener('open', onOpen);
     connection.removeEventListener('close', onClose);
     connection.removeEventListener('message', onMessage);
+    connection.removeEventListener('error', onError);
   }
 }
 

packages/core/src/env.d.ts

@@ -4,6 +4,8 @@ declare const WEBPACK_HASH: string;
 
 declare const RSBUILD_CLIENT_CONFIG: ClientConfig;
 
+declare const RSBUILD_RESOLVED_CLIENT_CONFIG: ClientConfig;
+
 declare const RSBUILD_DEV_LIVE_RELOAD: boolean;
 
 declare const RSBUILD_COMPILATION_NAME: string;

packages/core/src/provider/createCompiler.ts

@@ -1,5 +1,5 @@
-import { rspack } from '@rspack/core';
 import type { StatsCompilation } from '@rspack/core';
+import { rspack } from '@rspack/core';
 import {
   color,
   formatStats,
@@ -10,7 +10,7 @@ import {
 } from '../helpers';
 import { registerDevHook } from '../hooks';
 import { logger } from '../logger';
-import type { DevConfig, Rspack } from '../types';
+import type { DevConfig, Rspack, ServerConfig } from '../types';
 import { type InitConfigsOptions, initConfigs } from './initConfigs';
 
 export async function createCompiler(options: InitConfigsOptions): Promise<{
@@ -162,4 +162,6 @@ export type DevMiddlewareOptions = {
 
   /** whether use Server Side Render */
   serverSideRender?: boolean;
+
+  serverConfig: ServerConfig;
 };

packages/core/src/server/compilerDevMiddleware.ts

@@ -111,7 +111,10 @@ export class CompilerDevMiddleware {
   public async init(): Promise<void> {
     // start compiling
     const devMiddleware = await getDevMiddleware(this.compiler);
-    this.middleware = this.setupDevMiddleware(devMiddleware, this.publicPaths);
+    this.middleware = await this.setupDevMiddleware(
+      devMiddleware,
+      this.publicPaths,
+    );
     await this.socketServer.prepare();
   }
 
@@ -149,14 +152,12 @@ export class CompilerDevMiddleware {
     });
   }
 
-  private setupDevMiddleware(
+  private async setupDevMiddleware(
     devMiddleware: CustomDevMiddleware,
     publicPaths: string[],
-  ): DevMiddlewareAPI {
-    const {
-      devConfig,
-      serverConfig: { headers, base },
-    } = this;
+  ): Promise<DevMiddlewareAPI> {
+    const { devConfig, serverConfig } = this;
+    const { headers, base } = serverConfig;
 
     const callbacks = {
       onInvalid: (compilationId?: string, fileName?: string | null) => {
@@ -181,7 +182,7 @@ export class CompilerDevMiddleware {
 
     const clientPaths = getClientPaths(devConfig);
 
-    const middleware = devMiddleware({
+    const middleware = await devMiddleware({
       headers,
       publicPath: '/',
       stats: false,
@@ -194,6 +195,7 @@ export class CompilerDevMiddleware {
       // weak is enough in dev
       // https://developer.mozilla.org/en-US/docs/Web/HTTP/Conditional_requests#weak_validation
       etag: 'weak',
+      serverConfig,
     });
 
     const assetPrefixes = publicPaths

packages/core/src/server/devMiddleware.ts

@@ -4,6 +4,7 @@ import { applyToCompiler } from '../helpers';
 import type { DevMiddlewareOptions } from '../provider/createCompiler';
 import type { DevConfig, NextFunction } from '../types';
 import { getCompilationId } from './helper';
+import { getResolvedClientConfig } from './hmrFallback';
 
 type ServerCallbacks = {
   onInvalid: (compilationId?: string, fileName?: string | null) => void;
@@ -62,11 +63,13 @@ function applyHMREntry({
   compiler,
   clientPaths,
   clientConfig = {},
+  resolvedClientConfig = {},
   liveReload = true,
 }: {
   compiler: Compiler;
   clientPaths: string[];
   clientConfig: DevConfig['client'];
+  resolvedClientConfig: DevConfig['client'];
   liveReload: DevConfig['liveReload'];
 }) {
   if (!isClientCompiler(compiler)) {
@@ -76,6 +79,7 @@ function applyHMREntry({
   new compiler.webpack.DefinePlugin({
     RSBUILD_COMPILATION_NAME: JSON.stringify(getCompilationId(compiler)),
     RSBUILD_CLIENT_CONFIG: JSON.stringify(clientConfig),
+    RSBUILD_RESOLVED_CLIENT_CONFIG: JSON.stringify(resolvedClientConfig),
     RSBUILD_DEV_LIVE_RELOAD: liveReload,
   }).apply(compiler);
 
@@ -102,24 +106,37 @@ export type DevMiddlewareAPI = Middleware & {
  * - Inject the HMR client path into page （the HMR client rsbuild/server already provide）.
  * - Notify server when compiler hooks are triggered.
  */
-export type DevMiddleware = (options: DevMiddlewareOptions) => DevMiddlewareAPI;
+export type DevMiddleware = (
+  options: DevMiddlewareOptions,
+) => Promise<DevMiddlewareAPI>;
 
 export const getDevMiddleware = async (
   multiCompiler: Compiler | MultiCompiler,
 ): Promise<NonNullable<DevMiddleware>> => {
   const { default: rsbuildDevMiddleware } = await import(
     '../../compiled/rsbuild-dev-middleware/index.js'
   );
-  return (options) => {
-    const { clientPaths, clientConfig, callbacks, liveReload, ...restOptions } =
-      options;
+  return async (options) => {
+    const {
+      clientPaths,
+      clientConfig,
+      callbacks,
+      liveReload,
+      serverConfig,
+      ...restOptions
+    } = options;
+    const resolvedClientConfig = await getResolvedClientConfig(
+      clientConfig,
+      serverConfig,
+    );
 
     const setupCompiler = (compiler: Compiler) => {
       if (clientPaths) {
         applyHMREntry({
           compiler,
           clientPaths,
           clientConfig,
+          resolvedClientConfig,
           liveReload,
         });
       }

packages/core/src/server/hmrFallback.ts

@@ -0,0 +1,71 @@
+import { promises as dns } from 'node:dns';
+import type { DevConfig, ServerConfig } from '../types/config';
+
+type Hostname = {
+  host?: string;
+  name: string;
+};
+
+const wildcardHosts = new Set([
+  '0.0.0.0',
+  '::',
+  '0000:0000:0000:0000:0000:0000:0000:0000',
+]);
+
+export async function resolveHostname(
+  optionsHost: string | undefined,
+): Promise<Hostname> {
+  let host: string | undefined;
+  if (optionsHost === undefined) {
+    // Use a secure default
+    host = 'localhost';
+  } else {
+    host = optionsHost;
+  }
+
+  // Set host name to localhost when possible
+  let name = host === undefined || wildcardHosts.has(host) ? 'localhost' : host;
+
+  if (host === 'localhost') {
+    const localhostAddr = await getLocalhostAddressIfDiffersFromDNS();
+    if (localhostAddr) {
+      name = localhostAddr;
+    }
+  }
+
+  return { host, name };
+}
+
+/**
+ * Returns resolved localhost address when `dns.lookup` result differs from DNS
+ *
+ * `dns.lookup` result is same when defaultResultOrder is `verbatim`.
+ * Even if defaultResultOrder is `ipv4first`, `dns.lookup` result maybe same.
+ * For example, when IPv6 is not supported on that machine/network.
+ */
+export async function getLocalhostAddressIfDiffersFromDNS(): Promise<
+  string | undefined
+> {
+  const [nodeResult, dnsResult] = await Promise.all([
+    dns.lookup('localhost'),
+    dns.lookup('localhost', { verbatim: true }),
+  ]);
+  const isSame =
+    nodeResult.family === dnsResult.family &&
+    nodeResult.address === dnsResult.address;
+  return isSame ? undefined : nodeResult.address;
+}
+
+export async function getResolvedClientConfig(
+  clientConfig: DevConfig['client'],
+  serverConfig: ServerConfig,
+): Promise<DevConfig['client']> {
+  const resolvedServerHostname = (await resolveHostname(serverConfig.host))
+    .name;
+  const resolvedServerPort = serverConfig.port!;
+  return {
+    ...clientConfig,
+    host: resolvedServerHostname,
+    port: resolvedServerPort,
+  };
+}