refactor

bazelbuild · Oct 25, 2023 · e89950d · e89950d
1 parent 09ff174
commit e89950d
Show file tree

Hide file tree

Showing 7 changed files with 33 additions and 55 deletions.
diff --git a/.../java/com/google/devtools/build/lib/bazel/repository/downloader/DelegatingDownloader.java b/.../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);
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/downloader/DownloadManager.java b/src/main/java/com/google/devtools/build/lib/bazel/repository/downloader/DownloadManager.java
@@ -122,6 +122,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,
@@ -267,6 +268,7 @@ public Path download(
       try {
         downloader.download(
             rewrittenUrls,
+            headers,
             credentialFactory.create(rewrittenAuthHeaders),
             checksum,
             canonicalId,

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/downloader/Downloader.java b/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,

diff --git a/...a/com/google/devtools/build/lib/bazel/repository/downloader/HttpConnectorMultiplexer.java b/...a/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, Map.of(), StaticCredentials.EMPTY, Optional.empty());
   }
 
   /**
@@ -96,14 +96,18 @@ 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();
     }
+    Map<String, List<String>> baseHeaders = new HashMap<>(headers);
+    // REQUEST_HEADERS should not be overriable by user provided headers
+    baseHeaders.putAll(REQUEST_HEADERS);
+
     Function<URL, ImmutableMap<String, List<String>>> headerFunction =
-        getHeaderFunction(REQUEST_HEADERS, credentials);
+        getHeaderFunction(baseHeaders, credentials);
     URLConnection connection = connector.connect(url, headerFunction);
     return httpStreamFactory.create(
         connection,

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/downloader/HttpDownloader.java b/src/main/java/com/google/devtools/build/lib/bazel/repository/downloader/HttpDownloader.java
@@ -15,6 +15,8 @@
 package com.google.devtools.build.lib.bazel.repository.downloader;
 
 import com.google.auth.Credentials;
+import com.google.common.base.Optional;
+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 +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,
@@ -93,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);
@@ -152,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

diff --git a/.../com/google/devtools/build/lib/bazel/repository/starlark/StarlarkBaseExternalContext.java b/.../com/google/devtools/build/lib/bazel/repository/starlark/StarlarkBaseExternalContext.java
@@ -223,44 +223,18 @@ private static ImmutableMap<URI, Map<String, List<String>>> getAuthHeaders(
     return res;
   }
 
-  private static ImmutableMap<URI, Map<String, List<String>>> getHeaders(
-    Map<String, Map<String, List<String>>> auth) throws RepositoryFunctionException, EvalException {
-    ImmutableMap.Builder<URI, Map<String, List<String>>> headers = new ImmutableMap.Builder<>();
-    for (Map.Entry<String, Map<String, List<String>>> entry : auth.entrySet()) {
-      try {
-        URL url = new URL(entry.getKey());
-        Map<String, List<String>> headerMap = entry.getValue();
-        headers.put(url.toURI(), headerMap);
-      } catch (MalformedURLException e) {
-        throw new RepositoryFunctionException(e, Transience.PERSISTENT);
-      } catch (URISyntaxException e) {
-        throw new EvalException(e);
-      }
-    }
-    return headers.buildOrThrow();
-  }
-
-  private static ImmutableMap<String, Map<String, List<String>>> getHeaderContents(Dict<?, ?> x, String what)
+  private static ImmutableMap<String, List<String>> getHeaderContents(Dict<?, ?> x, String what)
   throws EvalException {
     // Dict.cast returns Dict<String, raw Dict>.
     @SuppressWarnings({"unchecked", "rawtypes"})
-    Map<String, Dict<?, ?>> urlHeaderMapUnchecked = (Map) Dict.cast(x, String.class, Dict.class, what);
-
-    ImmutableMap.Builder<String, Map<String, List<String>>> urlHeadersMap = new ImmutableMap.Builder<>();
-
-    for (Map.Entry<String, Dict<?, ?>> urlHeaderEntry : urlHeaderMapUnchecked.entrySet()) {
-
-      ImmutableMap.Builder<String, List<String>> headers = new ImmutableMap.Builder<>();
-      Dict<String, List> headersUnchecked = Dict.cast(urlHeaderEntry.getValue(), String.class, List.class, "header entry");
-
-      for (Map.Entry<String, List> headerEntry : headersUnchecked.entrySet()) {
-        List<String> headerValue = headerEntry.getValue().stream().map(r -> r.toString()).toList();
-        headers.put(headerEntry.getKey(), headerValue);
-      }
+    Dict<String, List> headersUnchecked = (Dict) Dict.cast(x, String.class, List.class, what);
+    ImmutableMap.Builder<String, List<String>> headers = new ImmutableMap.Builder<>();
 
-      urlHeadersMap.put(urlHeaderEntry.getKey(), headers.buildOrThrow());
+    for (Map.Entry<String, List> headerEntry : headersUnchecked.entrySet()) {
+      List<String> headerValue = headerEntry.getValue().stream().map(r -> r.toString()).toList();
+      headers.put(headerEntry.getKey(), headerValue);
     }
-    return urlHeadersMap.buildOrThrow();
+    return headers.buildOrThrow();
 }
 
   private static ImmutableList<String> checkAllUrls(Iterable<?> urlList) throws EvalException {
@@ -488,20 +462,15 @@ 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<URI, Map<String, List<String>>> headers = getHeaders(getHeaderContents(headersUnchecked, "headers"));
-
-    ImmutableMap<URI, Map<String, List<String>>> allHeaders = new ImmutableMap.Builder<URI, Map<String, List<String>>>()
-        .putAll(authHeaders)
-        .putAll(headers)
-        .buildOrThrow();
+    ImmutableMap<String, List<String>> headers = getHeaderContents(headersUnchecked, "headers");
 
     ImmutableList<URL> urls =
         getUrls(
@@ -536,7 +505,8 @@ public StructImpl download(
       downloadedPath =
           downloadManager.download(
               urls,
-              allHeaders,
+              headers,
+              authHeaders,
               checksum,
               canonicalId,
               Optional.<String>empty(),
@@ -696,12 +666,7 @@ public StructImpl downloadAndExtract(
     ImmutableMap<URI, Map<String, List<String>>> authHeaders =
         getAuthHeaders(getAuthContents(authUnchecked, "auth"));
 
-    ImmutableMap<URI, Map<String, List<String>>> headers = getHeaders(getHeaderContents(headersUnchecked, "headers"));
-
-    ImmutableMap<URI, Map<String, List<String>>> allHeaders = new ImmutableMap.Builder<URI, Map<String, List<String>>>()
-        .putAll(authHeaders)
-        .putAll(headers)
-        .buildOrThrow();
+    ImmutableMap<String, List<String>> headers = getHeaderContents(headersUnchecked, "headers");
 
     ImmutableList<URL> urls =
         getUrls(
@@ -750,7 +715,8 @@ public StructImpl downloadAndExtract(
       downloadedPath =
           downloadManager.download(
               urls,
-              allHeaders,
+              headers,
+              authHeaders,
               checksum,
               canonicalId,
               Optional.of(type),

diff --git a/src/main/java/com/google/devtools/build/lib/remote/downloader/GrpcRemoteDownloader.java b/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);
     }
   }