fix: use websocket to test server liveness before client reload (vite…

…js#17891)

Co-authored-by: sapphi-red <[email protected]>

packages/vite/src/client/client.ts

@@ -331,24 +331,28 @@ async function waitForSuccessfulPing(
   hostAndPath: string,
   ms = 1000,
 ) {
-  const pingHostProtocol = socketProtocol === 'wss' ? 'https' : 'http'
-
-  const ping = async () => {
-    // A fetch on a websocket URL will return a successful promise with status 400,
-    // but will reject a networking error.
-    // When running on middleware mode, it returns status 426, and a cors error happens if mode is not no-cors
-    try {
-      await fetch(`${pingHostProtocol}://${hostAndPath}`, {
-        mode: 'no-cors',
-        headers: {
-          // Custom headers won't be included in a request with no-cors so (ab)use one of the
-          // safelisted headers to identify the ping request
-          Accept: 'text/x-vite-ping',
-        },
-      })
-      return true
-    } catch {}
-    return false
+  async function ping() {
+    const socket = new WebSocket(
+      `${socketProtocol}://${hostAndPath}`,
+      'vite-ping',
+    )
+    return new Promise<boolean>((resolve) => {
+      function onOpen() {
+        resolve(true)
+        close()
+      }
+      function onError() {
+        resolve(false)
+        close()
+      }
+      function close() {
+        socket.removeEventListener('open', onOpen)
+        socket.removeEventListener('error', onError)
+        socket.close()
+      }
+      socket.addEventListener('open', onOpen)
+      socket.addEventListener('error', onError)
+    })
   }
 
   if (await ping()) {

packages/vite/src/node/server/ws.ts

@@ -133,7 +133,9 @@ export function createWebSocketServer(
     wss = new WebSocketServerRaw({ noServer: true })
     hmrServerWsListener = (req, socket, head) => {
       if (
-        req.headers['sec-websocket-protocol'] === HMR_HEADER &&
+        [HMR_HEADER, 'vite-ping'].includes(
+          req.headers['sec-websocket-protocol']!,
+        ) &&
         req.url === hmrBase
       ) {
         wss.handleUpgrade(req, socket as Socket, head, (ws) => {
@@ -157,17 +159,46 @@ export function createWebSocketServer(
       })
       res.end(body)
     }) as Parameters<typeof createHttpServer>[1]
+    // vite dev server in middleware mode
+    // need to call ws listen manually
     if (httpsOptions) {
       wsHttpServer = createHttpsServer(httpsOptions, route)
     } else {
       wsHttpServer = createHttpServer(route)
     }
-    // vite dev server in middleware mode
-    // need to call ws listen manually
-    wss = new WebSocketServerRaw({ server: wsHttpServer })
+    wss = new WebSocketServerRaw({ noServer: true })
+    wsHttpServer.on('upgrade', (req, socket, head) => {
+      const protocol = req.headers['sec-websocket-protocol']!
+      if (protocol === 'vite-ping' && server && !server.listening) {
+        // reject connection to tell the vite/client that the server is not ready
+        // if the http server is not listening
+        // because the ws server listens before the http server listens
+        req.destroy()
+        return
+      }
+      wss.handleUpgrade(req, socket as Socket, head, (ws) => {
+        wss.emit('connection', ws, req)
+      })
+    })
+    wsHttpServer.on('error', (e: Error & { code: string }) => {
+      if (e.code === 'EADDRINUSE') {
+        config.logger.error(
+          colors.red(`WebSocket server error: Port is already in use`),
+          { error: e },
+        )
+      } else {
+        config.logger.error(
+          colors.red(`WebSocket server error:\n${e.stack || e.message}`),
+          { error: e },
+        )
+      }
+    })
   }
 
   wss.on('connection', (socket) => {
+    if (socket.protocol === 'vite-ping') {
+      return
+    }
     socket.on('message', (raw) => {
       if (!customListeners.size) return
       let parsed: any

playground/client-reload/__tests__/client-reload.spec.ts

@@ -0,0 +1,75 @@
+import path from 'node:path'
+import { type ServerOptions, type ViteDevServer, createServer } from 'vite'
+import { afterEach, describe, expect, test } from 'vitest'
+import { hmrPorts, isServe, page, ports } from '~utils'
+
+let server: ViteDevServer
+
+afterEach(async () => {
+  await server?.close()
+})
+
+async function testClientReload(serverOptions: ServerOptions) {
+  // start server
+  server = await createServer({
+    root: path.resolve(import.meta.dirname, '..'),
+    logLevel: 'silent',
+    server: {
+      strictPort: true,
+      ...serverOptions,
+    },
+  })
+
+  await server.listen()
+  const serverUrl = server.resolvedUrls.local[0]
+
+  // open page and wait for connection
+  const connectedPromise = page.waitForEvent('console', {
+    predicate: (message) => message.text().includes('[vite] connected.'),
+    timeout: 5000,
+  })
+  await page.goto(serverUrl)
+  await connectedPromise
+
+  // input state
+  await page.locator('input').fill('hello')
+
+  // restart and wait for reconnection after reload
+  const reConnectedPromise = page.waitForEvent('console', {
+    predicate: (message) => message.text().includes('[vite] connected.'),
+    timeout: 5000,
+  })
+  await server.restart()
+  await reConnectedPromise
+  expect(await page.textContent('input')).toBe('')
+}
+
+describe.runIf(isServe)('client-reload', () => {
+  test('default', async () => {
+    await testClientReload({
+      port: ports['client-reload'],
+    })
+  })
+
+  test('custom hmr port', async () => {
+    await testClientReload({
+      port: ports['client-reload/hmr-port'],
+      hmr: {
+        port: hmrPorts['client-reload/hmr-port'],
+      },
+    })
+  })
+
+  test('custom hmr port and cross origin isolation', async () => {
+    await testClientReload({
+      port: ports['client-reload/cross-origin'],
+      hmr: {
+        port: hmrPorts['client-reload/cross-origin'],
+      },
+      headers: {
+        'Cross-Origin-Embedder-Policy': 'require-corp',
+        'Cross-Origin-Opener-Policy': 'same-origin',
+      },
+    })
+  })
+})

playground/client-reload/__tests__/serve.ts

@@ -0,0 +1,6 @@
+// do nothing here since server is managed inside spec
+export async function serve(): Promise<{ close(): Promise<void> }> {
+  return {
+    close: () => Promise.resolve(),
+  }
+}

playground/client-reload/index.html

@@ -0,0 +1,4 @@
+<body>
+  <h4>Test Client Reload</h4>
+  <input />
+</body>

playground/client-reload/package.json

@@ -0,0 +1,12 @@
+{
+  "name": "@vitejs/test-client-reload",
+  "private": true,
+  "version": "0.0.0",
+  "type": "module",
+  "scripts": {
+    "debug": "node --inspect-brk ../../packages/vite/bin/vite",
+    "dev": "vite",
+    "build": "vite build",
+    "preview": "vite preview"
+  }
+}

playground/client-reload/vite.config.ts

@@ -0,0 +1,5 @@
+import { defineConfig } from 'vite'
+
+export default defineConfig({
+  server: {},
+})

playground/js-sourcemap/test-ssr-dev.js

@@ -12,6 +12,7 @@ async function runTest() {
     server: {
       middlewareMode: true,
       hmr: false,
+      ws: false,
     },
     define: {
       __testDefineObject: '{ "hello": "test" }',

playground/ssr-html/test-network-imports.js

@@ -8,6 +8,7 @@ async function runTest(userRunner) {
     root: fileURLToPath(new URL('.', import.meta.url)),
     server: {
       middlewareMode: true,
+      ws: false,
     },
   })
   let mod

playground/ssr-html/test-stacktrace-runtime.js

@@ -10,6 +10,7 @@ const server = await createServer({
   root: fileURLToPath(new URL('.', import.meta.url)),
   server: {
     middlewareMode: true,
+    ws: false,
   },
 })
 

playground/ssr-html/test-stacktrace.js

@@ -29,6 +29,7 @@ const vite = await createServer({
   logLevel: isTest ? 'error' : 'info',
   server: {
     middlewareMode: true,
+    ws: false,
   },
   appType: 'custom',
 })

playground/test-utils.ts

@@ -47,6 +47,9 @@ export const ports = {
   'css/dynamic-import': 5007,
   'css/lightningcss-proxy': 5008,
   'backend-integration': 5009,
+  'client-reload': 5010,
+  'client-reload/hmr-port': 5011,
+  'client-reload/cross-origin': 5012,
 }
 export const hmrPorts = {
   'optimize-missing-deps': 24680,
@@ -58,6 +61,8 @@ export const hmrPorts = {
   'css/lightningcss-proxy': 24686,
   json: 24687,
   'ssr-conditions': 24688,
+  'client-reload/hmr-port': 24689,
+  'client-reload/cross-origin': 24690,
 }
 
 const hexToNameMap: Record<string, string> = {}