fix: validate websocket request (#7317)

Co-authored-by: Ari Perkkiö <[email protected]>

Author: hi-ogawa

packages/browser/src/client/client.ts

@@ -15,7 +15,7 @@ export const RPC_ID
 const METHOD = getBrowserState().method
 export const ENTRY_URL = `${
   location.protocol === 'https:' ? 'wss:' : 'ws:'
-}//${HOST}/__vitest_browser_api__?type=${PAGE_TYPE}&rpcId=${RPC_ID}&sessionId=${getBrowserState().sessionId}&projectName=${getBrowserState().config.name || ''}&method=${METHOD}`
+}//${HOST}/__vitest_browser_api__?type=${PAGE_TYPE}&rpcId=${RPC_ID}&sessionId=${getBrowserState().sessionId}&projectName=${getBrowserState().config.name || ''}&method=${METHOD}&token=${(window as any).VITEST_API_TOKEN}`
 
 let setCancel = (_: CancelReason) => {}
 export const onCancel = new Promise<CancelReason>((resolve) => {

packages/browser/src/client/public/esm-client-injector.js

@@ -30,6 +30,7 @@
     method: { __VITEST_METHOD__ },
     providedContext: { __VITEST_PROVIDED_CONTEXT__ },
   };
+  window.VITEST_API_TOKEN = { __VITEST_API_TOKEN__ };
 
   const config = __vitest_browser_runner__.config;
 

packages/browser/src/node/rpc.ts

@@ -10,7 +10,7 @@ import { ServerMockResolver } from '@vitest/mocker/node'
 import { createBirpc } from 'birpc'
 import { parse, stringify } from 'flatted'
 import { dirname } from 'pathe'
-import { createDebugger, isFileServingAllowed } from 'vitest/node'
+import { createDebugger, isFileServingAllowed, isValidApiRequest } from 'vitest/node'
 import { WebSocketServer } from 'ws'
 
 const debug = createDebugger('vitest:browser:api')
@@ -33,6 +33,11 @@ export function setupBrowserRpc(globalServer: ParentBrowserProject) {
       return
     }
 
+    if (!isValidApiRequest(vitest.config, request)) {
+      socket.destroy()
+      return
+    }
+
     const type = searchParams.get('type')
     const rpcId = searchParams.get('rpcId')
     const sessionId = searchParams.get('sessionId')

packages/browser/src/node/serverOrchestrator.ts

@@ -42,6 +42,7 @@ export async function resolveOrchestrator(
     __VITEST_SESSION_ID__: JSON.stringify(sessionId),
     __VITEST_TESTER_ID__: '"none"',
     __VITEST_PROVIDED_CONTEXT__: '{}',
+    __VITEST_API_TOKEN__: JSON.stringify(globalServer.vitest.config.api.token),
   })
 
   // disable CSP for the orchestrator as we are the ones controlling it

packages/browser/src/node/serverTester.ts

@@ -68,6 +68,7 @@ export async function resolveTester(
     __VITEST_SESSION_ID__: JSON.stringify(sessionId),
     __VITEST_TESTER_ID__: JSON.stringify(crypto.randomUUID()),
     __VITEST_PROVIDED_CONTEXT__: JSON.stringify(stringify(project.getProvidedContext())),
+    __VITEST_API_TOKEN__: JSON.stringify(globalServer.vitest.config.api.token),
   })
 
   const testerHtml = typeof browserProject.testerHtml === 'string'

packages/ui/client/constants.ts

@@ -4,6 +4,6 @@ export const PORT = import.meta.hot && !browserState ? '51204' : location.port
 export const HOST = [location.hostname, PORT].filter(Boolean).join(':')
 export const ENTRY_URL = `${
   location.protocol === 'https:' ? 'wss:' : 'ws:'
-}//${HOST}/__vitest_api__`
+}//${HOST}/__vitest_api__?token=${(window as any).VITEST_API_TOKEN}`
 export const isReport = !!window.METADATA_PATH
 export const BASE_PATH = isReport ? import.meta.env.BASE_URL : __BASE_PATH__

packages/ui/node/index.ts

@@ -1,5 +1,6 @@
 import type { Plugin } from 'vite'
 import type { Vitest } from 'vitest/node'
+import fs from 'node:fs'
 import { fileURLToPath } from 'node:url'
 import { toArray } from '@vitest/utils'
 import { basename, resolve } from 'pathe'
@@ -52,6 +53,28 @@ export default (ctx: Vitest): Plugin => {
         }
 
         const clientDist = resolve(fileURLToPath(import.meta.url), '../client')
+        const clientIndexHtml = fs.readFileSync(resolve(clientDist, 'index.html'), 'utf-8')
+
+        // serve index.html with api token
+        // eslint-disable-next-line prefer-arrow-callback
+        server.middlewares.use(function vitestUiHtmlMiddleware(req, res, next) {
+          if (req.url) {
+            const url = new URL(req.url, 'http://localhost')
+            if (url.pathname === base) {
+              const html = clientIndexHtml.replace(
+                '<!-- !LOAD_METADATA! -->',
+                `<script>window.VITEST_API_TOKEN = ${JSON.stringify(ctx.config.api.token)}</script>`,
+              )
+              res.setHeader('Cache-Control', 'no-cache, max-age=0, must-revalidate')
+              res.setHeader('Content-Type', 'text/html; charset=utf-8')
+              res.write(html)
+              res.end()
+              return
+            }
+          }
+          next()
+        })
+
         server.middlewares.use(
           base,
           sirv(clientDist, {

packages/vitest/rollup.config.js

@@ -72,6 +72,7 @@ const external = [
   'node:os',
   'node:stream',
   'node:vm',
+  'node:http',
   'inspector',
   'vite-node/source-map',
   'vite-node/client',

packages/vitest/src/api/check.ts

@@ -0,0 +1,22 @@
+import type { IncomingMessage } from 'node:http'
+import type { ResolvedConfig } from '../node/types/config'
+import crypto from 'node:crypto'
+
+export function isValidApiRequest(config: ResolvedConfig, req: IncomingMessage): boolean {
+  const url = new URL(req.url ?? '', 'http://localhost')
+
+  // validate token. token is injected in ui/tester/orchestrator html, which is cross origin proteced.
+  try {
+    const token = url.searchParams.get('token')
+    if (token && crypto.timingSafeEqual(
+      Buffer.from(token),
+      Buffer.from(config.api.token),
+    )) {
+      return true
+    }
+  }
+  // an error is thrown when the length is incorrect
+  catch {}
+
+  return false
+}

packages/vitest/src/api/setup.ts

@@ -1,5 +1,6 @@
 import type { File, TaskResultPack } from '@vitest/runner'
 
+import type { IncomingMessage } from 'node:http'
 import type { ViteDevServer } from 'vite'
 import type { WebSocket } from 'ws'
 import type { Vitest } from '../node/core'
@@ -21,6 +22,7 @@ import { API_PATH } from '../constants'
 import { getModuleGraph } from '../utils/graph'
 import { stringifyReplace } from '../utils/serialization'
 import { parseErrorStacktrace } from '../utils/source-map'
+import { isValidApiRequest } from './check'
 
 export function setup(ctx: Vitest, _server?: ViteDevServer) {
   const wss = new WebSocketServer({ noServer: true })
@@ -29,7 +31,7 @@ export function setup(ctx: Vitest, _server?: ViteDevServer) {
 
   const server = _server || ctx.server
 
-  server.httpServer?.on('upgrade', (request, socket, head) => {
+  server.httpServer?.on('upgrade', (request: IncomingMessage, socket, head) => {
     if (!request.url) {
       return
     }
@@ -39,6 +41,11 @@ export function setup(ctx: Vitest, _server?: ViteDevServer) {
       return
     }
 
+    if (!isValidApiRequest(ctx.config, request)) {
+      socket.destroy()
+      return
+    }
+
     wss.handleUpgrade(request, socket, head, (ws) => {
       wss.emit('connection', ws, request)
       setupClient(ws)