refactor: patch the webpack runtime using ast-grep (#351)

Co-authored-by: Dario Piotrowicz <dario@cloudflare.com>

Author: vicb

packages/cloudflare/src/cli/build/bundle-server.ts

@@ -8,6 +8,7 @@ import { getCrossPlatformPathRegex } from "@opennextjs/aws/utils/regex.js";
 import { build, Plugin } from "esbuild";
 
 import { patchVercelOgLibrary } from "./patches/ast/patch-vercel-og-library.js";
+import { patchWebpackRuntime } from "./patches/ast/webpack-runtime.js";
 import * as patches from "./patches/index.js";
 import { handleOptionalDependencies } from "./patches/plugins/optional-deps.js";
 import { fixRequire } from "./patches/plugins/require.js";
@@ -49,7 +50,7 @@ export async function bundleServer(buildOpts: BuildOptions): Promise<void> {
 
   console.log(`\x1b[35m⚙️ Bundling the OpenNext server...\n\x1b[0m`);
 
-  await patches.updateWebpackChunksFile(buildOpts);
+  await patchWebpackRuntime(buildOpts);
   patchVercelOgLibrary(buildOpts);
 
   const outputPath = path.join(outputDir, "server-functions", "default");

packages/cloudflare/src/cli/build/patches/ast/webpack-runtime.spec.ts

@@ -0,0 +1,64 @@
+import { describe, expect, test } from "vitest";
+
+import { patchCode } from "./util.js";
+import { buildInlineChunksRule } from "./webpack-runtime.js";
+
+describe("webpack runtime", () => {
+  test("patch runtime", () => {
+    const code = `
+        /******/ 		// require() chunk loading for javascript
+        /******/ 		__webpack_require__.f.require = (chunkId, promises) => {
+        /******/ 			// "1" is the signal for "already loaded"
+        /******/ 			if (!installedChunks[chunkId]) {
+        /******/ 				if (658 != chunkId) {
+        /******/ 					installChunk(require("./chunks/" + __webpack_require__.u(chunkId)));
+                    /******/
+                } else installedChunks[chunkId] = 1;
+                /******/
+            }
+            /******/
+        };
+    `;
+
+    expect(patchCode(code, buildInlineChunksRule([1, 2, 3]))).toMatchInlineSnapshot(`
+      "/******/ 		// require() chunk loading for javascript
+              /******/ 		__webpack_require__.f.require = (chunkId, _) => {
+        if (!installedChunks[chunkId]) {
+          switch (chunkId) {
+             case 1: installChunk(require("./chunks/1.js")); break;
+             case 2: installChunk(require("./chunks/2.js")); break;
+             case 3: installChunk(require("./chunks/3.js")); break;
+             case 658: installedChunks[chunkId] = 1; break;
+             default: throw new Error(\`Unknown chunk \${chunkId}\`);
+          }
+        }
+      }
+      ;
+          "
+    `);
+  });
+
+  test("patch minified runtime", () => {
+    const code = `
+    t.f.require=(o,n)=>{e[o]||(658!=o?r(require("./chunks/"+t.u(o))):e[o]=1)}
+    `;
+
+    expect(patchCode(code, buildInlineChunksRule([1, 2, 3]))).toMatchInlineSnapshot(
+      `
+      "t.f.require=(o, _) => {
+        if (!e[o]) {
+          switch (o) {
+             case 1: r(require("./chunks/1.js")); break;
+             case 2: r(require("./chunks/2.js")); break;
+             case 3: r(require("./chunks/3.js")); break;
+             case 658: e[o] = 1; break;
+             default: throw new Error(\`Unknown chunk \${o}\`);
+          }
+        }
+      }
+
+          "
+    `
+    );
+  });
+});

packages/cloudflare/src/cli/build/patches/ast/webpack-runtime.ts

@@ -0,0 +1,70 @@
+/**
+ * Inline dynamic requires in the webpack runtime.
+ *
+ * The webpack runtime has dynamic requires that would not be bundled by ESBuild:
+ *
+ *     installChunk(require("./chunks/" + __webpack_require__.u(chunkId)));
+ *
+ *  This patch unrolls the dynamic require for all the existing chunks:
+ *
+ *     switch (chunkId) {
+ *       case ID1: installChunk(require("./chunks/ID1"); break;
+ *       case ID2: installChunk(require("./chunks/ID2"); break;
+ *       // ...
+ *       case SELF_ID: installedChunks[chunkId] = 1; break;
+ *       default: throw new Error(`Unknown chunk ${chunkId}`);
+ *     }
+ */
+
+import { readdirSync, readFileSync, writeFileSync } from "node:fs";
+import { join } from "node:path";
+
+import { type BuildOptions, getPackagePath } from "@opennextjs/aws/build/helper.js";
+
+import { patchCode } from "./util.js";
+
+export function buildInlineChunksRule(chunks: number[]) {
+  return `
+rule:
+  pattern: ($CHUNK_ID, $_PROMISES) => { $$$ }
+  inside: {pattern: $_.$_.require = $$$_, stopBy: end}
+  all:
+    - has: {pattern: $INSTALL(require("./chunks/" + $$$)), stopBy: end}
+    - has: {pattern: $SELF_ID != $CHUNK_ID, stopBy: end}
+    - has: {pattern: "$INSTALLED_CHUNK[$CHUNK_ID] = 1", stopBy: end}
+fix: |
+  ($CHUNK_ID, _) => {
+    if (!$INSTALLED_CHUNK[$CHUNK_ID]) {
+      switch ($CHUNK_ID) {
+${chunks.map((chunk) => `         case ${chunk}: $INSTALL(require("./chunks/${chunk}.js")); break;`).join("\n")}
+         case $SELF_ID: $INSTALLED_CHUNK[$CHUNK_ID] = 1; break;
+         default: throw new Error(\`Unknown chunk \${$CHUNK_ID}\`);
+      }
+    }
+  }`;
+}
+
+/**
+ * Fixes the webpack-runtime.js file by removing its webpack dynamic requires.
+ */
+export async function patchWebpackRuntime(buildOpts: BuildOptions) {
+  const { outputDir } = buildOpts;
+
+  const dotNextServerDir = join(
+    outputDir,
+    "server-functions/default",
+    getPackagePath(buildOpts),
+    ".next/server"
+  );
+
+  const runtimeFile = join(dotNextServerDir, "webpack-runtime.js");
+  const runtimeCode = readFileSync(runtimeFile, "utf-8");
+  // Look for all the chunks.
+  const chunks = readdirSync(join(dotNextServerDir, "chunks"))
+    .filter((chunk) => /^\d+\.js$/.test(chunk))
+    .map((chunk) => {
+      return Number(chunk.replace(/\.js$/, ""));
+    });
+
+  writeFileSync(runtimeFile, patchCode(runtimeCode, buildInlineChunksRule(chunks)));
+}

packages/cloudflare/src/cli/build/patches/investigated/index.ts

@@ -1,4 +1,3 @@
 export * from "./copy-package-cli-files.js";
 export * from "./patch-cache.js";
 export * from "./patch-require.js";
-export * from "./update-webpack-chunks-file/index.js";