Add support for arbitrary headers to rctx.download[_and_extract]

fixes #17829

Closes #19501.

PiperOrigin-RevId: 588750878
Change-Id: I11c982864135d8e1c4420099ce27f67accd3fe20

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,14 @@ 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,23 @@ 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

@@ -16,6 +16,7 @@
 
 import com.google.auth.Credentials;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.common.io.ByteStreams;
 import com.google.devtools.build.lib.buildeventstream.FetchEvent;
@@ -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,8 @@ 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

@@ -281,6 +281,32 @@ 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()) {
+ ImmutableList<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 = ImmutableList.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 +603,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 = "''",
@@ -607,6 +638,7 @@ public Object download(
  Boolean allowFail,
  String canonicalId,
  Dict<?, ?> authUnchecked, // <String, Dict> expected
+ Dict<?, ?> headersUnchecked, // <String, List<String> | String> expected
  String integrity,
  Boolean block,
  StarlarkThread thread)
@@ -615,6 +647,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 +694,7 @@ public Object download(
  Future<Path> downloadFuture =
  downloadManager.startDownload(
  urls,
+ headers,
  authHeaders,
  checksum,
  canonicalId,
@@ -768,6 +803,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 +839,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 +895,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,15 @@ 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(),