fix: added server shutdown cleanup to proxy implementation

vercel · May 14, 2024 · 55552ba · 55552ba
1 parent 7f20310
commit 55552ba
Showing 1 changed file with 81 additions and 71 deletions.
diff --git a/packages/next/src/server/lib/router-utils/proxy-request.ts b/packages/next/src/server/lib/router-utils/proxy-request.ts
@@ -4,6 +4,7 @@ import type { NextUrlWithParsedQuery } from '../../request-meta'
 import url from 'url'
 import { stringifyQuery } from '../../server-route-utils'
 import { Duplex } from 'stream'
+import { DetachedPromise } from '../../../lib/detached-promise'
 
 export async function proxyRequest(
   req: IncomingMessage,
@@ -34,86 +35,95 @@ export async function proxyRequest(
     },
   })
 
-  await new Promise((proxyResolve, proxyReject) => {
-    let finished = false
+  let finished = false
 
-    // http-proxy does not properly detect a client disconnect in newer
-    // versions of Node.js. This is caused because it only listens for the
-    // `aborted` event on the our request object, but it also fully reads
-    // and closes the request object. Node **will not** fire `aborted` when
-    // the request is already closed. Listening for `close` on our response
-    // object will detect the disconnect, and we can abort the proxy's
-    // connection.
-    proxy.on('proxyReq', (proxyReq) => {
-      res.on('close', () => proxyReq.destroy())
-    })
-    proxy.on('proxyRes', (proxyRes) => {
-      if (res.destroyed) {
-        proxyRes.destroy()
-      } else {
-        res.on('close', () => proxyRes.destroy())
-      }
-    })
+  // http-proxy does not properly detect a client disconnect in newer
+  // versions of Node.js. This is caused because it only listens for the
+  // `aborted` event on the our request object, but it also fully reads
+  // and closes the request object. Node **will not** fire `aborted` when
+  // the request is already closed. Listening for `close` on our response
+  // object will detect the disconnect, and we can abort the proxy's
+  // connection.
+  proxy.on('proxyReq', (proxyReq) => {
+    res.on('close', () => proxyReq.destroy())
+  })
 
-    proxy.on('proxyRes', (proxyRes, innerReq, innerRes) => {
-      const cleanup = (err: any) => {
-        // cleanup event listeners to allow clean garbage collection
-        proxyRes.removeListener('error', cleanup)
-        proxyRes.removeListener('close', cleanup)
-        innerRes.removeListener('error', cleanup)
-        innerRes.removeListener('close', cleanup)
-
-        // destroy all source streams to propagate the caught event backward
-        innerReq.destroy(err)
-        proxyRes.destroy(err)
-      }
+  proxy.on('proxyRes', (proxyRes) => {
+    if (res.destroyed) {
+      proxyRes.destroy()
+    } else {
+      res.on('close', () => proxyRes.destroy())
+    }
+  })
 
-      proxyRes.once('error', cleanup)
-      proxyRes.once('close', cleanup)
-      innerRes.once('error', cleanup)
-      innerRes.once('close', cleanup)
-    })
+  proxy.on('proxyRes', (proxyRes, innerReq, innerRes) => {
+    const cleanup = (err: any) => {
+      // cleanup event listeners to allow clean garbage collection
+      proxyRes.removeListener('error', cleanup)
+      proxyRes.removeListener('close', cleanup)
+      innerRes.removeListener('error', cleanup)
+      innerRes.removeListener('close', cleanup)
+
+      // destroy all source streams to propagate the caught event backward
+      innerReq.destroy(err)
+      proxyRes.destroy(err)
+    }
 
-    proxy.on('error', (err) => {
-      console.error(`Failed to proxy ${target}`, err)
-      if (!finished) {
-        finished = true
-        proxyReject(err)
+    proxyRes.once('error', cleanup)
+    proxyRes.once('close', cleanup)
+    innerRes.once('error', cleanup)
+    innerRes.once('close', cleanup)
+  })
+
+  const detached = new DetachedPromise<boolean>()
+
+  // When the proxy finishes proxying the request, shut down the proxy.
+  detached.promise.finally(() => {
+    proxy.close()
+  })
 
-        if (!res.destroyed) {
-          if (!(res instanceof Duplex)) {
-            res.statusCode = 500
-          }
+  proxy.on('error', (err) => {
+    console.error(`Failed to proxy ${target}`, err)
+    if (!finished) {
+      finished = true
+      detached.reject(err)
 
-          res.end('Internal Server Error')
+      if (!res.destroyed) {
+        if (!(res instanceof Duplex)) {
+          res.statusCode = 500
         }
+
+        res.end('Internal Server Error')
       }
-    })
+    }
+  })
 
-    // if upgrade head is present treat as WebSocket request
-    if (upgradeHead || res instanceof Duplex) {
-      proxy.on('proxyReqWs', (proxyReq) => {
-        proxyReq.on('close', () => {
-          if (!finished) {
-            finished = true
-            proxyResolve(true)
-          }
-        })
-      })
-      proxy.ws(req, res, upgradeHead)
-      proxyResolve(true)
-    } else {
-      proxy.on('proxyReq', (proxyReq) => {
-        proxyReq.on('close', () => {
-          if (!finished) {
-            finished = true
-            proxyResolve(true)
-          }
-        })
+  // If upgrade head is present or the response is a Duplex stream, treat as
+  // WebSocket request.
+  if (upgradeHead || res instanceof Duplex) {
+    proxy.on('proxyReqWs', (proxyReq) => {
+      proxyReq.on('close', () => {
+        if (!finished) {
+          finished = true
+          detached.resolve(true)
+        }
       })
-      proxy.web(req, res, {
-        buffer: reqBody,
+    })
+    proxy.ws(req, res, upgradeHead)
+    detached.resolve(true)
+  } else {
+    proxy.on('proxyReq', (proxyReq) => {
+      proxyReq.on('close', () => {
+        if (!finished) {
+          finished = true
+          detached.resolve(true)
+        }
       })
-    }
-  })
+    })
+    proxy.web(req, res, {
+      buffer: reqBody,
+    })
+  }
+
+  return detached.promise
 }