Support setting CA and hostname verifier for proxy

2 files changed, 40 insertions, 3 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..d00ee6e6b04 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
@@ -41,6 +41,7 @@ public class FeedClientBuilderImpl implements FeedClientBuilder {
     final Map<String, Supplier<String>> requestHeaders = new HashMap<>();
     SSLContext sslContext;
     HostnameVerifier hostnameVerifier;
+    HostnameVerifier proxyHostnameVerifier;
     int connectionsPerEndpoint = 8;
     int maxStreamsPerConnection = 128;
     FeedClient.RetryStrategy retryStrategy = defaultRetryStrategy;
@@ -48,9 +49,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 +108,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() {
@@ -192,6 +202,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 +216,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 +262,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");
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 8d8a695e091..9261de7ea9b 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
@@ -174,9 +174,12 @@ class JettyCluster implements Cluster {
         Origin.Address address = new Origin.Address(b.proxy.getHost(), b.proxy.getPort());
         if (b.proxy.getScheme().equals("https")) {
             SslContextFactory.Client proxySslCtxFactory = new SslContextFactory.Client();
-            if (b.hostnameVerifier != null) proxySslCtxFactory.setHostnameVerifier(b.hostnameVerifier);
-            // Disable built-in hostname verification in the JDK's TLS implementation
-            proxySslCtxFactory.setEndpointIdentificationAlgorithm(null);
+            if (b.proxyHostnameVerifier != null) {
+                proxySslCtxFactory.setHostnameVerifier(b.proxyHostnameVerifier);
+                // Disable built-in hostname verification in the JDK's TLS implementation
+                proxySslCtxFactory.setEndpointIdentificationAlgorithm(null);
+            }
+            proxySslCtxFactory.setSslContext(b.constructProxySslContext());
             try { proxySslCtxFactory.start(); } catch (Exception e) { throw new IOException(e); }
             httpClient.getProxyConfiguration().addProxy(
                     new HttpProxy(address, proxySslCtxFactory, new Origin.Protocol(Collections.singletonList("h2"), false)));