Add support for overlays in source.json

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

@@ -15,10 +15,15 @@
 
 package com.google.devtools.build.lib.bazel.bzlmod;
 
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
 import net.starlark.java.eval.StarlarkInt;
 
 /**
@@ -77,6 +82,21 @@ public ArchiveRepoSpecBuilder setRemotePatches(ImmutableMap<String, String> remo
     return this;
   }
 
+  @CanIgnoreReturnValue
+  public ArchiveRepoSpecBuilder setOverlay(ImmutableMap<String, RemoteFile> overlay) {
+    Map<String, List<String>> remoteFiles =
+        overlay.entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Entry::getKey,
+                    e -> e.getValue().urls()));
+    Map<String, String> remoteFilesIntegrity =
+        overlay.entrySet().stream().collect(toImmutableMap(Entry::getKey, e -> e.getValue().integrity()));
+    attrBuilder.put("remote_file_urls", remoteFiles);
+    attrBuilder.put("remote_file_integrity", remoteFilesIntegrity);
+    return this;
+  }
+
   @CanIgnoreReturnValue
   public ArchiveRepoSpecBuilder setRemotePatchStrip(int remotePatchStrip) {
     attrBuilder.put("remote_patch_strip", StarlarkInt.of(remotePatchStrip));
@@ -98,4 +118,10 @@ public RepoSpec build() {
         .setAttributes(AttributeValues.create(attrBuilder.buildOrThrow()))
         .build();
   }
+
+  /**
+   * A simple pojo to track remote files that are offered at multiple urls (mirrors) with a single
+   * integrity. We split up the file here to simplify the dependency.
+   */
+  public record RemoteFile(String integrity, List<String> urls) {}
 }

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

@@ -17,6 +17,9 @@
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+
+
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
@@ -38,12 +41,15 @@
 import com.google.gson.JsonParseException;
 import java.io.FileNotFoundException;
 import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.net.MalformedURLException;
 import java.net.URI;
 import java.net.URL;
+import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
 /**
  * Represents a Bazel module registry that serves a list of module metadata from a static HTTP
@@ -222,6 +228,7 @@ private static class ArchiveSourceJson {
     String integrity;
     String stripPrefix;
     Map<String, String> patches;
+    Map<String, String> overlay;
     int patchStrip;
     String archiveType;
   }
@@ -412,11 +419,32 @@ private RepoSpec createArchiveRepoSpec(
       }
     }
 
+    Map<String, String> sourceJsonOverlay = sourceJson.overlay != null
+        ? ImmutableMap.copyOf(sourceJson.overlay)
+        : ImmutableMap.of();
+    ImmutableMap<String, ArchiveRepoSpecBuilder.RemoteFile> overlay = sourceJsonOverlay
+        .entrySet()
+        .stream()
+        .collect(toImmutableMap(
+            entry -> entry.getKey(),
+            entry -> new ArchiveRepoSpecBuilder.RemoteFile(
+                entry.getValue(), // integrity
+                // URLs in the registry itself are not mirrored.
+                ImmutableList.of(constructUrl(
+                    getUrl(),
+                    "modules",
+                    key.getName(),
+                    key.getVersion().toString(),
+                    entry.getKey())
+                ))
+            ));
+
     return new ArchiveRepoSpecBuilder()
         .setUrls(urls.build())
         .setIntegrity(sourceJson.integrity)
         .setStripPrefix(Strings.nullToEmpty(sourceJson.stripPrefix))
         .setRemotePatches(remotePatches.buildOrThrow())
+        .setOverlay(overlay)
         .setRemotePatchStrip(sourceJson.patchStrip)
         .setArchiveType(sourceJson.archiveType)
         .build();

src/test/java/com/google/devtools/build/lib/bazel/bzlmod/IndexRegistryTest.java

@@ -41,6 +41,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.io.Writer;
+import java.net.URI;
 import java.nio.file.Files;
 import java.util.ArrayList;
 import java.util.List;
@@ -182,6 +183,17 @@ public void testGetArchiveRepoSpec() throws Exception {
         "  },",
         "  \"patch_strip\": 3",
         "}");
+    server.serve(
+        "/modules/baz/3.0/source.json",
+        """
+        {
+            "url": "https://example.com/archive.jar?with=query",
+            "integrity": "sha256-bleh",
+            "overlay": {
+                "BUILD.bazel": "sha256-bleh-overlay"
+            }
+        }
+        """);
     server.start();
 
     Registry registry =
@@ -198,6 +210,7 @@ public void testGetArchiveRepoSpec() throws Exception {
                 .setIntegrity("sha256-blah")
                 .setStripPrefix("pref")
                 .setRemotePatches(ImmutableMap.of())
+                .setOverlay(ImmutableMap.of())
                 .setRemotePatchStrip(0)
                 .build());
     assertThat(registry.getRepoSpec(createModuleKey("bar", "2.0"), reporter))
@@ -216,6 +229,30 @@ public void testGetArchiveRepoSpec() throws Exception {
                         server.getUrl() + "/modules/bar/2.0/patches/2.fix-that.patch",
                             "sha256-kek"))
                 .setRemotePatchStrip(3)
+                .setOverlay(ImmutableMap.of())
+                .build());
+    assertThat(registry.getRepoSpec(createModuleKey("baz", "3.0"), reporter))
+        .isEqualTo(
+            new ArchiveRepoSpecBuilder()
+                .setUrls(
+                    ImmutableList.of(
+                        "https://mirror.bazel.build/example.com/archive.jar?with=query",
+                        "file:///home/bazel/mymirror/example.com/archive.jar?with=query",
+                        "https://example.com/archive.jar?with=query"))
+                .setIntegrity("sha256-bleh")
+                .setStripPrefix("")
+                .setOverlay(
+                    ImmutableMap.of(
+                        "BUILD.bazel",
+                        new ArchiveRepoSpecBuilder.RemoteFile(
+                            "sha256-bleh-overlay",
+                            // URLs in the registry itself are not mirrored.
+                            ImmutableList.of(
+                                server.getUrl() + "/modules/baz/3.0/BUILD.bazel"
+                            ))
+                    ))
+                .setRemotePatches(ImmutableMap.of())
+                .setRemotePatchStrip(0)
                 .build());
   }
 
@@ -264,6 +301,7 @@ public void testGetRepoInvalidRegistryJsonSpec() throws Exception {
                 .setIntegrity("sha256-blah")
                 .setStripPrefix("pref")
                 .setRemotePatches(ImmutableMap.of())
+                .setOverlay(ImmutableMap.of())
                 .setRemotePatchStrip(0)
                 .build());
   }
@@ -351,6 +389,7 @@ public void testArchiveWithExplicitType() throws Exception {
                 .setArchiveType("zip")
                 .setRemotePatches(ImmutableMap.of())
                 .setRemotePatchStrip(0)
+                .setOverlay(ImmutableMap.of())
                 .build());
   }
 

src/test/py/bazel/bzlmod/mod_command_test.py

@@ -397,13 +397,15 @@ def testShowModuleAndExtensionReposFromBaseModule(self):
     )
     self.assertRegex(stdout.pop(4), r'^  urls = \[".*"\],$')
     self.assertRegex(stdout.pop(4), r'^  integrity = ".*",$')
-    stdout.pop(11)
-    self.assertRegex(stdout.pop(16), r'^  path = ".*",$')
-    stdout.pop(29)
-    stdout.pop(39)
-    self.assertRegex(stdout.pop(44), r'^  urls = \[".*"\],$')
-    self.assertRegex(stdout.pop(44), r'^  integrity = ".*",$')
-    stdout.pop(51)
+    self.assertRegex(stdout.pop(19), r'^  path = ".*",$')
+    # lines after 'Rule data_repo defined at (most recent call last):'
+    stdout.pop(32)
+    stdout.pop(42)
+    self.assertRegex(stdout.pop(47), r'^  urls = \[".*"\],$')
+    self.assertRegex(stdout.pop(47), r'^  integrity = ".*",$')
+    # lines after '# Rule http_archive defined at (most recent call last):'
+    stdout.pop(13)
+    stdout.pop(55)
     self.assertListEqual(
         stdout,
         [
@@ -414,19 +416,21 @@ def testShowModuleAndExtensionReposFromBaseModule(self):
             # pop(4) -- urls=[...]
             # pop(4) -- integrity=...
             '  strip_prefix = "",',
+            '  remote_file_urls = {},',
+            '  remote_file_integrity = {},',
             '  remote_patches = {},',
             '  remote_patch_strip = 0,',
             ')',
             '# Rule bar~ instantiated at (most recent call last):',
             '#   <builtin> in <toplevel>',
             '# Rule http_archive defined at (most recent call last):',
-            # pop(11)
+            # pop(13)
             '',
             '## ext@1.0:',
             '# <builtin>',
             'local_repository(',
             '  name = "ext~",',
-            # pop(16) -- path=...
+            # pop(19) -- path=...
             ')',
             '# Rule ext~ instantiated at (most recent call last):',
             '#   <builtin> in <toplevel>',
@@ -440,7 +444,7 @@ def testShowModuleAndExtensionReposFromBaseModule(self):
             '# Rule ext~~ext~repo3 instantiated at (most recent call last):',
             '#   <builtin> in <toplevel>',
             '# Rule data_repo defined at (most recent call last):',
-            # pop(29)
+            # pop(32)
             '',
             '## @my_repo4:',
             '# <builtin>',
@@ -451,22 +455,24 @@ def testShowModuleAndExtensionReposFromBaseModule(self):
             '# Rule ext~~ext~repo4 instantiated at (most recent call last):',
             '#   <builtin> in <toplevel>',
             '# Rule data_repo defined at (most recent call last):',
-            # pop(39)
+            # pop(42)
             '',
             '## bar@2.0:',
             '# <builtin>',
             'http_archive(',
             '  name = "bar~",',
-            # pop(44) -- urls=[...]
-            # pop(44) -- integrity=...
+            # pop(47) -- urls=[...]
+            # pop(47) -- integrity=...
             '  strip_prefix = "",',
+            '  remote_file_urls = {},',
+            '  remote_file_integrity = {},',
             '  remote_patches = {},',
             '  remote_patch_strip = 0,',
             ')',
             '# Rule bar~ instantiated at (most recent call last):',
             '#   <builtin> in <toplevel>',
             '# Rule http_archive defined at (most recent call last):',
-            # pop(51)
+            # pop(55)
             '',
         ],
         'wrong output in the show query for module and extension-generated'