feat: support async for proxy.bypass (#18940)

Co-authored-by: 翠 / green <[email protected]>

Author: shulaoda

packages/vite/src/node/server/middlewares/proxy.ts

@@ -27,7 +27,13 @@ export interface ProxyOptions extends HttpProxy.ServerOptions {
     /** undefined for WebSocket upgrade requests */
     res: http.ServerResponse | undefined,
     options: ProxyOptions,
-  ) => void | null | undefined | false | string
+  ) =>
+    | void
+    | null
+    | undefined
+    | false
+    | string
+    | Promise<void | null | undefined | boolean | string>
   /**
    * rewrite the Origin header of a WebSocket request to match the target
    *
@@ -158,7 +164,7 @@ export function proxyMiddleware(
   })
 
   if (httpServer) {
-    httpServer.on('upgrade', (req, socket, head) => {
+    httpServer.on('upgrade', async (req, socket, head) => {
       const url = req.url!
       for (const context in proxies) {
         if (doesProxyContextMatchUrl(context, url)) {
@@ -169,14 +175,26 @@ export function proxyMiddleware(
             opts.target?.toString().startsWith('wss:')
           ) {
             if (opts.bypass) {
-              const bypassResult = opts.bypass(req, undefined, opts)
-              if (typeof bypassResult === 'string') {
-                req.url = bypassResult
-                debug?.(`bypass: ${req.url} -> ${bypassResult}`)
-                return
-              } else if (bypassResult === false) {
-                debug?.(`bypass: ${req.url} -> 404`)
-                socket.end('HTTP/1.1 404 Not Found\r\n\r\n', '')
+              try {
+                const bypassResult = await opts.bypass(req, undefined, opts)
+                if (typeof bypassResult === 'string') {
+                  debug?.(`bypass: ${req.url} -> ${bypassResult}`)
+                  req.url = bypassResult
+                  return
+                }
+                if (bypassResult === false) {
+                  debug?.(`bypass: ${req.url} -> 404`)
+                  socket.end('HTTP/1.1 404 Not Found\r\n\r\n', '')
+                  return
+                }
+              } catch (err) {
+                config.logger.error(
+                  `${colors.red(`ws proxy bypass error:`)}\n${err.stack}`,
+                  {
+                    timestamp: true,
+                    error: err,
+                  },
+                )
                 return
               }
             }
@@ -194,23 +212,29 @@ export function proxyMiddleware(
   }
 
   // Keep the named function. The name is visible in debug logs via `DEBUG=connect:dispatcher ...`
-  return function viteProxyMiddleware(req, res, next) {
+  return async function viteProxyMiddleware(req, res, next) {
     const url = req.url!
     for (const context in proxies) {
       if (doesProxyContextMatchUrl(context, url)) {
         const [proxy, opts] = proxies[context]
         const options: HttpProxy.ServerOptions = {}
 
         if (opts.bypass) {
-          const bypassResult = opts.bypass(req, res, opts)
-          if (typeof bypassResult === 'string') {
-            req.url = bypassResult
-            debug?.(`bypass: ${req.url} -> ${bypassResult}`)
-            return next()
-          } else if (bypassResult === false) {
-            debug?.(`bypass: ${req.url} -> 404`)
-            res.statusCode = 404
-            return res.end()
+          try {
+            const bypassResult = await opts.bypass(req, res, opts)
+            if (typeof bypassResult === 'string') {
+              debug?.(`bypass: ${req.url} -> ${bypassResult}`)
+              req.url = bypassResult
+              return next()
+            }
+            if (bypassResult === false) {
+              debug?.(`bypass: ${req.url} -> 404`)
+              res.statusCode = 404
+              return res.end()
+            }
+          } catch (e) {
+            debug?.(`bypass: ${req.url} -> ${e}`)
+            return next(e)
           }
         }
 

playground/proxy-bypass/__tests__/proxy-bypass.spec.ts

@@ -1,8 +1,19 @@
 import { expect, test, vi } from 'vitest'
-import { browserLogs } from '~utils'
+import { browserLogs, isServe, page, serverLogs } from '~utils'
 
 test('proxy-bypass', async () => {
   await vi.waitFor(() => {
     expect(browserLogs.join('\n')).toContain('status of 404 (Not Found)')
   })
 })
+
+test('async-proxy-bypass', async () => {
+  const content = await page.frame('async-response').content()
+  expect(content).toContain('Hello after 4 ms (async timeout)')
+})
+
+test.runIf(isServe)('async-proxy-bypass-with-error', async () => {
+  await vi.waitFor(() => {
+    expect(serverLogs.join('\n')).toContain('bypass error')
+  })
+})

playground/proxy-bypass/index.html

@@ -1,2 +1,4 @@
 root app<br />
 <iframe src="/nonExistentApp" style="border: 0"></iframe>
+<iframe src="/asyncResponse" name="async-response"></iframe>
+<iframe src="/asyncThrowingError"></iframe>

playground/proxy-bypass/vite.config.js

@@ -1,5 +1,7 @@
 import { defineConfig } from 'vite'
 
+const timeout = (ms) => new Promise((r) => setTimeout(r, ms))
+
 export default defineConfig({
   server: {
     port: 9606,
@@ -10,6 +12,22 @@ export default defineConfig({
           return false
         },
       },
+      '/asyncResponse': {
+        bypass: async (_, res) => {
+          await timeout(4)
+          res.writeHead(200, {
+            'Content-Type': 'text/plain',
+          })
+          res.end('Hello after 4 ms (async timeout)')
+          return '/asyncResponse'
+        },
+      },
+      '/asyncThrowingError': {
+        bypass: async () => {
+          await timeout(4)
+          throw new Error('bypass error')
+        },
+      },
     },
   },
 })