From 9df6e6beabf0d18988ec13b8b742d2aba29662f9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=BF=A0=20/=20green?= <green@sapphi.red>
Date: Tue, 21 Jan 2025 16:46:56 +0900
Subject: [PATCH] fix: `preview.allowedHosts` with specific values was not
 respected (#19246)

---
 packages/vite/src/node/preview.ts             |  2 +-
 packages/vite/src/node/server/index.ts        |  2 +-
 .../src/node/server/middlewares/hostCheck.ts  | 34 +++++++++++++------
 packages/vite/src/node/server/ws.ts           |  2 +-
 4 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/packages/vite/src/node/preview.ts b/packages/vite/src/node/preview.ts
index fe9cbf9c1bddab..3947fb1aa2ac0a 100644
--- a/packages/vite/src/node/preview.ts
+++ b/packages/vite/src/node/preview.ts
@@ -194,7 +194,7 @@ export async function preview(
   const { allowedHosts } = config.preview
   // no need to check for HTTPS as HTTPS is not vulnerable to DNS rebinding attacks
   if (allowedHosts !== true && !config.preview.https) {
-    app.use(hostCheckMiddleware(config))
+    app.use(hostCheckMiddleware(config, true))
   }
 
   // proxy
diff --git a/packages/vite/src/node/server/index.ts b/packages/vite/src/node/server/index.ts
index acfc7b00b4affb..9378a479fbb293 100644
--- a/packages/vite/src/node/server/index.ts
+++ b/packages/vite/src/node/server/index.ts
@@ -858,7 +858,7 @@ export async function _createServer(
   const { allowedHosts } = serverConfig
   // no need to check for HTTPS as HTTPS is not vulnerable to DNS rebinding attacks
   if (allowedHosts !== true && !serverConfig.https) {
-    middlewares.use(hostCheckMiddleware(config))
+    middlewares.use(hostCheckMiddleware(config, false))
   }
 
   middlewares.use(cachedTransformMiddleware(server))
diff --git a/packages/vite/src/node/server/middlewares/hostCheck.ts b/packages/vite/src/node/server/middlewares/hostCheck.ts
index 3433e01a3e7dff..2d706fba596c1b 100644
--- a/packages/vite/src/node/server/middlewares/hostCheck.ts
+++ b/packages/vite/src/node/server/middlewares/hostCheck.ts
@@ -3,7 +3,8 @@ import type { Connect } from 'dep-types/connect'
 import type { ResolvedConfig } from '../../config'
 import type { ResolvedPreviewOptions, ResolvedServerOptions } from '../..'
 
-const allowedHostsCache = new WeakMap<ResolvedConfig, Set<string>>()
+const allowedHostsServerCache = new WeakMap<ResolvedConfig, Set<string>>()
+const allowedHostsPreviewCache = new WeakMap<ResolvedConfig, Set<string>>()
 
 const isFileOrExtensionProtocolRE = /^(?:file|.+-extension):/i
 
@@ -118,48 +119,59 @@ export function isHostAllowedWithoutCache(
 
 /**
  * @param config resolved config
+ * @param isPreview whether it's for the preview server or not
  * @param host the value of host header. See [RFC 9110 7.2](https://datatracker.ietf.org/doc/html/rfc9110#name-host-and-authority).
  */
-export function isHostAllowed(config: ResolvedConfig, host: string): boolean {
-  if (config.server.allowedHosts === true) {
+export function isHostAllowed(
+  config: ResolvedConfig,
+  isPreview: boolean,
+  host: string,
+): boolean {
+  const allowedHosts = isPreview
+    ? config.preview.allowedHosts
+    : config.server.allowedHosts
+  if (allowedHosts === true) {
     return true
   }
 
-  if (!allowedHostsCache.has(config)) {
-    allowedHostsCache.set(config, new Set())
+  const cache = isPreview ? allowedHostsPreviewCache : allowedHostsServerCache
+  if (!cache.has(config)) {
+    cache.set(config, new Set())
   }
 
-  const allowedHosts = allowedHostsCache.get(config)!
-  if (allowedHosts.has(host)) {
+  const cachedAllowedHosts = cache.get(config)!
+  if (cachedAllowedHosts.has(host)) {
     return true
   }
 
   const result = isHostAllowedWithoutCache(
-    config.server.allowedHosts ?? [],
+    allowedHosts ?? [],
     config.additionalAllowedHosts,
     host,
   )
   if (result) {
-    allowedHosts.add(host)
+    cachedAllowedHosts.add(host)
   }
   return result
 }
 
 export function hostCheckMiddleware(
   config: ResolvedConfig,
+  isPreview: boolean,
 ): Connect.NextHandleFunction {
   // Keep the named function. The name is visible in debug logs via `DEBUG=connect:dispatcher ...`
   return function viteHostCheckMiddleware(req, res, next) {
     const hostHeader = req.headers.host
-    if (!hostHeader || !isHostAllowed(config, hostHeader)) {
+    if (!hostHeader || !isHostAllowed(config, isPreview, hostHeader)) {
       const hostname = hostHeader?.replace(/:\d+$/, '')
       const hostnameWithQuotes = JSON.stringify(hostname)
+      const optionName = `${isPreview ? 'preview' : 'server'}.allowedHosts`
       res.writeHead(403, {
         'Content-Type': 'text/plain',
       })
       res.end(
         `Blocked request. This host (${hostnameWithQuotes}) is not allowed.\n` +
-          `To allow this host, add ${hostnameWithQuotes} to \`server.allowedHosts\` in vite.config.js.`,
+          `To allow this host, add ${hostnameWithQuotes} to \`${optionName}\` in vite.config.js.`,
       )
       return
     }
diff --git a/packages/vite/src/node/server/ws.ts b/packages/vite/src/node/server/ws.ts
index 222f5f0e32d775..ca672d7f835dbe 100644
--- a/packages/vite/src/node/server/ws.ts
+++ b/packages/vite/src/node/server/ws.ts
@@ -155,7 +155,7 @@ export function createWebSocketServer(
 
   const shouldHandle = (req: IncomingMessage) => {
     const hostHeader = req.headers.host
-    if (!hostHeader || !isHostAllowed(config, hostHeader)) {
+    if (!hostHeader || !isHostAllowed(config, false, hostHeader)) {
       return false
     }