From 6ac725c0234d187cf9daceaa9df020889ab57b2c Mon Sep 17 00:00:00 2001
From: Sushain Cherivirala <sushain@skc.name>
Date: Tue, 27 Feb 2024 02:44:52 -0800
Subject: [PATCH] Passthrough HTTP headers to remote downloader service

Related to https://github.com/bazelbuild/bazel/issues/17829 and https://github.com/bazelbuild/bazel/commit/2697e0c7d798184cf80f8a5c16db3f0022d63256

I don't love this design but according to the Remote Asset API spec, this is an intended use of qualifiers: https://docs.google.com/document/d/10ari9WtTTSv9bqB_UU-oe2gBtaAA7HyQgkpP-RFP80c/edit#heading=h.sixrlhdnkfoa.

cc @Wyverald @jmillikin

Closes #21490.

PiperOrigin-RevId: 610688317
Change-Id: I272f63a6bc4ea432503003ee907ca012f6d641cf
---
 .../downloader/GrpcRemoteDownloader.java      | 25 +++++++++++++++++--
 .../downloader/GrpcRemoteDownloaderTest.java  | 13 +++++++++-
 2 files changed, 35 insertions(+), 3 deletions(-)

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
index ddc011ef99432d..ed8233f6487797 100644
--- 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
@@ -77,6 +77,11 @@ public class GrpcRemoteDownloader implements AutoCloseable, Downloader {
   private static final String QUALIFIER_CHECKSUM_SRI = "checksum.sri";
   private static final String QUALIFIER_CANONICAL_ID = "bazel.canonical_id";
 
+  // The `:` character is not permitted in an HTTP header name. So, we use it to
+  // delimit the qualifier prefix which denotes an HTTP header qualifer from the
+  // header name itself.
+  private static final String QUALIFIER_HTTP_HEADER_PREFIX = "http_header:";
+
   public GrpcRemoteDownloader(
       String buildRequestId,
       String commandId,
@@ -125,7 +130,7 @@ public void download(
         RemoteActionExecutionContext.create(metadata);
 
     final FetchBlobRequest request =
-        newFetchBlobRequest(options.remoteInstanceName, urls, checksum, canonicalId);
+        newFetchBlobRequest(options.remoteInstanceName, urls, checksum, canonicalId, headers);
     try {
       FetchBlobResponse response =
           retrier.execute(
@@ -169,12 +174,17 @@ public void download(
 
   @VisibleForTesting
   static FetchBlobRequest newFetchBlobRequest(
-      String instanceName, List<URL> urls, Optional<Checksum> checksum, String canonicalId) {
+      String instanceName,
+      List<URL> urls,
+      Optional<Checksum> checksum,
+      String canonicalId,
+      Map<String, List<String>> headers) {
     FetchBlobRequest.Builder requestBuilder =
         FetchBlobRequest.newBuilder().setInstanceName(instanceName);
     for (URL url : urls) {
       requestBuilder.addUris(url.toString());
     }
+
     if (checksum.isPresent()) {
       requestBuilder.addQualifiers(
           Qualifier.newBuilder()
@@ -182,11 +192,22 @@ static FetchBlobRequest newFetchBlobRequest(
               .setValue(checksum.get().toSubresourceIntegrity())
               .build());
     }
+
     if (!Strings.isNullOrEmpty(canonicalId)) {
       requestBuilder.addQualifiers(
           Qualifier.newBuilder().setName(QUALIFIER_CANONICAL_ID).setValue(canonicalId).build());
     }
 
+    for (Map.Entry<String, List<String>> entry : headers.entrySet()) {
+      // https://www.rfc-editor.org/rfc/rfc9110.html#name-field-order permits
+      // merging the field-values with a comma.
+      requestBuilder.addQualifiers(
+          Qualifier.newBuilder()
+              .setName(QUALIFIER_HTTP_HEADER_PREFIX + entry.getKey())
+              .setValue(String.join(",", entry.getValue()))
+              .build());
+    }
+
     return requestBuilder.build();
   }
 
diff --git a/src/test/java/com/google/devtools/build/lib/remote/downloader/GrpcRemoteDownloaderTest.java b/src/test/java/com/google/devtools/build/lib/remote/downloader/GrpcRemoteDownloaderTest.java
index f881d862800e83..4e58b6ed7962eb 100644
--- a/src/test/java/com/google/devtools/build/lib/remote/downloader/GrpcRemoteDownloaderTest.java
+++ b/src/test/java/com/google/devtools/build/lib/remote/downloader/GrpcRemoteDownloaderTest.java
@@ -351,7 +351,10 @@ public void testFetchBlobRequest() throws Exception {
             Optional.<Checksum>of(
                 Checksum.fromSubresourceIntegrity(
                     "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=")),
-            "canonical ID");
+            "canonical ID",
+            ImmutableMap.of(
+                "Authorization", ImmutableList.of("Basic Zm9vOmJhcg=="),
+                "X-Custom-Token", ImmutableList.of("foo", "bar")));
 
     assertThat(request)
         .isEqualTo(
@@ -366,6 +369,14 @@ public void testFetchBlobRequest() throws Exception {
                         .setValue("sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA="))
                 .addQualifiers(
                     Qualifier.newBuilder().setName("bazel.canonical_id").setValue("canonical ID"))
+                .addQualifiers(
+                    Qualifier.newBuilder()
+                        .setName("http_header:Authorization")
+                        .setValue("Basic Zm9vOmJhcg=="))
+                .addQualifiers(
+                    Qualifier.newBuilder()
+                        .setName("http_header:X-Custom-Token")
+                        .setValue("foo,bar"))
                 .build());
   }
 }