fix: verify token for HMR WebSocket connection

sapphi-red · sapphi-red · commit 029dcd6d77d3 · 2025-01-20T18:18:09.000+09:00
diff --git a/packages/vite/src/client/client.ts b/packages/vite/src/client/client.ts
@@ -19,6 +19,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...')
 
@@ -35,12 +36,16 @@ const socketHost = `${__HMR_HOSTNAME__ || importMetaUrl.hostname}:${
 const directSocketHost = __HMR_DIRECT_TARGET__
 const base = __BASE__ || '/'
 const hmrTimeout = __HMR_TIMEOUT__
+const wsToken = __WS_TOKEN__
 
 const transport = normalizeModuleRunnerTransport(
   (() => {
     let wsTransport = createWebSocketModuleRunnerTransport({
       createConnection: () =>
-        new WebSocket(`${socketProtocol}://${socketHost}`, 'vite-hmr'),
+        new WebSocket(
+          `${socketProtocol}://${socketHost}?token=${wsToken}`,
+          'vite-hmr',
+        ),
       pingInterval: hmrTimeout,
     })
 
@@ -54,7 +59,7 @@ const transport = normalizeModuleRunnerTransport(
             wsTransport = createWebSocketModuleRunnerTransport({
               createConnection: () =>
                 new WebSocket(
-                  `${socketProtocol}://${directSocketHost}`,
+                  `${socketProtocol}://${directSocketHost}?token=${wsToken}`,
                   'vite-hmr',
                 ),
               pingInterval: hmrTimeout,
@@ -241,7 +246,9 @@ async function handleMessage(payload: HotPayload) {
         if (hasDocument && !willUnload) {
           console.log(`[vite] server connection lost. Polling for restart...`)
           const socket = payload.data.webSocket as WebSocket
-          await waitForSuccessfulPing(socket.url)
+          const url = new URL(socket.url)
+          url.search = '' // remove query string including `token`
+          await waitForSuccessfulPing(url.href)
           location.reload()
         }
       }
diff --git a/packages/vite/src/node/config.ts b/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 { build } from 'esbuild'
@@ -519,6 +520,18 @@ export interface LegacyOptions {
    * https://github.com/vitejs/vite/discussions/14697.
    */
   proxySsrExternalModules?: boolean
+  /**
+   * In Vite 6.0.8 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+, 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 {
@@ -593,6 +606,17 @@ export type ResolvedConfig = Readonly<
     appType: AppType
     experimental: ExperimentalOptions
     environments: Record<string, ResolvedEnvironmentOptions>
+    /**
+     * 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
     /** @internal */
     fsDenyGlob: AnymatchFn
     /** @internal */
@@ -673,6 +697,7 @@ export const configDefaults = Object.freeze({
   },
   legacy: {
     proxySsrExternalModules: false,
+    skipWebSocketTokenCheck: false,
   },
   logLevel: 'info',
   customLogger: undefined,
@@ -1420,6 +1445,13 @@ export async function resolveConfig(
 
     environments: resolvedEnvironments,
 
+    // 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 --git a/packages/vite/src/node/plugins/clientInjections.ts b/packages/vite/src/node/plugins/clientInjections.ts
@@ -76,6 +76,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
@@ -90,6 +91,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 --git a/packages/vite/src/node/server/ws.ts b/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'
@@ -87,6 +88,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,
@@ -116,7 +140,6 @@ export function createWebSocketServer(
     }
   }
 
-  let wss: WebSocketServerRaw_
   let wsHttpServer: Server | undefined = undefined
 
   const hmr = isObject(config.server.hmr) && config.server.hmr
@@ -135,23 +158,64 @@ export function createWebSocketServer(
   const port = hmrPort || 24678
   const host = (hmr && hmr.host) || undefined
 
+  const shouldHandle = (req: IncomingMessage) => {
+    const protocol = req.headers['sec-websocket-protocol']!
+    // vite-ping is allowed to connect from anywhere
+    // because it needs to be connected before the client fetches the new `/@vite/client`
+    // this is fine because vite-ping does not receive / send any meaningful data
+    if (protocol === 'vite-ping') return true
+
+    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) => {
+      // vite-ping is allowed to connect from anywhere
+      // we close the connection immediately without connection event
+      // so that the client does not get included in `wss.clients`
+      if (isPing) {
+        ws.close(/* Normal Closure */ 1000)
+        return
+      }
+      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 protocol = req.headers['sec-websocket-protocol']!
+      const parsedUrl = new URL(`http://example.com${req.url!}`)
       if (
-        [HMR_HEADER, 'vite-ping'].includes(
-          req.headers['sec-websocket-protocol']!,
-        ) &&
-        req.url === hmrBase
+        [HMR_HEADER, 'vite-ping'].includes(protocol) &&
+        parsedUrl.pathname === hmrBase
       ) {
-        wss.handleUpgrade(req, socket as Socket, head, (ws) => {
-          wss.emit('connection', ws, req)
-        })
+        handleUpgrade(req, socket as Socket, head, protocol === 'vite-ping')
       }
     }
     wsServer.on('upgrade', hmrServerWsListener)
@@ -177,7 +241,6 @@ export function createWebSocketServer(
     } else {
       wsHttpServer = createHttpServer(route)
     }
-    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) {
@@ -187,9 +250,7 @@ export function createWebSocketServer(
         req.destroy()
         return
       }
-      wss.handleUpgrade(req, socket as Socket, head, (ws) => {
-        wss.emit('connection', ws, req)
-      })
+      handleUpgrade(req, socket as Socket, head, protocol === 'vite-ping')
     })
     wsHttpServer.on('error', (e: Error & { code: string }) => {
       if (e.code === 'EADDRINUSE') {
@@ -207,9 +268,6 @@ export function createWebSocketServer(
   }
 
   wss.on('connection', (socket) => {
-    if (socket.protocol === 'vite-ping') {
-      return
-    }
     socket.on('message', (raw) => {
       if (!customListeners.size) return
       let parsed: any
diff --git a/playground/fs-serve/__tests__/fs-serve.spec.ts b/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 --git a/playground/fs-serve/package.json b/playground/fs-serve/package.json
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
