Properly model proxy headers

2 files changed, 17 insertions, 9 deletions

diff --git a/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/FeedClientBuilderImpl.java b/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/FeedClientBuilderImpl.java
index d00ee6e6b04..3b7deb52b3b 100644
--- a/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/FeedClientBuilderImpl.java
+++ b/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/FeedClientBuilderImpl.java
@@ -39,6 +39,7 @@ public class FeedClientBuilderImpl implements FeedClientBuilder {
 
     List<URI> endpoints;
     final Map<String, Supplier<String>> requestHeaders = new HashMap<>();
+    final Map<String, Supplier<String>> proxyRequestHeaders = new HashMap<>();
     SSLContext sslContext;
     HostnameVerifier hostnameVerifier;
     HostnameVerifier proxyHostnameVerifier;
@@ -138,6 +139,18 @@ public class FeedClientBuilderImpl implements FeedClientBuilder {
         return this;
     }
 
+    @Override
+    public FeedClientBuilder addProxyRequestHeader(String name, String value) {
+        this.proxyRequestHeaders.put(requireNonNull(name), () -> requireNonNull(value));
+        return this;
+    }
+
+    @Override
+    public FeedClientBuilder addProxyRequestHeader(String name, Supplier<String> valueSupplier) {
+        this.proxyRequestHeaders.put(requireNonNull(name), requireNonNull(valueSupplier));
+        return this;
+    }
+
     /**
      * Overrides default retry strategy.
      * @see FeedClient.RetryStrategy
diff --git a/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/JettyCluster.java b/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/JettyCluster.java
index 9261de7ea9b..1a125ebfbb5 100644
--- a/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/JettyCluster.java
+++ b/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/JettyCluster.java
@@ -85,9 +85,7 @@ class JettyCluster implements Cluster {
                 Request jettyReq = client.newRequest(URI.create(endpoint.uri + req.path()))
                         .version(HttpVersion.HTTP_2)
                         .method(HttpMethod.fromString(req.method()))
-                        .headers(hs -> req.headers().forEach((k, v) -> {
-                            if (!isProxyHeader(k)) hs.add(k, v.get());
-                        }))
+                        .headers(hs -> req.headers().forEach((k, v) -> hs.add(k, v.get())))
                         .idleTimeout(IDLE_TIMEOUT.toMillis(), MILLISECONDS)
                         .timeout(reqTimeoutMillis, MILLISECONDS);
                 if (req.body() != null) {
@@ -187,20 +185,17 @@ class JettyCluster implements Cluster {
             httpClient.getProxyConfiguration().addProxy(
                     new HttpProxy(address, false, new Origin.Protocol(Collections.singletonList("h2c"), false)));
         }
-        Map<String, Supplier<String>> proxyHeaders = new TreeMap<>();
-        b.requestHeaders.forEach((k, v) -> { if (isProxyHeader(k)) proxyHeaders.put(k, v); });
-        if (!proxyHeaders.isEmpty()) {
+        Map<String, Supplier<String>> proxyHeadersCopy = new TreeMap<>(b.proxyRequestHeaders);
+        if (!proxyHeadersCopy.isEmpty()) {
             httpClient.getAuthenticationStore().addAuthenticationResult(new Authentication.Result() {
                 @Override public URI getURI() { return URI.create(endpointUri(b.proxy)); }
                 @Override public void apply(Request r) {
-                    r.headers(hs -> proxyHeaders.forEach((k, v) -> hs.add(k, v.get())));
+                    r.headers(hs -> proxyHeadersCopy.forEach((k, v) -> hs.add(k, v.get())));
                 }
             });
         }
     }
 
-    private static boolean isProxyHeader(String h) { return h.equalsIgnoreCase(HttpHeader.PROXY_AUTHORIZATION.asString()); }
-
     private static Endpoint findLeastBusyEndpoint(List<Endpoint> endpoints) {
         Endpoint leastBusy = endpoints.get(0);
         int minInflight = leastBusy.inflight.get();