fix: backport #7317 to v2 (#7318)

Author: hi-ogawa

.github/workflows/ci.yml

@@ -10,8 +10,6 @@ on:
       - main
 
   pull_request:
-    branches:
-      - main
 
   workflow_dispatch:
 

packages/browser/src/client/client.ts

@@ -14,7 +14,7 @@ export const SESSION_ID
     : getBrowserState().testerId
 export const ENTRY_URL = `${
   location.protocol === 'https:' ? 'wss:' : 'ws:'
-}//${HOST}/__vitest_browser_api__?type=${PAGE_TYPE}&sessionId=${SESSION_ID}`
+}//${HOST}/__vitest_browser_api__?type=${PAGE_TYPE}&sessionId=${SESSION_ID}&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

@@ -29,6 +29,7 @@
     provider: { __VITEST_PROVIDER__ },
     providedContext: { __VITEST_PROVIDED_CONTEXT__ },
   };
+  window.VITEST_API_TOKEN = { __VITEST_API_TOKEN__ };
 
   const config = __vitest_browser_runner__.config;
 

packages/browser/src/node/rpc.ts

@@ -8,7 +8,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')
@@ -32,6 +32,11 @@ export function setupBrowserRpc(server: BrowserServer) {
       return
     }
 
+    if (!isValidApiRequest(ctx.config, request)) {
+      socket.destroy()
+      return
+    }
+
     const type = searchParams.get('type') ?? 'tester'
     const sessionId = searchParams.get('sessionId') ?? '0'
 

packages/browser/src/node/serverOrchestrator.ts

@@ -32,6 +32,7 @@ export async function resolveOrchestrator(
     __VITEST_CONTEXT_ID__: JSON.stringify(contextId),
     __VITEST_TESTER_ID__: '"none"',
     __VITEST_PROVIDED_CONTEXT__: '{}',
+    __VITEST_API_TOKEN__: JSON.stringify(project.ctx.config.api.token),
   })
 
   // disable CSP for the orchestrator as we are the ones controlling it

packages/browser/src/node/serverTester.ts

@@ -52,6 +52,7 @@ export async function resolveTester(
     __VITEST_CONTEXT_ID__: JSON.stringify(contextId),
     __VITEST_TESTER_ID__: JSON.stringify(crypto.randomUUID()),
     __VITEST_PROVIDED_CONTEXT__: JSON.stringify(stringify(project.getProvidedContext())),
+    __VITEST_API_TOKEN__: JSON.stringify(project.ctx.config.api.token),
   })
 
   const testerHtml = typeof server.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,27 @@ 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
+        server.middlewares.use((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)

packages/vitest/src/node/config/resolveConfig.ts

@@ -9,6 +9,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'
@@ -584,7 +585,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 =>

packages/vitest/src/node/types/config.ts

@@ -1001,7 +1001,7 @@ export interface ResolvedConfig
 
   defines: Record<string, any>
 
-  api?: ApiConfig
+  api: ApiConfig & { token: string }
   cliExclude?: string[]
 
   benchmark?: Required<

packages/vitest/src/public/node.ts

@@ -2,6 +2,7 @@ import type { ModuleDiagnostic as _FileDiagnostic } from '../node/reporters/repo
 import { createServer as _createServer } from 'vite'
 import { TestModule as _TestFile } from '../node/reporters/reported-tasks'
 
+export { isValidApiRequest } from '../api/check'
 export { parseCLI } from '../node/cli/cac'
 export { startVitest } from '../node/cli/cli-api'
 export { resolveApiServerConfig, resolveConfig } from '../node/config/resolveConfig'

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),
     })
   })
 })