Skip to content

Commit 2697e0c

Browse files
Wyveraldthesayyn
Wyverald
and
thesayyn
committed
[7.1.0] Add support for arbitrary headers to rctx.download[_and_extract] (#20979)
fixes #17829 Closes #19501. PiperOrigin-RevId: 588750878 Change-Id: I11c982864135d8e1c4420099ce27f67accd3fe20 Co-authored-by: thesayyn <[email protected]>
1 parent 4e51534 commit 2697e0c

12 files changed

Lines changed: 319 additions & 16 deletions

File tree

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ public void setDelegate(@Nullable Downloader delegate) {
4747
@Override
4848
public void download(
4949
List<URL> urls,
50+
Map<String, List<String>> headers,
5051
Credentials credentials,
5152
Optional<Checksum> checksum,
5253
String canonicalId,
@@ -60,6 +61,14 @@ public void download(
6061
downloader = delegate;
6162
}
6263
downloader.download(
63-
urls, credentials, checksum, canonicalId, destination, eventHandler, clientEnv, type);
64+
urls,
65+
headers,
66+
credentials,
67+
checksum,
68+
canonicalId,
69+
destination,
70+
eventHandler,
71+
clientEnv,
72+
type);
6473
}
6574
}

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,7 @@ public void setCredentialFactory(CredentialFactory credentialFactory) {
116116

117117
public Future<Path> startDownload(
118118
List<URL> originalUrls,
119+
Map<String, List<String>> headers,
119120
Map<URI, Map<String, List<String>>> authHeaders,
120121
Optional<Checksum> checksum,
121122
String canonicalId,
@@ -129,6 +130,7 @@ public Future<Path> startDownload(
129130
try (SilentCloseable c = Profiler.instance().profile("fetching: " + context)) {
130131
return downloadInExecutor(
131132
originalUrls,
133+
headers,
132134
authHeaders,
133135
checksum,
134136
canonicalId,
@@ -154,6 +156,7 @@ public Path finalizeDownload(Future<Path> download) throws IOException, Interrup
154156

155157
public Path download(
156158
List<URL> originalUrls,
159+
Map<String, List<String>> headers,
157160
Map<URI, Map<String, List<String>>> authHeaders,
158161
Optional<Checksum> checksum,
159162
String canonicalId,
@@ -166,6 +169,7 @@ public Path download(
166169
Future<Path> future =
167170
startDownload(
168171
originalUrls,
172+
headers,
169173
authHeaders,
170174
checksum,
171175
canonicalId,
@@ -197,6 +201,7 @@ public Path download(
197201
*/
198202
private Path downloadInExecutor(
199203
List<URL> originalUrls,
204+
Map<String, List<String>> headers,
200205
Map<URI, Map<String, List<String>>> authHeaders,
201206
Optional<Checksum> checksum,
202207
String canonicalId,
@@ -339,6 +344,7 @@ private Path downloadInExecutor(
339344
try {
340345
downloader.download(
341346
rewrittenUrls,
347+
headers,
342348
credentialFactory.create(rewrittenAuthHeaders),
343349
checksum,
344350
canonicalId,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ public interface Downloader {
4242
*/
4343
void download(
4444
List<URL> urls,
45+
Map<String, List<String>> headers,
4546
Credentials credentials,
4647
Optional<Checksum> checksum,
4748
String canonicalId,

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

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ final class HttpConnectorMultiplexer {
7575
}
7676

7777
public HttpStream connect(URL url, Optional<Checksum> checksum) throws IOException {
78-
return connect(url, checksum, StaticCredentials.EMPTY, Optional.empty());
78+
return connect(url, checksum, ImmutableMap.of(), StaticCredentials.EMPTY, Optional.empty());
7979
}
8080

8181
/**
@@ -96,14 +96,23 @@ public HttpStream connect(URL url, Optional<Checksum> checksum) throws IOExcepti
9696
* @throws IllegalArgumentException if {@code urls} is empty or has an unsupported protocol
9797
*/
9898
public HttpStream connect(
99-
URL url, Optional<Checksum> checksum, Credentials credentials, Optional<String> type)
99+
URL url,
100+
Optional<Checksum> checksum,
101+
Map<String, List<String>> headers,
102+
Credentials credentials,
103+
Optional<String> type)
100104
throws IOException {
101105
Preconditions.checkArgument(HttpUtils.isUrlSupportedByDownloader(url));
102106
if (Thread.interrupted()) {
103107
throw new InterruptedIOException();
104108
}
109+
ImmutableMap.Builder<String, List<String>> baseHeaders = new ImmutableMap.Builder<>();
110+
baseHeaders.putAll(headers);
111+
// REQUEST_HEADERS should not be overridable by user provided headers
112+
baseHeaders.putAll(REQUEST_HEADERS);
113+
105114
Function<URL, ImmutableMap<String, List<String>>> headerFunction =
106-
getHeaderFunction(REQUEST_HEADERS, credentials, eventHandler);
115+
getHeaderFunction(baseHeaders.buildKeepingLast(), credentials, eventHandler);
107116
URLConnection connection = connector.connect(url, headerFunction);
108117
return httpStreamFactory.create(
109118
connection,

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
import com.google.auth.Credentials;
1818
import com.google.common.collect.ImmutableList;
19+
import com.google.common.collect.ImmutableMap;
1920
import com.google.common.collect.Iterables;
2021
import com.google.common.io.ByteStreams;
2122
import com.google.devtools.build.lib.buildeventstream.FetchEvent;
@@ -75,6 +76,7 @@ public void setMaxRetryTimeout(Duration maxRetryTimeout) {
7576
@Override
7677
public void download(
7778
List<URL> urls,
79+
Map<String, List<String>> headers,
7880
Credentials credentials,
7981
Optional<Checksum> checksum,
8082
String canonicalId,
@@ -94,7 +96,7 @@ public void download(
9496
for (URL url : urls) {
9597
SEMAPHORE.acquire();
9698

97-
try (HttpStream payload = multiplexer.connect(url, checksum, credentials, type);
99+
try (HttpStream payload = multiplexer.connect(url, checksum, headers, credentials, type);
98100
OutputStream out = destination.getOutputStream()) {
99101
try {
100102
ByteStreams.copy(payload, out);
@@ -153,7 +155,8 @@ public byte[] downloadAndReadOneUrl(
153155
ByteArrayOutputStream out = new ByteArrayOutputStream();
154156
SEMAPHORE.acquire();
155157
try (HttpStream payload =
156-
multiplexer.connect(url, Optional.empty(), credentials, Optional.empty())) {
158+
multiplexer.connect(
159+
url, Optional.empty(), ImmutableMap.of(), credentials, Optional.empty())) {
157160
ByteStreams.copy(payload, out);
158161
} catch (SocketTimeoutException e) {
159162
// SocketTimeoutExceptions are InterruptedIOExceptions; however they do not signify

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

Lines changed: 46 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -291,6 +291,32 @@ private static ImmutableMap<URI, Map<String, List<String>>> getAuthHeaders(
291291
return res;
292292
}
293293

294+
private static ImmutableMap<String, List<String>> getHeaderContents(Dict<?, ?> x, String what)
295+
throws EvalException {
296+
Dict<String, Object> headersUnchecked =
297+
(Dict<String, Object>) Dict.cast(x, String.class, Object.class, what);
298+
ImmutableMap.Builder<String, List<String>> headers = new ImmutableMap.Builder<>();
299+
300+
for (Map.Entry<String, Object> entry : headersUnchecked.entrySet()) {
301+
ImmutableList<String> headerValue;
302+
Object valueUnchecked = entry.getValue();
303+
if (valueUnchecked instanceof Sequence) {
304+
headerValue =
305+
Sequence.cast(valueUnchecked, String.class, "header values").getImmutableList();
306+
} else if (valueUnchecked instanceof String) {
307+
headerValue = ImmutableList.of(valueUnchecked.toString());
308+
} else {
309+
throw new EvalException(
310+
String.format(
311+
"%s argument must be a dict whose keys are string and whose values are either"
312+
+ " string or sequence of string",
313+
what));
314+
}
315+
headers.put(entry.getKey(), headerValue);
316+
}
317+
return headers.buildOrThrow();
318+
}
319+
294320
private static ImmutableList<String> checkAllUrls(Iterable<?> urlList) throws EvalException {
295321
ImmutableList.Builder<String> result = ImmutableList.builder();
296322

@@ -587,6 +613,11 @@ private StructImpl completeDownload(PendingDownload pendingDownload)
587613
defaultValue = "{}",
588614
named = true,
589615
doc = "An optional dict specifying authentication information for some of the URLs."),
616+
@Param(
617+
name = "headers",
618+
defaultValue = "{}",
619+
named = true,
620+
doc = "An optional dict specifying http headers for all URLs."),
590621
@Param(
591622
name = "integrity",
592623
defaultValue = "''",
@@ -617,6 +648,7 @@ public Object download(
617648
Boolean allowFail,
618649
String canonicalId,
619650
Dict<?, ?> authUnchecked, // <String, Dict> expected
651+
Dict<?, ?> headersUnchecked, // <String, List<String> | String> expected
620652
String integrity,
621653
Boolean block,
622654
StarlarkThread thread)
@@ -625,6 +657,8 @@ public Object download(
625657
ImmutableMap<URI, Map<String, List<String>>> authHeaders =
626658
getAuthHeaders(getAuthContents(authUnchecked, "auth"));
627659

660+
ImmutableMap<String, List<String>> headers = getHeaderContents(headersUnchecked, "headers");
661+
628662
ImmutableList<URL> urls =
629663
getUrls(
630664
url,
@@ -670,6 +704,7 @@ public Object download(
670704
Future<Path> downloadFuture =
671705
downloadManager.startDownload(
672706
urls,
707+
headers,
673708
authHeaders,
674709
checksum,
675710
canonicalId,
@@ -778,6 +813,11 @@ public Object download(
778813
defaultValue = "{}",
779814
named = true,
780815
doc = "An optional dict specifying authentication information for some of the URLs."),
816+
@Param(
817+
name = "headers",
818+
defaultValue = "{}",
819+
named = true,
820+
doc = "An optional dict specifying http headers for all URLs."),
781821
@Param(
782822
name = "integrity",
783823
defaultValue = "''",
@@ -809,13 +849,16 @@ public StructImpl downloadAndExtract(
809849
String stripPrefix,
810850
Boolean allowFail,
811851
String canonicalId,
812-
Dict<?, ?> auth, // <String, Dict> expected
852+
Dict<?, ?> authUnchecked, // <String, Dict> expected
853+
Dict<?, ?> headersUnchecked, // <String, List<String> | String> expected
813854
String integrity,
814855
Dict<?, ?> renameFiles, // <String, String> expected
815856
StarlarkThread thread)
816857
throws RepositoryFunctionException, InterruptedException, EvalException {
817858
ImmutableMap<URI, Map<String, List<String>>> authHeaders =
818-
getAuthHeaders(getAuthContents(auth, "auth"));
859+
getAuthHeaders(getAuthContents(authUnchecked, "auth"));
860+
861+
ImmutableMap<String, List<String>> headers = getHeaderContents(headersUnchecked, "headers");
819862

820863
ImmutableList<URL> urls =
821864
getUrls(
@@ -862,6 +905,7 @@ public StructImpl downloadAndExtract(
862905
Future<Path> pendingDownload =
863906
downloadManager.startDownload(
864907
urls,
908+
headers,
865909
authHeaders,
866910
checksum,
867911
canonicalId,

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,7 @@ public void close() {
110110
@Override
111111
public void download(
112112
List<URL> urls,
113+
Map<String, List<String>> headers,
113114
Credentials credentials,
114115
Optional<Checksum> checksum,
115116
String canonicalId,
@@ -154,7 +155,15 @@ public void download(
154155
eventHandler.handle(
155156
Event.warn("Remote Cache: " + Utils.grpcAwareErrorMessage(e, verboseFailures)));
156157
fallbackDownloader.download(
157-
urls, credentials, checksum, canonicalId, destination, eventHandler, clientEnv, type);
158+
urls,
159+
headers,
160+
credentials,
161+
checksum,
162+
canonicalId,
163+
destination,
164+
eventHandler,
165+
clientEnv,
166+
type);
158167
}
159168
}
160169

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

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,7 @@ public void downloadFrom1UrlOk() throws IOException, InterruptedException {
117117
Collections.singletonList(
118118
new URL(String.format("http://localhost:%d/foo", server.getLocalPort()))),
119119
Collections.emptyMap(),
120+
Collections.emptyMap(),
120121
Optional.empty(),
121122
"testCanonicalId",
122123
Optional.empty(),
@@ -181,6 +182,7 @@ public void downloadFrom2UrlsFirstOk() throws IOException, InterruptedException
181182
downloadManager.download(
182183
urls,
183184
Collections.emptyMap(),
185+
Collections.emptyMap(),
184186
Optional.empty(),
185187
"testCanonicalId",
186188
Optional.empty(),
@@ -248,6 +250,7 @@ public void downloadFrom2UrlsFirstSocketTimeoutOnBodyReadSecondOk()
248250
downloadManager.download(
249251
urls,
250252
Collections.emptyMap(),
253+
Collections.emptyMap(),
251254
Optional.empty(),
252255
"testCanonicalId",
253256
Optional.empty(),
@@ -317,6 +320,7 @@ public void downloadFrom2UrlsBothSocketTimeoutDuringBodyRead()
317320
downloadManager.download(
318321
urls,
319322
Collections.emptyMap(),
323+
Collections.emptyMap(),
320324
Optional.empty(),
321325
"testCanonicalId",
322326
Optional.empty(),
@@ -371,6 +375,7 @@ public void downloadOneUrl_ok() throws IOException, InterruptedException {
371375
httpDownloader.download(
372376
Collections.singletonList(
373377
new URL(String.format("http://localhost:%d/foo", server.getLocalPort()))),
378+
Collections.emptyMap(),
374379
StaticCredentials.EMPTY,
375380
Optional.empty(),
376381
"testCanonicalId",
@@ -410,6 +415,7 @@ public void downloadOneUrl_notFound() throws IOException, InterruptedException {
410415
httpDownloader.download(
411416
Collections.singletonList(
412417
new URL(String.format("http://localhost:%d/foo", server.getLocalPort()))),
418+
Collections.emptyMap(),
413419
StaticCredentials.EMPTY,
414420
Optional.empty(),
415421
"testCanonicalId",
@@ -470,6 +476,7 @@ public void downloadTwoUrls_firstNotFoundAndSecondOk() throws IOException, Inter
470476
Path destination = fs.getPath(workingDir.newFile().getAbsolutePath());
471477
httpDownloader.download(
472478
urls,
479+
Collections.emptyMap(),
473480
StaticCredentials.EMPTY,
474481
Optional.empty(),
475482
"testCanonicalId",
@@ -564,13 +571,14 @@ public void download_contentLengthMismatch_propagateErrorIfNotRetry() throws Exc
564571
throw new ContentLengthMismatchException(0, data.length);
565572
})
566573
.when(downloader)
567-
.download(any(), any(), any(), any(), any(), any(), any(), any());
574+
.download(any(), any(), any(), any(), any(), any(), any(), any(), any());
568575

569576
assertThrows(
570577
ContentLengthMismatchException.class,
571578
() ->
572579
downloadManager.download(
573580
ImmutableList.of(new URL("http://localhost")),
581+
Collections.emptyMap(),
574582
ImmutableMap.of(),
575583
Optional.empty(),
576584
"testCanonicalId",
@@ -597,20 +605,21 @@ public void download_contentLengthMismatch_retries() throws Exception {
597605
if (times.getAndIncrement() < 3) {
598606
throw new ContentLengthMismatchException(0, data.length);
599607
}
600-
Path output = invocationOnMock.getArgument(4, Path.class);
608+
Path output = invocationOnMock.getArgument(5, Path.class);
601609
try (OutputStream outputStream = output.getOutputStream()) {
602610
ByteStreams.copy(new ByteArrayInputStream(data), outputStream);
603611
}
604612

605613
return null;
606614
})
607615
.when(downloader)
608-
.download(any(), any(), any(), any(), any(), any(), any(), any());
616+
.download(any(), any(), any(), any(), any(), any(), any(), any(), any());
609617

610618
Path result =
611619
downloadManager.download(
612620
ImmutableList.of(new URL("http://localhost")),
613621
ImmutableMap.of(),
622+
ImmutableMap.of(),
614623
Optional.empty(),
615624
"testCanonicalId",
616625
Optional.empty(),

0 commit comments

Comments
 (0)