Feat: Generic Code Patching (#781)

* poc for code patching

* review fix

* fix build

* use ContentUpdater plugin

* added ast-grep code patcher

* added some manifests

* fix and add logs

* fix to apply patch one after the other

* review fix

* improve typing

* simplified CodePatcher types

* fix linting

* changed versions to be more user friendly

* review

* review fix

Author: conico974

packages/open-next/package.json

@@ -37,6 +37,7 @@
     "README.md"
   ],
   "dependencies": {
+    "@ast-grep/napi": "^0.35.0",
     "@aws-sdk/client-cloudfront": "3.398.0",
     "@aws-sdk/client-dynamodb": "^3.398.0",
     "@aws-sdk/client-lambda": "^3.398.0",
@@ -50,7 +51,8 @@
     "esbuild": "0.19.2",
     "express": "5.0.1",
     "path-to-regexp": "^6.3.0",
-    "urlpattern-polyfill": "^10.0.0"
+    "urlpattern-polyfill": "^10.0.0",
+    "yaml": "^2.7.0"
   },
   "devDependencies": {
     "@types/aws-lambda": "^8.10.109",

packages/open-next/src/adapters/config/util.ts

@@ -21,10 +21,14 @@ export function loadBuildId(nextDir: string) {
   return fs.readFileSync(filePath, "utf-8").trim();
 }
 
-export function loadHtmlPages(nextDir: string) {
+export function loadPagesManifest(nextDir: string) {
   const filePath = path.join(nextDir, "server/pages-manifest.json");
   const json = fs.readFileSync(filePath, "utf-8");
-  return Object.entries(JSON.parse(json))
+  return JSON.parse(json);
+}
+
+export function loadHtmlPages(nextDir: string) {
+  return Object.entries(loadPagesManifest(nextDir))
     .filter(([_, value]) => (value as string).endsWith(".html"))
     .map(([key]) => key);
 }

packages/open-next/src/build/copyTracedFiles.ts

@@ -13,7 +13,15 @@ import {
 } from "node:fs";
 import path from "node:path";
 
-import { loadConfig, loadPrerenderManifest } from "config/util.js";
+import {
+  loadAppPathsManifest,
+  loadBuildId,
+  loadConfig,
+  loadFunctionsConfigManifest,
+  loadMiddlewareManifest,
+  loadPagesManifest,
+  loadPrerenderManifest,
+} from "config/util.js";
 import { getCrossPlatformPathRegex } from "utils/regex.js";
 import logger from "../logger.js";
 import { MIDDLEWARE_TRACE_FILE } from "./constant.js";
@@ -50,6 +58,18 @@ interface CopyTracedFilesOptions {
   skipServerFiles?: boolean;
 }
 
+export function getManifests(nextDir: string) {
+  return {
+    buildId: loadBuildId(nextDir),
+    config: loadConfig(nextDir),
+    prerenderManifest: loadPrerenderManifest(nextDir),
+    pagesManifest: loadPagesManifest(nextDir),
+    appPathsManifest: loadAppPathsManifest(nextDir),
+    middlewareManifest: loadMiddlewareManifest(nextDir),
+    functionsConfigManifest: loadFunctionsConfigManifest(nextDir),
+  };
+}
+
 // eslint-disable-next-line sonarjs/cognitive-complexity
 export async function copyTracedFiles({
   buildOutputPath,
@@ -323,4 +343,9 @@ File ${fullFilePath} does not exist
   }
 
   logger.debug("copyTracedFiles:", Date.now() - tsStart, "ms");
+
+  return {
+    tracedFiles: Array.from(filesToCopy.values()),
+    manifests: getManifests(standaloneNextDir),
+  };
 }

packages/open-next/src/build/createServerBundle.ts

@@ -3,8 +3,10 @@ import path from "node:path";
 
 import type { FunctionOptions, SplittedFunctionOptions } from "types/open-next";
 
+import type { Plugin } from "esbuild";
 import logger from "../logger.js";
 import { minifyAll } from "../minimize-js.js";
+import { ContentUpdater } from "../plugins/content-updater.js";
 import { openNextReplacementPlugin } from "../plugins/replacement.js";
 import { openNextResolvePlugin } from "../plugins/resolve.js";
 import { getCrossPlatformPathRegex } from "../utils/regex.js";
@@ -14,8 +16,20 @@ import { copyTracedFiles } from "./copyTracedFiles.js";
 import { generateEdgeBundle } from "./edge/createEdgeBundle.js";
 import * as buildHelper from "./helper.js";
 import { installDependencies } from "./installDeps.js";
+import { type CodePatcher, applyCodePatches } from "./patch/codePatcher.js";
+
+interface CodeCustomization {
+  // These patches are meant to apply on user and next generated code
+  additionalCodePatches?: CodePatcher[];
+  // These plugins are meant to apply during the esbuild bundling process.
+  // This will only apply to OpenNext code.
+  additionalPlugins?: (contentUpdater: ContentUpdater) => Plugin[];
+}
 
-export async function createServerBundle(options: buildHelper.BuildOptions) {
+export async function createServerBundle(
+  options: buildHelper.BuildOptions,
+  codeCustomization?: CodeCustomization,
+) {
   const { config } = options;
   const foundRoutes = new Set<string>();
   // Get all functions to build
@@ -36,7 +50,7 @@ export async function createServerBundle(options: buildHelper.BuildOptions) {
     if (fnOptions.runtime === "edge") {
       await generateEdgeBundle(name, options, fnOptions);
     } else {
-      await generateBundle(name, options, fnOptions);
+      await generateBundle(name, options, fnOptions, codeCustomization);
     }
   });
 
@@ -101,6 +115,7 @@ async function generateBundle(
   name: string,
   options: buildHelper.BuildOptions,
   fnOptions: SplittedFunctionOptions,
+  codeCustomization?: CodeCustomization,
 ) {
   const { appPath, appBuildOutputPath, config, outputDir, monorepoRoot } =
     options;
@@ -153,14 +168,20 @@ async function generateBundle(
   buildHelper.copyEnvFile(appBuildOutputPath, packagePath, outputPath);
 
   // Copy all necessary traced files
-  await copyTracedFiles({
+  const { tracedFiles, manifests } = await copyTracedFiles({
     buildOutputPath: appBuildOutputPath,
     packagePath,
     outputDir: outputPath,
     routes: fnOptions.routes ?? ["app/page.tsx"],
     bundledNextServer: isBundled,
   });
 
+  const additionalCodePatches = codeCustomization?.additionalCodePatches ?? [];
+
+  await applyCodePatches(options, tracedFiles, manifests, [
+    ...additionalCodePatches,
+  ]);
+
   // Build Lambda code
   // note: bundle in OpenNext package b/c the adapter relies on the
   //       "serverless-http" package which is not a dependency in user's
@@ -179,6 +200,12 @@ async function generateBundle(
 
   const disableRouting = isBefore13413 || config.middleware?.external;
 
+  const updater = new ContentUpdater(options);
+
+  const additionalPlugins = codeCustomization?.additionalPlugins
+    ? codeCustomization.additionalPlugins(updater)
+    : [];
+
   const plugins = [
     openNextReplacementPlugin({
       name: `requestHandlerOverride ${name}`,
@@ -204,6 +231,9 @@ async function generateBundle(
       fnName: name,
       overrides,
     }),
+    ...additionalPlugins,
+    // The content updater plugin must be the last plugin
+    updater.plugin,
   ];
 
   const outfileExt = fnOptions.runtime === "deno" ? "ts" : "mjs";

packages/open-next/src/build/edge/createEdgeBundle.ts

@@ -2,7 +2,7 @@ import { mkdirSync } from "node:fs";
 
 import fs from "node:fs";
 import path from "node:path";
-import { build } from "esbuild";
+import { type Plugin, build } from "esbuild";
 import type { MiddlewareInfo } from "types/next-types";
 import type {
   IncludedConverter,
@@ -16,6 +16,7 @@ import type {
 import { loadMiddlewareManifest } from "config/util.js";
 import type { OriginResolver } from "types/overrides.js";
 import logger from "../../logger.js";
+import { ContentUpdater } from "../../plugins/content-updater.js";
 import { openNextEdgePlugins } from "../../plugins/edge.js";
 import { openNextExternalMiddlewarePlugin } from "../../plugins/externalMiddleware.js";
 import { openNextReplacementPlugin } from "../../plugins/replacement.js";
@@ -39,6 +40,7 @@ interface BuildEdgeBundleOptions {
   additionalExternals?: string[];
   onlyBuildOnce?: boolean;
   name: string;
+  additionalPlugins?: (contentUpdater: ContentUpdater) => Plugin[];
 }
 
 export async function buildEdgeBundle({
@@ -53,13 +55,18 @@ export async function buildEdgeBundle({
   additionalExternals,
   onlyBuildOnce,
   name,
+  additionalPlugins: additionalPluginsFn,
 }: BuildEdgeBundleOptions) {
   const isInCloudfare = await isEdgeRuntime(overrides);
   function override<T extends keyof Override>(target: T) {
     return typeof overrides?.[target] === "string"
       ? overrides[target]
       : undefined;
   }
+  const contentUpdater = new ContentUpdater(options);
+  const additionalPlugins = additionalPluginsFn
+    ? additionalPluginsFn(contentUpdater)
+    : [];
   await esbuildAsync(
     {
       entryPoints: [entrypoint],
@@ -98,6 +105,9 @@ export async function buildEdgeBundle({
           nextDir: path.join(options.appBuildOutputPath, ".next"),
           isInCloudfare,
         }),
+        ...additionalPlugins,
+        // The content updater plugin must be the last plugin
+        contentUpdater.plugin,
       ],
       treeShaking: true,
       alias: {
@@ -173,6 +183,7 @@ export async function generateEdgeBundle(
   name: string,
   options: BuildOptions,
   fnOptions: SplittedFunctionOptions,
+  additionalPlugins: (contentUpdater: ContentUpdater) => Plugin[] = () => [],
 ) {
   logger.info(`Generating edge bundle for: ${name}`);
 
@@ -226,5 +237,6 @@ export async function generateEdgeBundle(
     overrides: fnOptions.override,
     additionalExternals: options.config.edgeExternals,
     name,
+    additionalPlugins,
   });
 }

packages/open-next/src/build/patch/astCodePatcher.ts

@@ -0,0 +1,114 @@
+// Mostly copied from the cloudflare adapter
+import { readFileSync } from "node:fs";
+
+import {
+  type Edit,
+  Lang,
+  type NapiConfig,
+  type SgNode,
+  parse,
+} from "@ast-grep/napi";
+import yaml from "yaml";
+import type { PatchCodeFn } from "./codePatcher";
+
+/**
+ * fix has the same meaning as in yaml rules
+ * see https://ast-grep.github.io/guide/rewrite-code.html#using-fix-in-yaml-rule
+ */
+export type RuleConfig = NapiConfig & { fix?: string };
+
+/**
+ * Returns the `Edit`s and `Match`es for an ast-grep rule in yaml format
+ *
+ * The rule must have a `fix` to rewrite the matched node.
+ *
+ * Tip: use https://ast-grep.github.io/playground.html to create rules.
+ *
+ * @param rule The rule. Either a yaml string or an instance of `RuleConfig`
+ * @param root The root node
+ * @param once only apply once
+ * @returns A list of edits and a list of matches.
+ */
+export function applyRule(
+  rule: string | RuleConfig,
+  root: SgNode,
+  { once = false } = {},
+) {
+  const ruleConfig: RuleConfig =
+    typeof rule === "string" ? yaml.parse(rule) : rule;
+  if (ruleConfig.transform) {
+    throw new Error("transform is not supported");
+  }
+  if (!ruleConfig.fix) {
+    throw new Error("no fix to apply");
+  }
+
+  const fix = ruleConfig.fix;
+
+  const matches = once
+    ? [root.find(ruleConfig)].filter((m) => m !== null)
+    : root.findAll(ruleConfig);
+
+  const edits: Edit[] = [];
+
+  matches.forEach((match) => {
+    edits.push(
+      match.replace(
+        // Replace known placeholders by their value
+        fix
+          .replace(/\$\$\$([A-Z0-9_]+)/g, (_m, name) =>
+            match
+              .getMultipleMatches(name)
+              .map((n) => n.text())
+              .join(""),
+          )
+          .replace(
+            /\$([A-Z0-9_]+)/g,
+            (m, name) => match.getMatch(name)?.text() ?? m,
+          ),
+      ),
+    );
+  });
+
+  return { edits, matches };
+}
+
+/**
+ * Parse a file and obtain its root.
+ *
+ * @param path The file path
+ * @param lang The language to parse. Defaults to TypeScript.
+ * @returns The root for the file.
+ */
+export function parseFile(path: string, lang = Lang.TypeScript) {
+  return parse(lang, readFileSync(path, { encoding: "utf-8" })).root();
+}
+
+/**
+ * Patches the code from by applying the rule.
+ *
+ * This function is mainly for on off edits and tests,
+ * use `getRuleEdits` to apply multiple rules.
+ *
+ * @param code The source code
+ * @param rule The astgrep rule (yaml or NapiConfig)
+ * @param lang The language used by the source code
+ * @param lang Whether to apply the rule only once
+ * @returns The patched code
+ */
+export function patchCode(
+  code: string,
+  rule: string | RuleConfig,
+  { lang = Lang.TypeScript, once = false } = {},
+): string {
+  const node = parse(lang, code).root();
+  const { edits } = applyRule(rule, node, { once });
+  return node.commitEdits(edits);
+}
+
+export function createPatchCode(
+  rule: string | RuleConfig,
+  lang = Lang.TypeScript,
+): PatchCodeFn {
+  return async ({ code }) => patchCode(code, rule, { lang });
+}