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

…dren (vercel#72342)

Based on:
- vercel#72341 

---

Creates 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.

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 add a
salt.

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('')