fix: validate websocket request (#7317)

hi-ogawa · AriPerkkio · web-flow · commit 191ef9e34c86 · 2025-02-02T10:31:33.000+02:00
Co-authored-by: Ari Perkkiö &lt;ari.perkkio@gmail.com&gt;
diff --git a/packages/browser/src/client/client.ts b/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) => {
diff --git a/packages/browser/src/client/public/esm-client-injector.js b/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;
 
diff --git a/packages/browser/src/node/rpc.ts b/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')
diff --git a/packages/browser/src/node/serverOrchestrator.ts b/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
diff --git a/packages/browser/src/node/serverTester.ts b/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'
diff --git a/packages/ui/client/constants.ts b/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__
diff --git a/packages/ui/node/index.ts b/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, {
diff --git a/packages/vitest/rollup.config.js b/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',
diff --git a/packages/vitest/src/api/check.ts b/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
+}
diff --git a/packages/vitest/src/api/setup.ts b/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)
diff --git a/packages/vitest/src/node/config/resolveConfig.ts b/packages/vitest/src/node/config/resolveConfig.ts
@@ -8,6 +8,7 @@ import type {
 } from '../types/config'
 import type { BaseCoverageOptions, CoverageReporterWithOptions } from '../types/coverage'
 import type { BuiltinPool, ForksOptions, PoolOptions, ThreadsOptions } from '../types/pool-options'
+import crypto from 'node:crypto'
 import { toArray } from '@vitest/utils'
 import { resolveModule } from 'local-pkg'
 import { normalize, relative, resolve } from 'pathe'
@@ -629,7 +630,8 @@ export function resolveConfig(
   }
 
   // the server has been created, we don't need to override vite.server options
-  resolved.api = resolveApiServerConfig(options, defaultPort)
+  const api = resolveApiServerConfig(options, defaultPort)
+  resolved.api = { ...api, token: crypto.randomUUID() }
 
   if (options.related) {
     resolved.related = toArray(options.related).map(file =>
diff --git a/packages/vitest/src/node/types/config.ts b/packages/vitest/src/node/types/config.ts
@@ -1014,7 +1014,7 @@ export interface ResolvedConfig
 
   defines: Record<string, any>
 
-  api?: ApiConfig
+  api: ApiConfig & { token: string }
   cliExclude?: string[]
 
   project: string[]
diff --git a/packages/vitest/src/public/node.ts b/packages/vitest/src/public/node.ts
@@ -5,6 +5,7 @@ import { TestModule as _TestFile } from '../node/reporters/reported-tasks'
 
 export const version = Vitest.version
 
+export { isValidApiRequest } from '../api/check'
 export { parseCLI } from '../node/cli/cac'
 export type { CliParseOptions } from '../node/cli/cac'
 export { startVitest } from '../node/cli/cli-api'
diff --git a/test/config/test/override.test.ts b/test/config/test/override.test.ts
@@ -249,6 +249,7 @@ describe('correctly defines api flag', () => {
     expect(c.server.config.server.middlewareMode).toBe(true)
     expect(c.config.api).toEqual({
       middlewareMode: true,
+      token: expect.any(String),
     })
   })
 
@@ -262,6 +263,7 @@ describe('correctly defines api flag', () => {
     expect(c.server.config.server.port).toBe(4321)
     expect(c.config.api).toEqual({
       port: 4321,
+      token: expect.any(String),
     })
   })
 })
diff --git a/test/ui/test/ui.spec.ts b/test/ui/test/ui.spec.ts
@@ -30,6 +30,36 @@ test.describe('ui', () => {
     await vitest?.close()
   })
 
+  test('security', async ({ page }) => {
+    await page.goto('https://example.com/')
+
+    // request html
+    const htmlResult = await page.evaluate(async (pageUrl) => {
+      try {
+        const res = await fetch(pageUrl)
+        return res.status
+      }
+      catch (e) {
+        return e instanceof Error ? e.message : e
+      }
+    }, pageUrl)
+    expect(htmlResult).toBe('Failed to fetch')
+
+    // request websocket
+    const wsResult = await page.evaluate(async (pageUrl) => {
+      const ws = new WebSocket(new URL('/__vitest_api__', pageUrl))
+      return new Promise((resolve) => {
+        ws.addEventListener('open', () => {
+          resolve('open')
+        })
+        ws.addEventListener('error', () => {
+          resolve('error')
+        })
+      })
+    }, pageUrl)
+    expect(wsResult).toBe('error')
+  })
+
   test('basic', async ({ page }) => {
     const pageErrors: unknown[] = []
     page.on('pageerror', error => pageErrors.push(error))
diff --git a/test/ui/tsconfig.json b/test/ui/tsconfig.json
@@ -1,6 +1,7 @@
 {
   "extends": "../../tsconfig.base.json",
   "compilerOptions": {
+    "lib": ["ESNext", "DOM"],
     "baseUrl": "../..",
     "paths": {
       "vitest/node": [
