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 Path download(
       List<URL> originalUrls,
+      Map<String, List<String>> headers,
       Map<URI, Map<String, List<String>>> authHeaders,
       Optional<Checksum> checksum,
       String canonicalId,
@@ -258,6 +259,7 @@ public Path download(
       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;
@@ -74,6 +75,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,
@@ -93,7 +95,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);
@@ -152,7 +154,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

@@ -77,6 +77,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import net.starlark.java.annot.Param;
 import net.starlark.java.annot.ParamType;
@@ -223,6 +224,20 @@ 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.cast returns Dict<String, raw Dict>.
+    @SuppressWarnings({"unchecked", "rawtypes"})
+    Dict<String, Sequence> headersUnchecked = (Dict) Dict.cast(x, String.class, Sequence.class, what);
+    ImmutableMap.Builder<String, List<String>> headers = new ImmutableMap.Builder<>();
+
+    for (Map.Entry<String, Sequence> headerEntry : headersUnchecked.entrySet()) {
+      List<String> headerValue = Sequence.cast(headerEntry.getValue(), String.class, "header values").getImmutableList();
+      headers.put(headerEntry.getKey(), headerValue);
+    }
+    return headers.buildOrThrow();
+}
+
   private static ImmutableList<String> checkAllUrls(Iterable<?> urlList) throws EvalException {
     ImmutableList.Builder<String> result = ImmutableList.builder();
 
@@ -265,9 +280,9 @@ private static ImmutableList<URL> getUrls(
             new IOException("Unsupported protocol: " + url.getProtocol()), Transience.PERSISTENT);
       }
       if (!checksumGiven) {
-        if (!Ascii.equalsIgnoreCase("http", url.getProtocol())) {
+        //if (!Ascii.equalsIgnoreCase("http", url.getProtocol())) {
           urls.add(url);
-        }
+        // }
       } else {
         urls.add(url);
       }
@@ -424,6 +439,11 @@ private StructImpl calculateDownloadResult(Optional<Checksum> checksum, Path dow
             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 some of the URLs."),
         @Param(
             name = "integrity",
             defaultValue = "''",
@@ -443,13 +463,16 @@ public StructImpl download(
       Boolean executable,
       Boolean allowFail,
       String canonicalId,
-      Dict<?, ?> authUnchecked, // <String, Dict> expected
+      Dict<?, ?> authUnchecked, // <String, List<String>> expected
+      Dict<?, ?> headersUnchecked, // <String, Dict> expected
       String integrity,
       StarlarkThread thread)
       throws RepositoryFunctionException, EvalException, InterruptedException {
     ImmutableMap<URI, Map<String, List<String>>> authHeaders =
         getAuthHeaders(getAuthContents(authUnchecked, "auth"));
 
+    ImmutableMap<String, List<String>> headers = getHeaderContents(headersUnchecked, "headers");
+
     ImmutableList<URL> urls =
         getUrls(
             url,
@@ -483,6 +506,7 @@ public StructImpl download(
       downloadedPath =
           downloadManager.download(
               urls,
+              headers,
               authHeaders,
               checksum,
               canonicalId,
@@ -595,6 +619,11 @@ public StructImpl 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 some of the URLs."),
         @Param(
             name = "integrity",
             defaultValue = "''",
@@ -626,13 +655,16 @@ public StructImpl downloadAndExtract(
       String stripPrefix,
       Boolean allowFail,
       String canonicalId,
-      Dict<?, ?> auth, // <String, Dict> expected
+      Dict<?, ?> authUnchecked, // <String, Dict> expected
+      Dict<?, ?> headersUnchecked, // <String, Dict> 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(
@@ -681,6 +713,7 @@ public StructImpl downloadAndExtract(
       downloadedPath =
           downloadManager.download(
               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 =

src/test/shell/bazel/remote_helpers.sh

@@ -179,6 +179,25 @@ function serve_timeout() {
   cd -
 }
 
+# Serves a HTTP 200 Ok response with headers dumped into the body
+function serve_file_header_dump() {
+  file_name=served_file.$$
+  cat $1 > "${TEST_TMPDIR}/$file_name"
+  nc_log="${TEST_TMPDIR}/nc.log"
+  rm -f $nc_log
+  touch $nc_log
+  cd "${TEST_TMPDIR}"
+  port_file=server-port.$$
+  rm -f $port_file
+  python3 $python_server always $file_name --dump_headers ${2:-"headers.json"} > $port_file &
+  nc_pid=$!
+  while ! grep started $port_file; do sleep 1; done
+  nc_port=$(head -n 1 $port_file)
+  fileserver_port=$nc_port
+  wait_for_server_startup
+  cd -
+}
+
 # Waits for the SimpleHTTPServer to actually start up before the test is run.
 # Otherwise the entire test can run before the server starts listening for
 # connections, which causes flakes.