From b75b2f02c9dc33c649c8d90c952be9bfee70c60b Mon Sep 17 00:00:00 2001
From: Shu Ding <g@shud.in>
Date: Fri, 5 Nov 2021 23:59:46 +0100
Subject: [PATCH] Improve error handling in the SSR middleware (#31057)

This PR improves error handling in the SSR middleware. Previously the response was sent out synchronously, and and errors were silently swallowed. There was no `.catch` for `renderToHTML`. This changes the middleware to be asynchronous, which waits until the initial Document to be rendered correctly and then starts the streaming.

With this change we can also send correct status code when there're immediate errors before Fizz.

## Bug

- [ ] Related issues linked using `fixes #number`
- [ ] Integration tests added
- [ ] Errors have helpful link attached, see `contributing.md`

## Feature

- [ ] Implements an existing feature request or RFC. Make sure the feature request has been accepted for implementation before opening a PR.
- [ ] Related issues linked using `fixes #number`
- [ ] Integration tests added
- [ ] Documentation added
- [ ] Telemetry added. In case of a feature if it's used or not.
- [ ] Errors have helpful link attached, see `contributing.md`

## Documentation / Examples

- [ ] Make sure the linting passes by running `yarn lint`
---
 .../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.'
       )
     })
   },