[Per-Segment Prefetch] Parent segments must be prefetched before children

packages/next/src/server/app-render/collect-segment-data.tsx

@@ -50,6 +50,11 @@ type TreePrefetch = {
 type SegmentPrefetch = {
   rsc: React.ReactNode | null
   loading: LoadingModuleData
+
+  // Access tokens for the child segments.
+  slots: null | {
+    [parallelRouteKey: string]: string
+  }
 }
 
 export async function collectSegmentData(
@@ -89,6 +94,7 @@ export async function collectSegmentData(
     serverConsumerManifest,
     [],
     '',
+    '',
     resultMap,
     collectedTasks
   )
@@ -118,10 +124,17 @@ function collectSegmentDataImpl(
   serverConsumerManifest: any,
   segmentPath: Array<[string, Segment]>,
   segmentPathStr: string,
+  accessToken: string,
   segmentBufferMap: Map<string, Buffer>,
   collectedTasks: Array<Promise<void>>
 ): TreePrefetch {
-  const slots: { [parallelRouteKey: string]: TreePrefetch } = {}
+  // Metadata about the segment. Sent as part of the tree prefetch. Null if
+  // there are no children.
+  let slotMetadata: { [parallelRouteKey: string]: TreePrefetch } | null = null
+
+  // Access tokens for the child segments. Sent as part of layout's data. Null
+  // if there are no children.
+  let childAccessTokens: { [parallelRouteKey: string]: string } | null = null
 
   const children = route[1]
   for (const parallelRouteKey in children) {
@@ -135,17 +148,31 @@ function collectSegmentDataImpl(
       '/' +
       encodeChildSegmentAsFilesystemSafePathname(parallelRouteKey, childSegment)
 
+    // Create an access token for each child slot
+    const childAccessToken = createSegmentAccessToken(
+      segmentPathStr,
+      parallelRouteKey
+    )
     const childTree = collectSegmentDataImpl(
       childRoute,
       fullPageDataBuffer,
       clientModules,
       serverConsumerManifest,
       childSegmentPath,
       childSegmentPathStr,
+      childAccessToken,
       segmentBufferMap,
       collectedTasks
     )
-    slots[parallelRouteKey] = childTree
+    if (slotMetadata === null) {
+      slotMetadata = {}
+    }
+    slotMetadata[parallelRouteKey] = childTree
+
+    if (childAccessTokens === null) {
+      childAccessTokens = {}
+    }
+    childAccessTokens[parallelRouteKey] = childAccessToken
   }
 
   // Spawn a task to render the segment data to a stream.
@@ -156,6 +183,8 @@ function collectSegmentDataImpl(
       serverConsumerManifest,
       segmentPath,
       segmentPathStr,
+      accessToken,
+      childAccessTokens,
       segmentBufferMap
     )
   )
@@ -165,8 +194,8 @@ function collectSegmentDataImpl(
   const segment = route[0]
   const isRootLayout = route[4]
   return {
-    key: segmentPathStr,
-    slots,
+    key: segmentPathStr === '' ? '/' : segmentPathStr,
+    slots: slotMetadata,
     extra: [segment, isRootLayout === true],
   }
 }
@@ -177,6 +206,8 @@ async function renderSegmentDataToStream(
   serverConsumerManifest: any,
   segmentPath: Array<[string, Segment]>,
   segmentPathStr: string,
+  accessToken: string,
+  childAccessTokens: { [parallelRouteKey: string]: string } | null,
   segmentBufferMap: Map<string, Buffer>
 ) {
   // Create a new Flight response that contains data only for this segment.
@@ -196,6 +227,7 @@ async function renderSegmentDataToStream(
         fullPageDataBuffer={fullPageDataBuffer}
         serverConsumerManifest={serverConsumerManifest}
         segmentPath={segmentPath}
+        childAccessTokens={childAccessTokens}
       />,
       clientModules,
       {
@@ -211,7 +243,15 @@ async function renderSegmentDataToStream(
     if (segmentPathStr === '') {
       segmentBufferMap.set('/', segmentBuffer)
     } else {
-      segmentBufferMap.set(segmentPathStr, segmentBuffer)
+      // The access token is appended to the end of the segment name. To request
+      // a segment, the client sends a header like:
+      //
+      //   Next-Router-Segment-Prefetch: /path/to/segment.accesstoken
+      //
+      // The segment path is provided by the tree prefetch, and the access
+      // token is provided in the parent layout's data.
+      const fullPath = `${segmentPathStr}.${accessToken}`
+      segmentBufferMap.set(fullPath, segmentBuffer)
     }
   } catch {
     // If there are any errors, then we skip the segment. The effect is that
@@ -223,10 +263,12 @@ async function PickSegment({
   fullPageDataBuffer,
   serverConsumerManifest,
   segmentPath,
+  childAccessTokens,
 }: {
   fullPageDataBuffer: Buffer
   serverConsumerManifest: any
   segmentPath: Array<[string, Segment]>
+  childAccessTokens: { [parallelRouteKey: string]: string } | null
 }): Promise<SegmentPrefetch | null> {
   // We're currently rendering a Flight response for a segment prefetch.
   // Decode the Flight stream for the whole page, then pick out the data for the
@@ -288,6 +330,7 @@ async function PickSegment({
   return {
     rsc,
     loading,
+    slots: childAccessTokens,
   }
 }
 
@@ -370,3 +413,26 @@ function encodeParamValue(segment: string): string {
   // We also add a $ prefix so it doesn't collide with the simple case.
   return '$' + Buffer.from(segment, 'utf-8').toString('base64url')
 }
+
+function createSegmentAccessToken(
+  parentSegmentPathStr: string,
+  parallelRouteKey: string
+): string {
+  // Create an access token that the client passes when requesting a segment.
+  // The token is sent to the client as part of the parent layout's data.
+  //
+  // The token is hash of the parent segment path and the parallel route key. A
+  // subtle detail here is that it does *not* include the value of the segment
+  // itself — a shared layout must produce the same access tokens for its
+  // children regardless of their segment values, so that the client only has to
+  // fetch the layout once.
+  //
+  // TODO: Because this only affects prefetches, this doesn't need to be secure.
+  // It's just for obfuscation. But eventually we will use this technique when
+  // performing dynamic navigations, to support auth checks in a layout that
+  // conditionally renders its slots. At that point we'll need to create an
+  // actual cryptographic hash with a salt.
+  return Buffer.from(parentSegmentPathStr + parallelRouteKey, 'utf-8')
+    .toString('hex')
+    .slice(0, 7)
+}

test/e2e/app-dir/ppr-navigations/simple/per-segment-prefetching.test.ts

@@ -14,44 +14,62 @@ describe('per segment prefetching', () => {
   // these types of requests. This tests that the server responds correctly.
   // TODO: Replace with e2e tests once more is implemented.
 
-  it('prefetch a route tree', async () => {
-    const response = await next.fetch('/en', {
+  function prefetch(pageUrl, segmentPath) {
+    return next.fetch(pageUrl, {
       headers: {
         RSC: '1',
         'Next-Router-Prefetch': '1',
-        'Next-Router-Segment-Prefetch': '/_tree',
+        'Next-Router-Segment-Prefetch': segmentPath,
       },
     })
-    expect(response.status).toBe(200)
-    const responseText = await response.text()
-    // This is a basic check to ensure that the name of an expected field is
-    // somewhere in the Flight stream.
-    expect(responseText).toInclude('"staleTime"')
-  })
+  }
 
-  it('prefetch an individual segment', async () => {
-    const response = await next.fetch('/en', {
-      headers: {
-        RSC: '1',
-        'Next-Router-Prefetch': '1',
-        'Next-Router-Segment-Prefetch': '/',
-      },
-    })
-    expect(response.status).toBe(200)
-    const responseText = await response.text()
-    // This is a basic check to ensure that the name of an expected field is
-    // somewhere in the Flight stream.
-    expect(responseText).toInclude('"rsc"')
+  function extractPseudoJSONFromFlightResponse(flightText: string) {
+    // This is a cheat that takes advantage of the fact that the roots of the
+    // Flight responses in this test are JSON. Again, this is just a temporary
+    // smoke test until the client part is implemented; we shouldn't rely on
+    // this as a general testing strategy.
+    const match = flightText.match(/^0:(.*)$/m)
+    if (match) {
+      return JSON.parse(match[1])
+    }
+    return null
+  }
+
+  it('basic prefetching flow', async () => {
+    // To perform a prefetch a page, the client first fetches the route tree.
+    // The response is used to construct prefetches of individual segments.
+    const routeTreeResponse = await prefetch('/en', '/_tree')
+    const routeTreeResponseText = await routeTreeResponse.text()
+    const routeTree = extractPseudoJSONFromFlightResponse(routeTreeResponseText)
+
+    // The root segment is a shared segment. Demonstrate that fetching the root
+    // segment for two different pages results in the same response.
+    const enResponse = await prefetch('/en', '/')
+    const enResponseText = await enResponse.text()
+    const frResponse = await prefetch('/fr', '/')
+    const frResponseText = await frResponse.text()
+    expect(enResponseText).toEqual(frResponseText)
+
+    // Now use both the tree response and the root segment data to construct a
+    // request for the child segment.
+    const childSegmentPath = routeTree.tree.slots.children.key
+    const childToken =
+      extractPseudoJSONFromFlightResponse(enResponseText).slots.children
+
+    // The access token, which we extracted from the response for its parent
+    // segment, is appended to the end of the segment path.
+    const fullChildSegmentPath = `${childSegmentPath}.${childToken}`
+    const childResponse = await prefetch('/en', fullChildSegmentPath)
+    const childResponseText = await childResponse.text()
+
+    // Confirm that the prefetch was successful. This is a basic check to ensure
+    // that the name of an expected field is somewhere in the Flight stream.
+    expect(childResponseText).toInclude('"rsc"')
   })
 
   it('respond with 404 if the segment does not have prefetch data', async () => {
-    const response = await next.fetch('/en', {
-      headers: {
-        RSC: '1',
-        'Next-Router-Prefetch': '1',
-        'Next-Router-Segment-Prefetch': '/does-not-exist',
-      },
-    })
+    const response = await prefetch('/en', '/does-not-exist')
     expect(response.status).toBe(404)
     const responseText = await response.text()
     expect(responseText.trim()).toBe('')