diff options
3 files changed, 26 insertions, 9 deletions
diff --git a/vespa-feed-client-api/src/main/java/ai/vespa/feed/client/FeedClientBuilder.java b/vespa-feed-client-api/src/main/java/ai/vespa/feed/client/FeedClientBuilder.java index 3be773c19f5..b5b6874ded9 100644 --- a/vespa-feed-client-api/src/main/java/ai/vespa/feed/client/FeedClientBuilder.java +++ b/vespa-feed-client-api/src/main/java/ai/vespa/feed/client/FeedClientBuilder.java @@ -84,6 +84,15 @@ public interface FeedClientBuilder { */ FeedClientBuilder addRequestHeader(String name, Supplier<String> valueSupplier); + /** Adds HTTP request header to all proxy requests. */ + FeedClientBuilder addProxyRequestHeader(String name, String value); + + /** + * Adds HTTP request header to all proxy requests. Value {@link Supplier} is invoked for each HTTP request, + * i.e. value can be dynamically updated for each new proxy connection. + */ + FeedClientBuilder addProxyRequestHeader(String name, Supplier<String> valueSupplier); + /** * Overrides default retry strategy. * @see FeedClient.RetryStrategy 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(); |