From d53450e522df70a23f4bd7704fd068b8de24c854 Mon Sep 17 00:00:00 2001
From: Chi Wang <chiwang@google.com>
Date: Mon, 10 Oct 2022 14:40:17 +0200
Subject: [PATCH] Update GrpcRemoteDownloader to only include relevant headers.

---
 .../downloader/GrpcRemoteDownloader.java       | 18 ++++++++++++++----
 .../downloader/GrpcRemoteDownloaderTest.java   |  6 ------
 2 files changed, 14 insertions(+), 10 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 43b2316de00063..f38d112c39ccf9 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
@@ -23,6 +23,7 @@
 import build.bazel.remote.execution.v2.RequestMetadata;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
 import com.google.devtools.build.lib.bazel.repository.downloader.Downloader;
@@ -72,7 +73,8 @@ public class GrpcRemoteDownloader implements AutoCloseable, Downloader {
   private final RemoteCacheClient cacheClient;
   private final RemoteOptions options;
   private final boolean verboseFailures;
-  @Nullable private final Downloader fallbackDownloader;
+  @Nullable
+  private final Downloader fallbackDownloader;
 
   private final AtomicBoolean closed = new AtomicBoolean();
 
@@ -198,7 +200,7 @@ static FetchBlobRequest newFetchBlobRequest(
       requestBuilder.addQualifiers(
           Qualifier.newBuilder()
               .setName(QUALIFIER_AUTH_HEADERS)
-              .setValue(authHeadersJson(authHeaders, includeAllHeaders))
+              .setValue(authHeadersJson(urls, authHeaders, includeAllHeaders))
               .build());
     }
 
@@ -224,10 +226,18 @@ private OutputStream newOutputStream(
     return out;
   }
 
-  private static String authHeadersJson(
+  private static String authHeadersJson(List<URL> urls,
       Map<URI, Map<String, List<String>>> authHeaders, boolean includeAllHeaders) {
+    ImmutableSet<String> hostSet = urls.stream().map(URL::getHost)
+        .collect(ImmutableSet.toImmutableSet());
     Map<String, JsonObject> subObjects = new TreeMap<>();
     for (Map.Entry<URI, Map<String, List<String>>> entry : authHeaders.entrySet()) {
+      URI uri = entry.getKey();
+      // Only add headers that are relevant to the hosts.
+      if (!hostSet.contains(uri.getHost())) {
+        continue;
+      }
+
       JsonObject subObject = new JsonObject();
       Map<String, List<String>> orderedHeaders = new TreeMap<>(entry.getValue());
       for (Map.Entry<String, List<String>> subEntry : orderedHeaders.entrySet()) {
@@ -244,7 +254,7 @@ private static String authHeadersJson(
           }
         }
       }
-      subObjects.put(entry.getKey().toString(), subObject);
+      subObjects.put(uri.toString(), subObject);
     }
 
     JsonObject authHeadersJson = new JsonObject();
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 62440310443cad..9ca60e4710bbae 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
@@ -371,9 +371,6 @@ public void testFetchBlobRequest() throws Exception {
             + "\"http://example.com\":{"
             + "\"Another-Header\":\"another header content\","
             + "\"Some-Header\":\"some header content\""
-            + "},"
-            + "\"http://example.org\":{"
-            + "\"Org-Header\":\"org header content\""
             + "}"
             + "}";
 
@@ -427,9 +424,6 @@ public void testFetchBlobRequestWithAllHeaders() throws Exception {
             + "\"http://example.com\":{"
             + "\"Another-Header\":[\"another header content\",\"even more header content\"],"
             + "\"Some-Header\":[\"some header content\"]"
-            + "},"
-            + "\"http://example.org\":{"
-            + "\"Org-Header\":[\"org header content\",\"and a second one\",\"and a third one\"]"
             + "}"
             + "}";