Apply single_version_override patches to module files

BUILD

@@ -92,6 +92,13 @@ genrule(
         "MODULE.bazel",
         "//third_party/googleapis:MODULE.bazel",
         "//third_party/remoteapis:MODULE.bazel",
+        "//third_party:BUILD",
+        "//third_party:rules_jvm_external_6.0.patch",
+        "//third_party:protobuf_21.7.patch",
+        "//third_party/upb:BUILD",
+        "//third_party/upb:00_remove_toolchain_transition.patch",
+        "//third_party/grpc:BUILD",
+        "//third_party/grpc:00_disable_layering_check.patch",
     ],
     outs = ["MODULE.bazel.lock.dist"],
     cmd = " && ".join([

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD

@@ -240,6 +240,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:blaze_directories",
         "//src/main/java/com/google/devtools/build/lib/analysis:blaze_version_info",
         "//src/main/java/com/google/devtools/build/lib/bazel:bazel_version",
+        "//src/main/java/com/google/devtools/build/lib/bazel/repository",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository:repository_options",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository/cache",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository/downloader",
@@ -258,6 +259,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/skyframe:precomputed_value",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
+        "//src/main/java/com/google/devtools/build/lib/vfs/inmemoryfs",
         "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
         "//src/main/java/net/starlark/java/annot",
         "//src/main/java/net/starlark/java/eval",
@@ -268,6 +270,7 @@ java_library(
         "//third_party:gson",
         "//third_party:guava",
         "//third_party:jsr305",
+        "@maven//:io_github_java_diff_utils_java_diff_utils",
     ],
 )
 

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunction.java

@@ -20,14 +20,17 @@
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
+import com.github.difflib.patch.PatchFailedException;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
 import com.google.devtools.build.lib.actions.FileValue;
 import com.google.devtools.build.lib.bazel.bzlmod.CompiledModuleFile.IncludeStatement;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.NonRootModuleFileValue;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.RootModuleFileValue;
+import com.google.devtools.build.lib.bazel.repository.PatchUtil;
 import com.google.devtools.build.lib.bazel.repository.downloader.Checksum.MissingChecksumException;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelConstants;
@@ -50,11 +53,13 @@
 import com.google.devtools.build.lib.skyframe.PackageLookupValue;
 import com.google.devtools.build.lib.skyframe.PrecomputedValue;
 import com.google.devtools.build.lib.skyframe.PrecomputedValue.Precomputed;
+import com.google.devtools.build.lib.vfs.DigestHashFunction;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.Root;
 import com.google.devtools.build.lib.vfs.RootedPath;
+import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem;
 import com.google.devtools.build.skyframe.SkyFunction;
 import com.google.devtools.build.skyframe.SkyFunctionException;
 import com.google.devtools.build.skyframe.SkyFunctionException.Transience;
@@ -616,11 +621,15 @@ private GetModuleFileResult getModuleFile(
     StoredEventHandler downloadEventHandler = new StoredEventHandler();
     for (Registry registry : registryObjects) {
       try {
-        Optional<ModuleFile> moduleFile = registry.getModuleFile(key, downloadEventHandler);
-        if (moduleFile.isEmpty()) {
+        Optional<ModuleFile> maybeModuleFile = registry.getModuleFile(key, downloadEventHandler);
+        if (maybeModuleFile.isEmpty()) {
           continue;
         }
-        return new GetModuleFileResult(moduleFile.get(), registry, downloadEventHandler);
+        ModuleFile moduleFile = maybePatchModuleFile(maybeModuleFile.get(), override, env);
+        if (moduleFile == null) {
+          return null;
+        }
+        return new GetModuleFileResult(moduleFile, registry, downloadEventHandler);
       } catch (MissingChecksumException e) {
         throw new ModuleFileFunctionException(
             ExternalDepsException.withCause(Code.BAD_LOCKFILE, e));
@@ -633,6 +642,113 @@ private GetModuleFileResult getModuleFile(
     throw errorf(Code.MODULE_NOT_FOUND, "module not found in registries: %s", key);
   }
 
+  /**
+   * Applies any patches specified in registry overrides.
+   *
+   * <p>This allows users to modify MODULE.bazel files and thus influence resolution and visibility
+   * for modules via patches without having to replace the entire module via a non-registry
+   * override.
+   *
+   * <p>Note: Only patch files from the main repo are applied, all other patches are ignored. This
+   * is necessary as we can't load other repos during resolution (unless they are subject to a
+   * non-registry override). Patch commands are also not supported as they cannot be selectively
+   * applied to the module file only.
+   */
+  private ModuleFile maybePatchModuleFile(
+      ModuleFile moduleFile, ModuleOverride override, Environment env)
+      throws InterruptedException, ModuleFileFunctionException {
+    if (!(override instanceof SingleVersionOverride singleVersionOverride)) {
+      return moduleFile;
+    }
+    var patchesInMainRepo =
+        singleVersionOverride.getPatches().stream()
+            .filter(label -> label.getRepository().isMain())
+            .collect(toImmutableList());
+    if (patchesInMainRepo.isEmpty()) {
+      return moduleFile;
+    }
+
+    // Get the patch paths.
+    var patchPackageLookupKeys =
+        patchesInMainRepo.stream()
+            .map(Label::getPackageIdentifier)
+            .map(pkg -> (SkyKey) PackageLookupValue.key(pkg))
+            .collect(toImmutableList());
+    var patchPackageLookupResult = env.getValuesAndExceptions(patchPackageLookupKeys);
+    if (env.valuesMissing()) {
+      return null;
+    }
+    List<RootedPath> patchPaths = new ArrayList<>();
+    for (int i = 0; i < patchesInMainRepo.size(); i++) {
+      var pkgLookupValue =
+          (PackageLookupValue) patchPackageLookupResult.get(patchPackageLookupKeys.get(i));
+      if (pkgLookupValue == null) {
+        return null;
+      }
+      if (!pkgLookupValue.packageExists()) {
+        Label label = patchesInMainRepo.get(i);
+        throw errorf(
+            Code.BAD_MODULE,
+            "unable to load package for single_version_override patch %s: %s",
+            label,
+            PackageLookupFunction.explainNoBuildFileValue(label.getPackageIdentifier(), env));
+      }
+      patchPaths.add(
+          RootedPath.toRootedPath(
+              pkgLookupValue.getRoot(), patchesInMainRepo.get(i).toPathFragment()));
+    }
+
+    // Register dependencies on the patch files and ensure that they exist.
+    var fileKeys = Iterables.transform(patchPaths, FileValue::key);
+    var fileValueResult = env.getValuesAndExceptions(fileKeys);
+    if (env.valuesMissing()) {
+      return null;
+    }
+    for (var key : fileKeys) {
+      var fileValue = (FileValue) fileValueResult.get(key);
+      if (fileValue == null) {
+        return null;
+      }
+      if (!fileValue.isFile()) {
+        throw errorf(
+            Code.BAD_MODULE,
+            "error reading single_version_override patch %s: not a regular file",
+            key.argument());
+      }
+    }
+
+    // Apply the patches to the module file only.
+    InMemoryFileSystem patchFs = new InMemoryFileSystem(DigestHashFunction.SHA256);
+    try {
+      Path moduleRoot = patchFs.getPath("/module");
+      moduleRoot.createDirectoryAndParents();
+      Path moduleFilePath = moduleRoot.getChild("MODULE.bazel");
+      FileSystemUtils.writeContent(moduleFilePath, moduleFile.getContent());
+      for (var patchPath : patchPaths) {
+        try {
+          PatchUtil.applyToSingleFile(
+              patchPath.asPath(),
+              singleVersionOverride.getPatchStrip(),
+              moduleRoot,
+              moduleFilePath);
+        } catch (PatchFailedException e) {
+          throw errorf(
+              Code.BAD_MODULE,
+              "error applying single_version_override patch %s to module file: %s",
+              patchPath.asPath(),
+              e.getMessage());
+        }
+      }
+      return ModuleFile.create(
+          FileSystemUtils.readContent(moduleFilePath), moduleFile.getLocation());
+    } catch (IOException e) {
+      throw errorf(
+          Code.BAD_MODULE,
+          "error applying single_version_override patches to module file: %s",
+          e.getMessage());
+    }
+  }
+
   private static byte[] readModuleFile(Path path) throws ModuleFileFunctionException {
     try {
       return FileSystemUtils.readWithKnownFileSize(path, path.getFileSize());

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileGlobals.java

@@ -784,15 +784,20 @@ public void include(String label, StarlarkThread thread)
             doc =
                 "A list of labels pointing to patch files to apply for this module. The patch files"
                     + " must exist in the source tree of the top level project. They are applied in"
-                    + " the list order.",
+                    + " the list order."
+                    + ""
+                    + "<p>If a patch makes changes to the MODULE.bazel file, these changes will"
+                    + " only be effective if the patch file is provided by the root module.",
             allowedTypes = {@ParamType(type = Iterable.class, generic1 = String.class)},
             named = true,
             positional = false,
             defaultValue = "[]"),
         @Param(
             name = "patch_cmds",
             doc =
-                "Sequence of Bash commands to be applied on Linux/Macos after patches are applied.",
+                "Sequence of Bash commands to be applied on Linux/Macos after patches are applied."
+                    + ""
+                    + "<p>Changes to the MODULE.bazel file will not be effective.",
             allowedTypes = {@ParamType(type = Iterable.class, generic1 = String.class)},
             named = true,
             positional = false,

src/main/java/com/google/devtools/build/lib/bazel/repository/PatchUtil.java

@@ -417,10 +417,32 @@ private static void checkFilesStatusForPatching(
    *
    * @param patchFile the patch file to apply
    * @param strip the number of leading components to strip from file path in the patch file
-   * @param outputDirectory the repository directory to apply the patch file
+   * @param outputDirectory the directory to apply the patch file to
    */
   public static void apply(Path patchFile, int strip, Path outputDirectory)
       throws IOException, PatchFailedException {
+    applyInternal(patchFile, strip, outputDirectory, /* singleFile= */ null);
+  }
+
+  /**
+   * Apply a patch file under a directory, skipping all parts of the patch file that do not apply to
+   * the given single file.
+   *
+   * @param patchFile the patch file to apply
+   * @param strip the number of leading components to strip from file path in the patch file
+   * @param outputDirectory the directory to apply the patch file to
+   * @param singleFile only apply the parts of the patch file that apply to this file. Renaming the
+   *     file is not supported in this case.
+   */
+  public static void applyToSingleFile(
+      Path patchFile, int strip, Path outputDirectory, Path singleFile)
+      throws IOException, PatchFailedException {
+    applyInternal(patchFile, strip, outputDirectory, singleFile);
+  }
+
+  private static void applyInternal(
+      Path patchFile, int strip, Path outputDirectory, @Nullable Path singleFile)
+      throws IOException, PatchFailedException {
     if (!patchFile.exists()) {
       throw new PatchFailedException("Cannot find patch file: " + patchFile.getPathString());
     }
@@ -561,15 +583,25 @@ public static void apply(Path patchFile, int strip, Path outputDirectory)
                 patchContent, header, oldLineCount, newLineCount, patchStartLocation);
 
             if (isRenaming) {
-              checkFilesStatusForRenaming(
-                  oldFile, newFile, oldFileStr, newFileStr, patchStartLocation);
+              if (singleFile != null) {
+                if (singleFile.equals(newFile) || singleFile.equals(oldFile)) {
+                  throw new PatchFailedException(
+                      "Renaming %s while applying patches to it as a single file is not supported."
+                          .formatted(singleFile));
+                }
+              } else {
+                checkFilesStatusForRenaming(
+                    oldFile, newFile, oldFileStr, newFileStr, patchStartLocation);
+              }
             }
 
-            Patch<String> patch = UnifiedDiffUtils.parseUnifiedDiff(patchContent);
-            checkFilesStatusForPatching(
-                patch, oldFile, newFile, oldFileStr, newFileStr, patchStartLocation);
+            if (singleFile == null || (singleFile.equals(newFile) && singleFile.equals(oldFile))) {
+              Patch<String> patch = UnifiedDiffUtils.parseUnifiedDiff(patchContent);
+              checkFilesStatusForPatching(
+                  patch, oldFile, newFile, oldFileStr, newFileStr, patchStartLocation);
 
-            applyPatchToFile(patch, oldFile, newFile, isRenaming, filePermission);
+              applyPatchToFile(patch, oldFile, newFile, isRenaming, filePermission);
+            }
           }
 
           patchContent.clear();