feat: add support for arbitrary headers rctx.download[_and_extract]

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

@@ -47,6 +47,7 @@ public void setDelegate(@Nullable Downloader delegate) {
   @Override
   public void download(
       List<URL> urls,
+      Map<String, List<String>> headers,
       Credentials credentials,
       Optional<Checksum> checksum,
       String canonicalId,
@@ -60,6 +61,6 @@ public void download(
       downloader = delegate;
     }
     downloader.download(
-        urls, credentials, checksum, canonicalId, destination, eventHandler, clientEnv, type);
+        urls, headers, credentials, checksum, canonicalId, destination, eventHandler, clientEnv, type);
   }
 }

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

@@ -116,6 +116,7 @@ public void setCredentialFactory(CredentialFactory credentialFactory) {
 
   public Future<Path> startDownload(
       List<URL> originalUrls,
+      Map<String, List<String>> headers,
       Map<URI, Map<String, List<String>>> authHeaders,
       Optional<Checksum> checksum,
       String canonicalId,
@@ -129,6 +130,7 @@ public Future<Path> startDownload(
           try (SilentCloseable c = Profiler.instance().profile("fetching: " + context)) {
             return downloadInExecutor(
                 originalUrls,
+                headers,
                 authHeaders,
                 checksum,
                 canonicalId,
@@ -154,6 +156,7 @@ public Path finalizeDownload(Future<Path> download) throws IOException, Interrup
 
   public Path download(
       List<URL> originalUrls,
+      Map<String, List<String>> headers,
       Map<URI, Map<String, List<String>>> authHeaders,
       Optional<Checksum> checksum,
       String canonicalId,
@@ -166,6 +169,7 @@ public Path download(
     Future<Path> future =
         startDownload(
             originalUrls,
+            headers,
             authHeaders,
             checksum,
             canonicalId,
@@ -197,6 +201,7 @@ public Path download(
    */
   private Path downloadInExecutor(
       List<URL> originalUrls,
+      Map<String, List<String>> headers,
       Map<URI, Map<String, List<String>>> authHeaders,
       Optional<Checksum> checksum,
       String canonicalId,
@@ -339,6 +344,7 @@ private Path downloadInExecutor(
       try {
         downloader.download(
             rewrittenUrls,
+            headers,
             credentialFactory.create(rewrittenAuthHeaders),
             checksum,
             canonicalId,

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

@@ -42,6 +42,7 @@ public interface Downloader {
    */
   void download(
       List<URL> urls,
+      Map<String, List<String>> headers,
       Credentials credentials,
       Optional<Checksum> checksum,
       String canonicalId,

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

@@ -75,7 +75,7 @@ final class HttpConnectorMultiplexer {
   }
 
   public HttpStream connect(URL url, Optional<Checksum> checksum) throws IOException {
-    return connect(url, checksum, StaticCredentials.EMPTY, Optional.empty());
+    return connect(url, checksum, ImmutableMap.of(), StaticCredentials.EMPTY, Optional.empty());
   }
 
   /**
@@ -96,14 +96,19 @@ public HttpStream connect(URL url, Optional<Checksum> checksum) throws IOExcepti
    * @throws IllegalArgumentException if {@code urls} is empty or has an unsupported protocol
    */
   public HttpStream connect(
-      URL url, Optional<Checksum> checksum, Credentials credentials, Optional<String> type)
+      URL url, Optional<Checksum> checksum, Map<String, List<String>> headers, Credentials credentials, Optional<String> type)
       throws IOException {
     Preconditions.checkArgument(HttpUtils.isUrlSupportedByDownloader(url));
     if (Thread.interrupted()) {
       throw new InterruptedIOException();
     }
+    ImmutableMap.Builder<String, List<String>> baseHeaders = new ImmutableMap.Builder();
+    baseHeaders.putAll(headers);
+    // REQUEST_HEADERS should not be overridable by user provided headers
+    baseHeaders.putAll(REQUEST_HEADERS);
+
     Function<URL, ImmutableMap<String, List<String>>> headerFunction =
-        getHeaderFunction(REQUEST_HEADERS, credentials, eventHandler);
+        getHeaderFunction(baseHeaders.buildKeepingLast(), credentials, eventHandler);
     URLConnection connection = connector.connect(url, headerFunction);
     return httpStreamFactory.create(
         connection,

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

@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.bazel.repository.downloader;
 
 import com.google.auth.Credentials;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.io.ByteStreams;
@@ -75,6 +76,7 @@ public void setMaxRetryTimeout(Duration maxRetryTimeout) {
   @Override
   public void download(
       List<URL> urls,
+      Map<String, List<String>> headers,
       Credentials credentials,
       Optional<Checksum> checksum,
       String canonicalId,
@@ -94,7 +96,7 @@ public void download(
     for (URL url : urls) {
       SEMAPHORE.acquire();
 
-      try (HttpStream payload = multiplexer.connect(url, checksum, credentials, type);
+      try (HttpStream payload = multiplexer.connect(url, checksum, headers, credentials, type);
           OutputStream out = destination.getOutputStream()) {
         try {
           ByteStreams.copy(payload, out);
@@ -153,7 +155,7 @@ public byte[] downloadAndReadOneUrl(
     ByteArrayOutputStream out = new ByteArrayOutputStream();
     SEMAPHORE.acquire();
     try (HttpStream payload =
-        multiplexer.connect(url, Optional.empty(), credentials, Optional.empty())) {
+        multiplexer.connect(url, Optional.empty(), ImmutableMap.of(), credentials, Optional.empty())) {
       ByteStreams.copy(payload, out);
     } catch (SocketTimeoutException e) {
       // SocketTimeoutExceptions are InterruptedIOExceptions; however they do not signify

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

@@ -80,6 +80,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.stream.Collectors;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
@@ -281,6 +282,27 @@ private static ImmutableMap<URI, Map<String, List<String>>> getAuthHeaders(
     return res;
   }
 
+  private static ImmutableMap<String, List<String>> getHeaderContents(Dict<?, ?> x, String what)
+      throws EvalException {
+    Dict<String, Object> headersUnchecked = (Dict<String, Object>) Dict.cast(x, String.class, Object.class, what);
+    ImmutableMap.Builder<String, List<String>> headers = new ImmutableMap.Builder<>();
+
+    for (Map.Entry<String, Object> entry : headersUnchecked.entrySet()) {
+      List<String> headerValue;
+      Object valueUnchecked = entry.getValue();
+      if (valueUnchecked instanceof Sequence) {
+        headerValue = Sequence.cast(valueUnchecked, String.class, "header values").getImmutableList();
+      } else if (valueUnchecked instanceof String) {
+        headerValue = List.of(valueUnchecked.toString());
+      } else {
+        throw new EvalException(
+          String.format("%s argument must be a dict whose keys are string and whose values are either string or sequence of string", what));
+      }
+      headers.put(entry.getKey(), headerValue);
+    }
+    return headers.buildOrThrow();
+  }
+
   private static ImmutableList<String> checkAllUrls(Iterable<?> urlList) throws EvalException {
     ImmutableList.Builder<String> result = ImmutableList.builder();
 
@@ -577,6 +599,11 @@ private StructImpl completeDownload(PendingDownload pendingDownload)
             defaultValue = "{}",
             named = true,
             doc = "An optional dict specifying authentication information for some of the URLs."),
+        @Param(
+              name = "headers",
+              defaultValue = "{}",
+              named = true,
+              doc = "An optional dict specifying http headers for all URLs."),
         @Param(
             name = "integrity",
             defaultValue = "''",
@@ -606,7 +633,8 @@ public Object download(
       Boolean executable,
       Boolean allowFail,
       String canonicalId,
-      Dict<?, ?> authUnchecked, // <String, Dict> expected
+      Dict<?, ?> authUnchecked, // <String, Dict> expected 
+      Dict<?, ?> headersUnchecked, // <String, List<String> | String> expected
       String integrity,
       Boolean block,
       StarlarkThread thread)
@@ -615,6 +643,8 @@ public Object download(
     ImmutableMap<URI, Map<String, List<String>>> authHeaders =
         getAuthHeaders(getAuthContents(authUnchecked, "auth"));
 
+    ImmutableMap<String, List<String>> headers = getHeaderContents(headersUnchecked, "headers");
+
     ImmutableList<URL> urls =
         getUrls(
             url,
@@ -660,6 +690,7 @@ public Object download(
       Future<Path> downloadFuture =
           downloadManager.startDownload(
               urls,
+              headers,
               authHeaders,
               checksum,
               canonicalId,
@@ -768,6 +799,11 @@ public Object download(
             defaultValue = "{}",
             named = true,
             doc = "An optional dict specifying authentication information for some of the URLs."),
+        @Param(
+              name = "headers",
+              defaultValue = "{}",
+              named = true,
+              doc = "An optional dict specifying http headers for all URLs."),
         @Param(
             name = "integrity",
             defaultValue = "''",
@@ -799,13 +835,16 @@ public StructImpl downloadAndExtract(
       String stripPrefix,
       Boolean allowFail,
       String canonicalId,
-      Dict<?, ?> auth, // <String, Dict> expected
+      Dict<?, ?> authUnchecked, // <String, Dict> expected
+      Dict<?, ?> headersUnchecked, // <String, List<String> | String> expected
       String integrity,
       Dict<?, ?> renameFiles, // <String, String> expected
       StarlarkThread thread)
       throws RepositoryFunctionException, InterruptedException, EvalException {
     ImmutableMap<URI, Map<String, List<String>>> authHeaders =
-        getAuthHeaders(getAuthContents(auth, "auth"));
+        getAuthHeaders(getAuthContents(authUnchecked, "auth"));
+
+    ImmutableMap<String, List<String>> headers = getHeaderContents(headersUnchecked, "headers");
 
     ImmutableList<URL> urls =
         getUrls(
@@ -852,6 +891,7 @@ public StructImpl downloadAndExtract(
       Future<Path> pendingDownload =
           downloadManager.startDownload(
               urls,
+              headers,
               authHeaders,
               checksum,
               canonicalId,

src/main/java/com/google/devtools/build/lib/remote/downloader/GrpcRemoteDownloader.java

@@ -110,6 +110,7 @@ public void close() {
   @Override
   public void download(
       List<URL> urls,
+      Map<String, List<String>> headers,
       Credentials credentials,
       Optional<Checksum> checksum,
       String canonicalId,
@@ -154,7 +155,7 @@ public void download(
       eventHandler.handle(
           Event.warn("Remote Cache: " + Utils.grpcAwareErrorMessage(e, verboseFailures)));
       fallbackDownloader.download(
-          urls, credentials, checksum, canonicalId, destination, eventHandler, clientEnv, type);
+          urls, headers, credentials, checksum, canonicalId, destination, eventHandler, clientEnv, type);
     }
   }
 

src/test/java/com/google/devtools/build/lib/bazel/repository/downloader/HttpDownloaderTest.java

@@ -117,6 +117,7 @@ public void downloadFrom1UrlOk() throws IOException, InterruptedException {
               Collections.singletonList(
                   new URL(String.format("http://localhost:%d/foo", server.getLocalPort()))),
               Collections.emptyMap(),
+              Collections.emptyMap(),
               Optional.empty(),
               "testCanonicalId",
               Optional.empty(),
@@ -181,6 +182,7 @@ public void downloadFrom2UrlsFirstOk() throws IOException, InterruptedException
           downloadManager.download(
               urls,
               Collections.emptyMap(),
+              Collections.emptyMap(),
               Optional.empty(),
               "testCanonicalId",
               Optional.empty(),
@@ -248,6 +250,7 @@ public void downloadFrom2UrlsFirstSocketTimeoutOnBodyReadSecondOk()
           downloadManager.download(
               urls,
               Collections.emptyMap(),
+              Collections.emptyMap(),
               Optional.empty(),
               "testCanonicalId",
               Optional.empty(),
@@ -317,6 +320,7 @@ public void downloadFrom2UrlsBothSocketTimeoutDuringBodyRead()
         downloadManager.download(
             urls,
             Collections.emptyMap(),
+            Collections.emptyMap(),
             Optional.empty(),
             "testCanonicalId",
             Optional.empty(),
@@ -371,6 +375,7 @@ public void downloadOneUrl_ok() throws IOException, InterruptedException {
       httpDownloader.download(
           Collections.singletonList(
               new URL(String.format("http://localhost:%d/foo", server.getLocalPort()))),
+          Collections.emptyMap(),
           StaticCredentials.EMPTY,
           Optional.empty(),
           "testCanonicalId",
@@ -410,6 +415,7 @@ public void downloadOneUrl_notFound() throws IOException, InterruptedException {
               httpDownloader.download(
                   Collections.singletonList(
                       new URL(String.format("http://localhost:%d/foo", server.getLocalPort()))),
+                  Collections.emptyMap(),
                   StaticCredentials.EMPTY,
                   Optional.empty(),
                   "testCanonicalId",
@@ -470,6 +476,7 @@ public void downloadTwoUrls_firstNotFoundAndSecondOk() throws IOException, Inter
       Path destination = fs.getPath(workingDir.newFile().getAbsolutePath());
       httpDownloader.download(
           urls,
+          Collections.emptyMap(),
           StaticCredentials.EMPTY,
           Optional.empty(),
           "testCanonicalId",
@@ -564,13 +571,14 @@ public void download_contentLengthMismatch_propagateErrorIfNotRetry() throws Exc
                   throw new ContentLengthMismatchException(0, data.length);
                 })
         .when(downloader)
-        .download(any(), any(), any(), any(), any(), any(), any(), any());
+        .download(any(), any(), any(), any(), any(), any(), any(), any(), any());
 
     assertThrows(
         ContentLengthMismatchException.class,
         () ->
             downloadManager.download(
                 ImmutableList.of(new URL("http://localhost")),
+                Collections.emptyMap(),
                 ImmutableMap.of(),
                 Optional.empty(),
                 "testCanonicalId",
@@ -597,20 +605,21 @@ public void download_contentLengthMismatch_retries() throws Exception {
                   if (times.getAndIncrement() < 3) {
                     throw new ContentLengthMismatchException(0, data.length);
                   }
-                  Path output = invocationOnMock.getArgument(4, Path.class);
+                  Path output = invocationOnMock.getArgument(5, Path.class);
                   try (OutputStream outputStream = output.getOutputStream()) {
                     ByteStreams.copy(new ByteArrayInputStream(data), outputStream);
                   }
 
                   return null;
                 })
         .when(downloader)
-        .download(any(), any(), any(), any(), any(), any(), any(), any());
+        .download(any(), any(), any(), any(), any(), any(), any(), any(), any());
 
     Path result =
         downloadManager.download(
             ImmutableList.of(new URL("http://localhost")),
             ImmutableMap.of(),
+            ImmutableMap.of(),
             Optional.empty(),
             "testCanonicalId",
             Optional.empty(),

src/test/java/com/google/devtools/build/lib/remote/downloader/GrpcRemoteDownloaderTest.java

@@ -178,6 +178,7 @@ private static byte[] downloadBlob(
     final Path destination = scratch.resolve("output file path");
     downloader.download(
         urls,
+        ImmutableMap.of(),
         StaticCredentials.EMPTY,
         checksum,
         canonicalId,
@@ -240,13 +241,13 @@ public void fetchBlob(
             invocation -> {
               List<URL> urls = invocation.getArgument(0);
               if (urls.equals(ImmutableList.of(new URL("http://example.com/content.txt")))) {
-                Path output = invocation.getArgument(4);
+                Path output = invocation.getArgument(5);
                 FileSystemUtils.writeContent(output, content);
               }
               return null;
             })
         .when(fallbackDownloader)
-        .download(any(), any(), any(), any(), any(), any(), any(), any());
+        .download(any(), any(), any(), any(), any(), any(), any(), any(), any());
     final GrpcRemoteDownloader downloader = newDownloader(cacheClient, fallbackDownloader);
 
     final byte[] downloaded =