fix: verify token for HMR WebSocket connection

Author: sapphi-red

Diff for: packages/vite/src/client/client.ts

@@ -15,6 +15,7 @@ declare const __HMR_DIRECT_TARGET__: string
 declare const __HMR_BASE__: string
 declare const __HMR_TIMEOUT__: number
 declare const __HMR_ENABLE_OVERLAY__: boolean
+declare const __WS_TOKEN__: string
 
 console.debug('[vite] connecting...')
 
@@ -30,6 +31,7 @@ const socketHost = `${__HMR_HOSTNAME__ || importMetaUrl.hostname}:${
 }${__HMR_BASE__}`
 const directSocketHost = __HMR_DIRECT_TARGET__
 const base = __BASE__ || '/'
+const wsToken = __WS_TOKEN__
 
 let socket: WebSocket
 try {
@@ -74,7 +76,10 @@ function setupWebSocket(
   hostAndPath: string,
   onCloseWithoutOpen?: () => void,
 ) {
-  const socket = new WebSocket(`${protocol}://${hostAndPath}`, 'vite-hmr')
+  const socket = new WebSocket(
+    `${protocol}://${hostAndPath}?token=${wsToken}`,
+    'vite-hmr',
+  )
   let isOpened = false
 
   socket.addEventListener(

Diff for: packages/vite/src/node/config.ts

@@ -5,6 +5,7 @@ import { pathToFileURL } from 'node:url'
 import { promisify } from 'node:util'
 import { performance } from 'node:perf_hooks'
 import { createRequire } from 'node:module'
+import crypto from 'node:crypto'
 import colors from 'picocolors'
 import type { Alias, AliasOptions } from 'dep-types/alias'
 import aliasPlugin from '@rollup/plugin-alias'
@@ -341,6 +342,18 @@ export interface LegacyOptions {
    * https://github.com/vitejs/vite/discussions/14697.
    */
   proxySsrExternalModules?: boolean
+  /**
+   * In Vite 6.0.8 / 5.4.11 and below, WebSocket server was able to connect from any web pages. However,
+   * that could be exploited by a malicious web page.
+   *
+   * In Vite 6.0.9+ / 5.4.12+, the WebSocket server now requires a token to connect from a web page.
+   * But this may break some plugins and frameworks that connects to the WebSocket server
+   * on their own. Enabling this option will make Vite skip the token check.
+   *
+   * **We do not recommend enabling this option unless you are sure that you are fine with
+   * that security weakness.**
+   */
+  skipWebSocketTokenCheck?: boolean
 }
 
 export interface ResolvedWorkerOptions {
@@ -400,6 +413,17 @@ export type ResolvedConfig = Readonly<
     worker: ResolvedWorkerOptions
     appType: AppType
     experimental: ExperimentalOptions
+    /**
+     * The token to connect to the WebSocket server from browsers.
+     *
+     * We recommend using `import.meta.hot` rather than connecting
+     * to the WebSocket server directly.
+     * If you have a usecase that requires connecting to the WebSocket
+     * server, please create an issue so that we can discuss.
+     *
+     * @deprecated
+     */
+    webSocketToken: string
   } & PluginHookUtils
 >
 
@@ -828,6 +852,12 @@ export async function resolveConfig(
       hmrPartialAccept: false,
       ...config.experimental,
     },
+    // random 72 bits (12 base64 chars)
+    // at least 64bits is recommended
+    // https://owasp.org/www-community/vulnerabilities/Insufficient_Session-ID_Length
+    webSocketToken: Buffer.from(
+      crypto.getRandomValues(new Uint8Array(9)),
+    ).toString('base64url'),
     getSortedPlugins: undefined!,
     getSortedPluginHooks: undefined!,
   }

Diff for: packages/vite/src/node/plugins/clientInjections.ts

@@ -72,6 +72,7 @@ export function clientInjectionsPlugin(config: ResolvedConfig): Plugin {
       const hmrTimeoutReplacement = escapeReplacement(timeout)
       const hmrEnableOverlayReplacement = escapeReplacement(overlay)
       const hmrConfigNameReplacement = escapeReplacement(hmrConfigName)
+      const wsTokenReplacement = escapeReplacement(config.webSocketToken)
 
       injectConfigValues = (code: string) => {
         return code
@@ -87,6 +88,7 @@ export function clientInjectionsPlugin(config: ResolvedConfig): Plugin {
           .replace(`__HMR_TIMEOUT__`, hmrTimeoutReplacement)
           .replace(`__HMR_ENABLE_OVERLAY__`, hmrEnableOverlayReplacement)
           .replace(`__HMR_CONFIG_NAME__`, hmrConfigNameReplacement)
+          .replace(`__WS_TOKEN__`, wsTokenReplacement)
       }
     },
     async transform(code, id, options) {

Diff for: packages/vite/src/node/server/ws.ts

@@ -5,6 +5,7 @@ import type { ServerOptions as HttpsServerOptions } from 'node:https'
 import { createServer as createHttpsServer } from 'node:https'
 import type { Socket } from 'node:net'
 import type { Duplex } from 'node:stream'
+import crypto from 'node:crypto'
 import colors from 'picocolors'
 import type { WebSocket as WebSocketRaw } from 'ws'
 import { WebSocketServer as WebSocketServerRaw_ } from 'ws'
@@ -89,6 +90,29 @@ function noop() {
   // noop
 }
 
+// we only allow websockets to be connected if it has a valid token
+// this is to prevent untrusted origins to connect to the server
+// for example, Cross-site WebSocket hijacking
+//
+// we should check the token before calling wss.handleUpgrade
+// otherwise untrusted ws clients will be included in wss.clients
+//
+// using the query params means the token might be logged out in server or middleware logs
+// but we assume that is not an issue since the token is regenerated for each process
+function hasValidToken(config: ResolvedConfig, url: URL) {
+  const token = url.searchParams.get('token')
+  if (!token) return false
+
+  try {
+    const isValidToken = crypto.timingSafeEqual(
+      Buffer.from(token),
+      Buffer.from(config.webSocketToken),
+    )
+    return isValidToken
+  } catch {} // an error is thrown when the length is incorrect
+  return false
+}
+
 export function createWebSocketServer(
   server: HttpServer | null,
   config: ResolvedConfig,
@@ -110,7 +134,6 @@ export function createWebSocketServer(
     }
   }
 
-  let wss: WebSocketServerRaw_
   let wsHttpServer: Server | undefined = undefined
 
   const hmr = isObject(config.server.hmr) && config.server.hmr
@@ -129,21 +152,50 @@ export function createWebSocketServer(
   const port = hmrPort || 24678
   const host = (hmr && hmr.host) || undefined
 
+  const shouldHandle = (req: IncomingMessage) => {
+    if (config.legacy?.skipWebSocketTokenCheck) {
+      return true
+    }
+
+    // If the Origin header is set, this request might be coming from a browser.
+    // Browsers always sets the Origin header for WebSocket connections.
+    if (req.headers.origin) {
+      const parsedUrl = new URL(`http://example.com${req.url!}`)
+      return hasValidToken(config, parsedUrl)
+    }
+
+    // We allow non-browser requests to connect without a token
+    // for backward compat and convenience
+    // This is fine because if you can sent a request without the SOP limitation,
+    // you can also send a normal HTTP request to the server.
+    return true
+  }
+  const handleUpgrade = (
+    req: IncomingMessage,
+    socket: Duplex,
+    head: Buffer,
+    _isPing: boolean,
+  ) => {
+    wss.handleUpgrade(req, socket as Socket, head, (ws) => {
+      wss.emit('connection', ws, req)
+    })
+  }
+  const wss: WebSocketServerRaw_ = new WebSocketServerRaw({ noServer: true })
+  wss.shouldHandle = shouldHandle
+
   if (wsServer) {
     let hmrBase = config.base
     const hmrPath = hmr ? hmr.path : undefined
     if (hmrPath) {
       hmrBase = path.posix.join(hmrBase, hmrPath)
     }
-    wss = new WebSocketServerRaw({ noServer: true })
     hmrServerWsListener = (req, socket, head) => {
+      const parsedUrl = new URL(`http://example.com${req.url!}`)
       if (
         req.headers['sec-websocket-protocol'] === HMR_HEADER &&
-        req.url === hmrBase
+        parsedUrl.pathname === hmrBase
       ) {
-        wss.handleUpgrade(req, socket as Socket, head, (ws) => {
-          wss.emit('connection', ws, req)
-        })
+        handleUpgrade(req, socket as Socket, head, false)
       }
     }
     wsServer.on('upgrade', hmrServerWsListener)
@@ -167,9 +219,22 @@ export function createWebSocketServer(
     } else {
       wsHttpServer = createHttpServer(route)
     }
-    // vite dev server in middleware mode
-    // need to call ws listen manually
-    wss = new WebSocketServerRaw({ server: wsHttpServer })
+    wsHttpServer.on('upgrade', (req, socket, head) => {
+      handleUpgrade(req, socket as Socket, head, false)
+    })
+    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) => {

Diff for: playground/fs-serve/__tests__/fs-serve.spec.ts

@@ -8,8 +8,9 @@ import {
   test,
 } from 'vitest'
 import type { Page } from 'playwright-chromium'
+import WebSocket from 'ws'
 import testJSON from '../safe.json'
-import { browser, isServe, page, viteTestUrl } from '~utils'
+import { browser, isServe, page, viteServer, viteTestUrl } from '~utils'
 
 const getViteTestIndexHtmlUrl = () => {
   const srcPrefix = viteTestUrl.endsWith('/') ? '' : '/'
@@ -139,6 +140,51 @@ describe('cross origin', () => {
     }, url)
   }
 
+  const connectWebSocketFromPage = async (page: Page, url: string) => {
+    return await page.evaluate(async (url: string) => {
+      try {
+        const ws = new globalThis.WebSocket(url, ['vite-hmr'])
+        await new Promise<void>((resolve, reject) => {
+          ws.addEventListener('open', () => {
+            resolve()
+            ws.close()
+          })
+          ws.addEventListener('error', () => {
+            reject()
+          })
+        })
+        return true
+      } catch {
+        return false
+      }
+    }, url)
+  }
+
+  const connectWebSocketFromServer = async (
+    url: string,
+    origin: string | undefined,
+  ) => {
+    try {
+      const ws = new WebSocket(url, ['vite-hmr'], {
+        headers: {
+          ...(origin ? { Origin: origin } : undefined),
+        },
+      })
+      await new Promise<void>((resolve, reject) => {
+        ws.addEventListener('open', () => {
+          resolve()
+          ws.close()
+        })
+        ws.addEventListener('error', () => {
+          reject()
+        })
+      })
+      return true
+    } catch {
+      return false
+    }
+  }
+
   describe('allowed for same origin', () => {
     beforeEach(async () => {
       await page.goto(getViteTestIndexHtmlUrl())
@@ -156,6 +202,23 @@ describe('cross origin', () => {
       )
       expect(status).toBe(200)
     })
+
+    test.runIf(isServe)('connect WebSocket with valid token', async () => {
+      const token = viteServer.config.webSocketToken
+      const result = await connectWebSocketFromPage(
+        page,
+        `${viteTestUrl}?token=${token}`,
+      )
+      expect(result).toBe(true)
+    })
+
+    test.runIf(isServe)(
+      'connect WebSocket without a token without the origin header',
+      async () => {
+        const result = await connectWebSocketFromServer(viteTestUrl, undefined)
+        expect(result).toBe(true)
+      },
+    )
   })
 
   describe('denied for different origin', async () => {
@@ -180,5 +243,31 @@ describe('cross origin', () => {
       )
       expect(status).not.toBe(200)
     })
+
+    test.runIf(isServe)('connect WebSocket without token', async () => {
+      const result = await connectWebSocketFromPage(page, viteTestUrl)
+      expect(result).toBe(false)
+
+      const result2 = await connectWebSocketFromPage(
+        page,
+        `${viteTestUrl}?token=`,
+      )
+      expect(result2).toBe(false)
+    })
+
+    test.runIf(isServe)('connect WebSocket with invalid token', async () => {
+      const token = viteServer.config.webSocketToken
+      const result = await connectWebSocketFromPage(
+        page,
+        `${viteTestUrl}?token=${'t'.repeat(token.length)}`,
+      )
+      expect(result).toBe(false)
+
+      const result2 = await connectWebSocketFromPage(
+        page,
+        `${viteTestUrl}?token=${'t'.repeat(token.length)}t`, // different length
+      )
+      expect(result2).toBe(false)
+    })
   })
 })

Diff for: playground/fs-serve/package.json

@@ -14,5 +14,8 @@
     "dev:deny": "vite root --config ./root/vite.config-deny.js",
     "build:deny": "vite build root --config ./root/vite.config-deny.js",
     "preview:deny": "vite preview root --config ./root/vite.config-deny.js"
+  },
+  "devDependencies": {
+    "ws": "^8.18.0"
   }
 }