From a4c293493c87ee53e7289b53babdee5e9817c1f6 Mon Sep 17 00:00:00 2001
From: Shu Ding <g@shud.in>
Date: Fri, 5 Nov 2021 23:11:07 +0100
Subject: [PATCH] fix uncaught errors in middleware

---
 .../next-middleware-ssr-loader/index.ts       | 36 +++++++++----------
 .../test/index.test.js                        |  2 +-
 2 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/packages/next/build/webpack/loaders/next-middleware-ssr-loader/index.ts b/packages/next/build/webpack/loaders/next-middleware-ssr-loader/index.ts
index 1f367ac3eadebd1..07f6dfd0861826f 100644
--- a/packages/next/build/webpack/loaders/next-middleware-ssr-loader/index.ts
+++ b/packages/next/build/webpack/loaders/next-middleware-ssr-loader/index.ts
@@ -57,10 +57,6 @@ export default async function middlewareRSCLoader(this: any) {
           throw new Error('Your page must export a \`default\` component')
         }
 
-        function renderError(err, status) {
-          return new Response(err.toString(), {status})
-        }
-
         function wrapReadable(readable) {
           const encoder = new TextEncoder()
           const transformStream = new TransformStream()
@@ -104,15 +100,10 @@ export default async function middlewareRSCLoader(this: any) {
         const Component = Page`
         }
 
-        function render(request) {
+        async function render(request) {
           const url = request.nextUrl
           const query = Object.fromEntries(url.searchParams)
 
-          if (Document.getInitialProps) {
-            const err = new Error('Document.getInitialProps is not supported with server components, please remove it from pages/_document')
-            return renderError(err, 500)
-          }
-
           // Preflight request
           if (request.method === 'HEAD') {
             return new Response('OK.', {
@@ -172,18 +163,27 @@ export default async function middlewareRSCLoader(this: any) {
           const writer = transformStream.writable.getWriter()
           const encoder = new TextEncoder()
 
-          renderToHTML(
-            { url: url.pathname },
-            {},
-            url.pathname,
-            query,
-            renderOpts
-          ).then(result => {
+          try {
+            const result = await renderToHTML(
+              { url: url.pathname },
+              {},
+              url.pathname,
+              query,
+              renderOpts
+            )
             result.pipe({
               write: str => writer.write(encoder.encode(str)),
               end: () => writer.close()
             })
-          })
+          } catch (err) {
+            return new Response(
+              (err || 'An error occurred while rendering ' + url.pathname + '.').toString(),
+              {
+                status: 500,
+                headers: { 'x-middleware-ssr': '1' }
+              }
+            )
+          }
 
           return new Response(transformStream.readable, {
             headers: { 'x-middleware-ssr': '1' }
diff --git a/test/integration/react-streaming-and-server-components/test/index.test.js b/test/integration/react-streaming-and-server-components/test/index.test.js
index 613c90d350d5e74..c8f5e4999bb21e1 100644
--- a/test/integration/react-streaming-and-server-components/test/index.test.js
+++ b/test/integration/react-streaming-and-server-components/test/index.test.js
@@ -190,7 +190,7 @@ const documentSuite = {
 
       expect(res.status).toBe(500)
       expect(html).toContain(
-        'Document.getInitialProps is not supported with server components, please remove it from pages/_document'
+        'Error: `getInitialProps` in Document component is not supported with `concurrentFeatures` enabled.'
       )
     })
   },