Add render prop support to <Main> (#30156)

Adds support for render props to the `<Main>` component, when using the [functional custom `Document`](#28515) style. This allows you to write something like this: ```tsx export default function Document() { const jsxStyleRegistry = createStyleRegistry() return ( <Html> <Head /> <body> <Main> {content => ( <StyledJsxWrapper registry={jsxStyleRegistry}> {content} </StyledJsxWrapper> )} </Main> <NextScript /> </body> </Html> ) } ``` In functional document components, this allows the `<App>` to be wrapped, similar to `enhanceApp` (which is only available via `getInitialProps`, which is not supported by functional document components). The primary use for this is for integrating with 3rd party CSS-in-JS libraries, allowing them to attach an `useFlush` handler to [support React 18](reactwg/react-18#110): ```tsx import { unstable_useFlush as useFlush } from 'next/document' export default function StyledJsxWrapper({ children, registry }) { useFlush(() => { /* ... */ }) return ( <StyleRegistry registry={registry}> {children} </StyleRegistry> ) } ``` Support for `useFlush` will be added in a follow up PR.
vercel · Nov 7, 2021 · ad98178 · ad98178
1 parent f8e6661
commit ad98178
Show file tree

Hide file tree

Showing 8 changed files with 122 additions and 34 deletions.
diff --git a/packages/next/pages/_document.tsx b/packages/next/pages/_document.tsx
@@ -1,8 +1,5 @@
 import React, { Component, ReactElement, ReactNode, useContext } from 'react'
-import {
-  BODY_RENDER_TARGET,
-  OPTIMIZED_FONT_PROVIDERS,
-} from '../shared/lib/constants'
+import { OPTIMIZED_FONT_PROVIDERS } from '../shared/lib/constants'
 import {
   DocumentContext,
   DocumentInitialProps,
@@ -763,13 +760,18 @@ export class Head extends Component<
   }
 }
 
-export function Main() {
-  const { inAmpMode, docComponentsRendered } = useContext(HtmlContext)
-
+export function Main({
+  children,
+}: {
+  children?: (content: JSX.Element) => JSX.Element
+}) {
+  const { inAmpMode, docComponentsRendered, useMainContent } =
+    useContext(HtmlContext)
+  const content = useMainContent(children)
   docComponentsRendered.Main = true
 
-  if (inAmpMode) return <>{BODY_RENDER_TARGET}</>
-  return <div id="__next">{BODY_RENDER_TARGET}</div>
+  if (inAmpMode) return content
+  return <div id="__next">{content}</div>
 }
 
 export class NextScript extends Component<OriginProps> {

diff --git a/packages/next/server/render.tsx b/packages/next/server/render.tsx
@@ -20,7 +20,6 @@ import { GetServerSideProps, GetStaticProps, PreviewData } from '../types'
 import { isInAmpMode } from '../shared/lib/amp'
 import { AmpStateContext } from '../shared/lib/amp-context'
 import {
-  BODY_RENDER_TARGET,
   SERVER_PROPS_ID,
   STATIC_PROPS_ID,
   STATIC_STATUS_PAGES,
@@ -935,6 +934,20 @@ export async function renderToHTML(
     }
   }
 
+  const appWrappers: Array<(content: JSX.Element) => JSX.Element> = []
+  const getWrappedApp = (app: JSX.Element) => {
+    // Prevent wrappers from reading/writing props by rendering inside an
+    // opaque component. Wrappers should use context instead.
+    const InnerApp = () => app
+    return (
+      <AppContainer>
+        {appWrappers.reduce((innerContent, fn) => {
+          return fn(innerContent)
+        }, <InnerApp />)}
+      </AppContainer>
+    )
+  }
+
   /**
    * Rules of Static & Dynamic HTML:
    *
@@ -976,13 +989,13 @@ export async function renderToHTML(
           enhanceComponents(options, App, Component)
 
         const html = ReactDOMServer.renderToString(
-          <AppContainer>
+          getWrappedApp(
             <EnhancedApp
               Component={EnhancedComponent}
               router={router}
               {...props}
             />
-          </AppContainer>
+          )
         )
         return { html, head }
       }
@@ -1002,33 +1015,51 @@ export async function renderToHTML(
       }
 
       return {
-        bodyResult: piperFromArray([docProps.html]),
+        bodyResult: () => piperFromArray([docProps.html]),
         documentElement: (htmlProps: HtmlProps) => (
           <Document {...htmlProps} {...docProps} />
         ),
+        useMainContent: (fn?: (content: JSX.Element) => JSX.Element) => {
+          if (fn) {
+            throw new Error(
+              'The `children` property is not supported by non-functional custom Document components'
+            )
+          }
+          // @ts-ignore
+          return <next-js-internal-body-render-target />
+        },
         head: docProps.head,
         headTags: await headTags(documentCtx),
         styles: docProps.styles,
       }
     } else {
-      const content =
-        ctx.err && ErrorDebug ? (
-          <ErrorDebug error={ctx.err} />
-        ) : (
-          <AppContainer>
-            <App {...props} Component={Component} router={router} />
-          </AppContainer>
-        )
+      const bodyResult = async () => {
+        const content =
+          ctx.err && ErrorDebug ? (
+            <ErrorDebug error={ctx.err} />
+          ) : (
+            getWrappedApp(
+              <App {...props} Component={Component} router={router} />
+            )
+          )
 
-      const bodyResult = concurrentFeatures
-        ? process.browser
-          ? await renderToReadableStream(content)
-          : await renderToNodeStream(content, generateStaticHTML)
-        : piperFromArray([ReactDOMServer.renderToString(content)])
+        return concurrentFeatures
+          ? process.browser
+            ? await renderToReadableStream(content)
+            : await renderToNodeStream(content, generateStaticHTML)
+          : piperFromArray([ReactDOMServer.renderToString(content)])
+      }
 
       return {
         bodyResult,
         documentElement: () => (Document as any)(),
+        useMainContent: (fn?: (content: JSX.Element) => JSX.Element) => {
+          if (fn) {
+            appWrappers.push(fn)
+          }
+          // @ts-ignore
+          return <next-js-internal-body-render-target />
+        },
         head,
         headTags: [],
         styles: jsxStyleRegistry.styles(),
@@ -1056,8 +1087,8 @@ export async function renderToHTML(
   }
 
   const hybridAmp = ampState.hybrid
-
   const docComponentsRendered: DocumentProps['docComponentsRendered'] = {}
+
   const {
     assetPrefix,
     buildId,
@@ -1123,6 +1154,7 @@ export async function renderToHTML(
     head: documentResult.head,
     headTags: documentResult.headTags,
     styles: documentResult.styles,
+    useMainContent: documentResult.useMainContent,
     useMaybeDeferContent,
   }
 
@@ -1181,20 +1213,20 @@ export async function renderToHTML(
     }
   }
 
-  const renderTargetIdx = documentHTML.indexOf(BODY_RENDER_TARGET)
+  const [renderTargetPrefix, renderTargetSuffix] = documentHTML.split(
+    /<next-js-internal-body-render-target><\/next-js-internal-body-render-target>/
+  )
   const prefix: Array<string> = []
   prefix.push('<!DOCTYPE html>')
-  prefix.push(documentHTML.substring(0, renderTargetIdx))
+  prefix.push(renderTargetPrefix)
   if (inAmpMode) {
     prefix.push('<!-- __NEXT_DATA__ -->')
   }
 
   let pipers: Array<NodeWritablePiper> = [
     piperFromArray(prefix),
-    documentResult.bodyResult,
-    piperFromArray([
-      documentHTML.substring(renderTargetIdx + BODY_RENDER_TARGET.length),
-    ]),
+    await documentResult.bodyResult(),
+    piperFromArray([renderTargetSuffix]),
   ]
 
   const postProcessors: Array<((html: string) => Promise<string>) | null> = (

diff --git a/packages/next/shared/lib/constants.ts b/packages/next/shared/lib/constants.ts
@@ -23,7 +23,6 @@ export const BLOCKED_PAGES = ['/_document', '/_app', '/_error']
 export const CLIENT_PUBLIC_FILES_PATH = 'public'
 export const CLIENT_STATIC_FILES_PATH = 'static'
 export const CLIENT_STATIC_FILES_RUNTIME = 'runtime'
-export const BODY_RENDER_TARGET = '__NEXT_BODY_RENDER_TARGET__'
 export const STRING_LITERAL_DROP_BUNDLE = '__NEXT_DROP_CLIENT_FILE__'
 
 // server/middleware-flight-manifest.js

diff --git a/packages/next/shared/lib/utils.ts b/packages/next/shared/lib/utils.ts
@@ -221,6 +221,7 @@ export type HtmlProps = {
   styles?: React.ReactElement[] | React.ReactFragment
   head?: Array<JSX.Element | null>
   useMaybeDeferContent: MaybeDeferContentHook
+  useMainContent: (fn?: (content: JSX.Element) => JSX.Element) => JSX.Element
 }
 
 /**

diff --git a/test/integration/document-functional-render-prop/lib/context.js b/test/integration/document-functional-render-prop/lib/context.js
@@ -0,0 +1,3 @@
+import { createContext } from 'react'
+
+export default createContext(null)
diff --git a/test/integration/document-functional-render-prop/pages/_document.js b/test/integration/document-functional-render-prop/pages/_document.js
@@ -0,0 +1,20 @@
+import { Html, Head, Main, NextScript } from 'next/document'
+import Context from '../lib/context'
+
+export default function Document() {
+  return (
+    <Html>
+      <Head />
+      <body>
+        <Main>
+          {(children) => (
+            <Context.Provider value="from render prop">
+              {children}
+            </Context.Provider>
+          )}
+        </Main>
+        <NextScript />
+      </body>
+    </Html>
+  )
+}
diff --git a/test/integration/document-functional-render-prop/pages/index.js b/test/integration/document-functional-render-prop/pages/index.js
@@ -0,0 +1,7 @@
+import { useContext } from 'react'
+import Context from '../lib/context'
+
+export default function MainRenderProp() {
+  const value = useContext(Context)
+  return <span>{value}</span>
+}
diff --git a/test/integration/document-functional-render-prop/tests/index.test.js b/test/integration/document-functional-render-prop/tests/index.test.js
@@ -0,0 +1,24 @@
+/* eslint-env jest */
+
+import { join } from 'path'
+import { findPort, launchApp, killApp, renderViaHTTP } from 'next-test-utils'
+
+const appDir = join(__dirname, '..')
+let appPort
+let app
+
+describe('Functional Custom Document', () => {
+  describe('development mode', () => {
+    beforeAll(async () => {
+      appPort = await findPort()
+      app = await launchApp(appDir, appPort)
+    })
+
+    afterAll(() => killApp(app))
+
+    it('supports render props', async () => {
+      const html = await renderViaHTTP(appPort, '/')
+      expect(html).toMatch(/<span>from render prop<\/span>/)
+    })
+  })
+})