diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkBaseExternalContext.java b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkBaseExternalContext.java index 2d9f73b4fe25a3..e23e2978bdd025 100644 --- a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkBaseExternalContext.java +++ b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkBaseExternalContext.java @@ -223,6 +223,46 @@ private static ImmutableMap>> getAuthHeaders( return res; } + private static ImmutableMap>> getHeaders( + Map>> auth) throws RepositoryFunctionException, EvalException { + ImmutableMap.Builder>> headers = new ImmutableMap.Builder<>(); + for (Map.Entry>> entry : auth.entrySet()) { + try { + URL url = new URL(entry.getKey()); + Map> 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>> getHeaderContents(Dict x, String what) + throws EvalException { + // Dict.cast returns Dict. + @SuppressWarnings({"unchecked", "rawtypes"}) + Map> urlHeaderMapUnchecked = (Map) Dict.cast(x, String.class, Dict.class, what); + + ImmutableMap.Builder>> urlHeadersMap = new ImmutableMap.Builder<>(); + + for (Map.Entry> urlHeaderEntry : urlHeaderMapUnchecked.entrySet()) { + + ImmutableMap.Builder> headers = new ImmutableMap.Builder<>(); + Dict headersUnchecked = Dict.cast(urlHeaderEntry.getValue(), String.class, List.class, "header entry"); + + for (Map.Entry headerEntry : headersUnchecked.entrySet()) { + List headerValue = headerEntry.getValue().stream().map(r -> r.toString()).toList(); + headers.put(headerEntry.getKey(), headerValue); + } + + urlHeadersMap.put(urlHeaderEntry.getKey(), headers.buildOrThrow()); + } + return urlHeadersMap.buildOrThrow(); +} + private static ImmutableList checkAllUrls(Iterable urlList) throws EvalException { ImmutableList.Builder result = ImmutableList.builder(); @@ -424,6 +464,11 @@ private StructImpl calculateDownloadResult(Optional 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 = "''", @@ -444,12 +489,20 @@ public StructImpl download( Boolean allowFail, String canonicalId, Dict authUnchecked, // expected + Dict headersUnchecked, // expected String integrity, StarlarkThread thread) throws RepositoryFunctionException, EvalException, InterruptedException { ImmutableMap>> authHeaders = getAuthHeaders(getAuthContents(authUnchecked, "auth")); + ImmutableMap>> headers = getHeaders(getHeaderContents(headersUnchecked, "headers")); + + ImmutableMap>> allHeaders = new ImmutableMap.Builder>>() + .putAll(authHeaders) + .putAll(headers) + .buildOrThrow(); + ImmutableList urls = getUrls( url, @@ -483,7 +536,7 @@ public StructImpl download( downloadedPath = downloadManager.download( urls, - authHeaders, + allHeaders, checksum, canonicalId, Optional.absent(), @@ -598,6 +651,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 = "''", @@ -629,13 +687,21 @@ public StructImpl downloadAndExtract( String stripPrefix, Boolean allowFail, String canonicalId, - Dict auth, // expected + Dict authUnchecked, // expected + Dict headersUnchecked, // expected String integrity, Dict renameFiles, // expected StarlarkThread thread) throws RepositoryFunctionException, InterruptedException, EvalException { ImmutableMap>> authHeaders = - getAuthHeaders(getAuthContents(auth, "auth")); + getAuthHeaders(getAuthContents(authUnchecked, "auth")); + + ImmutableMap>> headers = getHeaders(getHeaderContents(headersUnchecked, "headers")); + + ImmutableMap>> allHeaders = new ImmutableMap.Builder>>() + .putAll(authHeaders) + .putAll(headers) + .buildOrThrow(); ImmutableList urls = getUrls( @@ -684,7 +750,7 @@ public StructImpl downloadAndExtract( downloadedPath = downloadManager.download( urls, - authHeaders, + allHeaders, checksum, canonicalId, Optional.of(type),