1 files changed, 47 insertions, 0 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 197b7721eca..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,8 +39,10 @@ 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;
     int connectionsPerEndpoint = 8;
     int maxStreamsPerConnection = 128;
     FeedClient.RetryStrategy retryStrategy = defaultRetryStrategy;
@@ -48,9 +50,11 @@ public class FeedClientBuilderImpl implements FeedClientBuilder {
     Path certificateFile;
     Path privateKeyFile;
     Path caCertificatesFile;
+    Path proxyCaCertificatesFile;
     Collection<X509Certificate> certificate;
     PrivateKey privateKey;
     Collection<X509Certificate> caCertificates;
+    Collection<X509Certificate> proxyCaCertificates;
     boolean benchmark = true;
     boolean dryrun = false;
     boolean speedTest = false;
@@ -105,6 +109,13 @@ public class FeedClientBuilderImpl implements FeedClientBuilder {
         return this;
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public FeedClientBuilder setProxyHostnameVerifier(HostnameVerifier verifier) {
+        this.proxyHostnameVerifier = requireNonNull(verifier);
+        return this;
+    }
+
     /** Turns off benchmarking. Attempting to get {@link FeedClient#stats()} will result in an exception. */
     @Override
     public FeedClientBuilderImpl noBenchmarking() {
@@ -128,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
@@ -192,6 +215,13 @@ public class FeedClientBuilderImpl implements FeedClientBuilder {
         return this;
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public FeedClientBuilderImpl setProxyCaCertificatesFile(Path caCertificatesFile) {
+        this.proxyCaCertificatesFile = caCertificatesFile;
+        return this;
+    }
+
     /** Overrides JVM default SSL truststore */
     @Override
     public FeedClientBuilderImpl setCaCertificates(Collection<X509Certificate> caCertificates) {
@@ -199,6 +229,13 @@ public class FeedClientBuilderImpl implements FeedClientBuilder {
         return this;
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public FeedClientBuilder setProxyCaCertificates(Collection<X509Certificate> caCertificates) {
+        this.proxyCaCertificates = caCertificates;
+        return null;
+    }
+
     @Override
     public FeedClientBuilderImpl setProxy(URI uri) {
         this.proxy = uri;
@@ -238,6 +275,16 @@ public class FeedClientBuilderImpl implements FeedClientBuilder {
         return sslContextBuilder.build();
     }
 
+    SSLContext constructProxySslContext() throws IOException {
+        SslContextBuilder b = new SslContextBuilder();
+        if (proxyCaCertificatesFile != null) {
+            b.withCaCertificates(proxyCaCertificatesFile);
+        } else if (proxyCaCertificates != null) {
+            b.withCaCertificates(proxyCaCertificates);
+        }
+        return b.build();
+    }
+
     private void validateConfiguration() {
         if (endpoints == null) {
             throw new IllegalArgumentException("At least one endpoint must be provided");