Output caching headers on RSC/prefetch responses (#2559)

Author: SamyPesse

bun.lockb

@@ -8,7 +8,8 @@
     },
     "packageManager": "bun@1.1.18",
     "patchedDependencies": {
-        "@vercel/next@4.3.15": "patches/@vercel%2Fnext@4.3.15.patch"
+        "@vercel/next@4.3.15": "patches/@vercel%2Fnext@4.3.15.patch",
+        "@cloudflare/next-on-pages@1.13.5": "patches/@cloudflare%2Fnext-on-pages@1.13.5.patch"
     },
     "private": true,
     "scripts": {

package.json

@@ -5,7 +5,7 @@
     "scripts": {
         "dev": "env-cmd --silent -f ../../.env.local next dev",
         "build": "next build",
-        "build:cloudflare": "next-on-pages",
+        "build:cloudflare": "next-on-pages --custom-entrypoint=./src/cloudflare-entrypoint.ts",
         "start": "next start",
         "lint": "next lint",
         "typecheck": "tsc --noEmit",

packages/gitbook/package.json

@@ -0,0 +1,18 @@
+// @ts-ignore
+import nextOnPagesHandler from '@cloudflare/next-on-pages/fetch-handler';
+
+import { withMiddlewareHeadersStorage } from './lib/middleware';
+
+/**
+ * We use a custom entrypoint until we can move to opennext (https://github.com/opennextjs/opennextjs-cloudflare/issues/92).
+ * There is a bug in next-on-pages where headers can't be set on the response in the middleware for RSC requests (https://github.com/cloudflare/next-on-pages/issues/897).
+ */
+export default {
+    async fetch(request, env, ctx) {
+        const response = await withMiddlewareHeadersStorage(() =>
+            nextOnPagesHandler.fetch(request, env, ctx),
+        );
+
+        return response;
+    },
+} as ExportedHandler<{ ASSETS: Fetcher }>;

packages/gitbook/src/cloudflare-entrypoint.ts

@@ -1,3 +1,45 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
+
+/**
+ * Set a header on the middleware response.
+ * We do this because of https://github.com/opennextjs/opennextjs-cloudflare/issues/92
+ * It can be removed as soon as we move to opennext where hopefully this is fixed.
+ */
+export function setMiddlewareHeader(response: Response, name: string, value: string) {
+    const responseHeadersLocalStorage =
+        // @ts-ignore
+        globalThis.responseHeadersLocalStorage as AsyncLocalStorage<Headers> | undefined;
+    const responseHeaders = responseHeadersLocalStorage?.getStore();
+    response.headers.set(name, value);
+
+    if (responseHeaders) {
+        responseHeaders.set(name, value);
+    }
+}
+
+/**
+ * Wrap some middleware with a the storage to store headers.
+ */
+export async function withMiddlewareHeadersStorage(
+    handler: () => Promise<Response>,
+): Promise<Response> {
+    const responseHeadersLocalStorage =
+        // @ts-ignore
+        (globalThis.responseHeadersLocalStorage as AsyncLocalStorage<Headers>) ??
+        new AsyncLocalStorage<Headers>();
+    // @ts-ignore
+    globalThis.responseHeadersLocalStorage = responseHeadersLocalStorage;
+
+    const responseHeaders = new Headers();
+    const response = await responseHeadersLocalStorage.run(responseHeaders, handler);
+
+    for (const [name, value] of responseHeaders.entries()) {
+        response.headers.set(name, value);
+    }
+
+    return response;
+}
+
 /**
  * For a given GitBook URL, return a list of alternative URLs that could be matched against to lookup the content.
  * The approach is optimized to aim at reusing cached lookup results as much as possible.

packages/gitbook/src/lib/middleware.ts

@@ -20,7 +20,7 @@ import {
 import { race } from '@/lib/async';
 import { buildVersion } from '@/lib/build';
 import { createContentSecurityPolicyNonce, getContentSecurityPolicy } from '@/lib/csp';
-import { getURLLookupAlternatives, normalizeURL } from '@/lib/middleware';
+import { getURLLookupAlternatives, normalizeURL, setMiddlewareHeader } from '@/lib/middleware';
 import {
     VisitorAuthCookieValue,
     getVisitorAuthCookieName,
@@ -253,43 +253,31 @@ export async function middleware(request: NextRequest) {
         resolved.cookies,
     );
 
-    response.headers.set('x-gitbook-version', buildVersion());
+    setMiddlewareHeader(response, 'x-gitbook-version', buildVersion());
 
     // Add Content Security Policy header
-    response.headers.set('content-security-policy', csp);
+    setMiddlewareHeader(response, 'content-security-policy', csp);
     // Basic security headers
-    response.headers.set('strict-transport-security', 'max-age=31536000');
-    response.headers.set('referrer-policy', 'no-referrer-when-downgrade');
-    response.headers.set('x-content-type-options', 'nosniff');
-
-    const isPrefetch = request.headers.has('x-middleware-prefetch');
-
-    if (isPrefetch) {
-        // To avoid cache poisoning, we don't cache prefetch requests
-        response.headers.set(
-            'cache-control',
-            'private, no-cache, no-store, max-age=0, must-revalidate',
-        );
-    } else {
-        if (typeof resolved.cacheMaxAge === 'number') {
-            const cacheControl = `public, max-age=0, s-maxage=${resolved.cacheMaxAge}, stale-if-error=0`;
-
-            if (
-                process.env.GITBOOK_OUTPUT_CACHE === 'true' &&
-                process.env.NODE_ENV !== 'development'
-            ) {
-                response.headers.set('cache-control', cacheControl);
-                response.headers.set('Cloudflare-CDN-Cache-Control', cacheControl);
-            } else {
-                response.headers.set('x-gitbook-cache-control', cacheControl);
-            }
+    setMiddlewareHeader(response, 'strict-transport-security', 'max-age=31536000');
+    setMiddlewareHeader(response, 'referrer-policy', 'no-referrer-when-downgrade');
+    setMiddlewareHeader(response, 'x-content-type-options', 'nosniff');
+
+    if (typeof resolved.cacheMaxAge === 'number') {
+        const cacheControl = `public, max-age=0, s-maxage=${resolved.cacheMaxAge}, stale-if-error=0`;
+
+        if (process.env.GITBOOK_OUTPUT_CACHE === 'true' && process.env.NODE_ENV !== 'development') {
+            setMiddlewareHeader(response, 'cache-control', cacheControl);
+            setMiddlewareHeader(response, 'Cloudflare-CDN-Cache-Control', cacheControl);
+        } else {
+            setMiddlewareHeader(response, 'x-gitbook-cache-control', cacheControl);
         }
+    }
+    // }
 
-        if (resolved.cacheTags && resolved.cacheTags.length > 0) {
-            const headerCacheTag = resolved.cacheTags.join(',');
-            response.headers.set('cache-tag', headerCacheTag);
-            response.headers.set('x-gitbook-cache-tag', headerCacheTag);
-        }
+    if (resolved.cacheTags && resolved.cacheTags.length > 0) {
+        const headerCacheTag = resolved.cacheTags.join(',');
+        setMiddlewareHeader(response, 'cache-tag', headerCacheTag);
+        setMiddlewareHeader(response, 'x-gitbook-cache-tag', headerCacheTag);
     }
 
     return response;

packages/gitbook/src/middleware.ts

@@ -0,0 +1,12 @@
+diff --git a/dist/index.js b/dist/index.js
+index 32fec63484ec332eb291a7253e5e168223627535..653dee64794140bafe57219356c712b197c17530 100644
+--- a/dist/index.js
++++ b/dist/index.js
+@@ -6983,6 +6983,7 @@ async function buildWorkerFile({ vercelConfig, vercelOutput }, {
+       outfile: outputFile,
+       allowOverwrite: true,
+       bundle: true,
++      external: ["node:*", "cloudflare:*"],
+       plugins: [
+         {
+           name: "custom-entrypoint-import-plugin",