Merge pull request #27751 from vespa-engine/bjorncs/vespa-feed-client

Bjorncs/vespa feed client

6 files changed, 102 insertions, 27 deletions

diff --git a/vespa-feed-client-api/abi-spec.json b/vespa-feed-client-api/abi-spec.json
index 469928d819c..e8c2b4a3c9e 100644
--- a/vespa-feed-client-api/abi-spec.json
+++ b/vespa-feed-client-api/abi-spec.json
@@ -146,9 +146,12 @@
       "public abstract ai.vespa.feed.client.FeedClientBuilder setMaxStreamPerConnection(int)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setSslContext(javax.net.ssl.SSLContext)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setHostnameVerifier(javax.net.ssl.HostnameVerifier)",
+      "public abstract ai.vespa.feed.client.FeedClientBuilder setProxyHostnameVerifier(javax.net.ssl.HostnameVerifier)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder noBenchmarking()",
       "public abstract ai.vespa.feed.client.FeedClientBuilder addRequestHeader(java.lang.String, java.lang.String)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder addRequestHeader(java.lang.String, java.util.function.Supplier)",
+      "public abstract ai.vespa.feed.client.FeedClientBuilder addProxyRequestHeader(java.lang.String, java.lang.String)",
+      "public abstract ai.vespa.feed.client.FeedClientBuilder addProxyRequestHeader(java.lang.String, java.util.function.Supplier)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setRetryStrategy(ai.vespa.feed.client.FeedClient$RetryStrategy)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setCircuitBreaker(ai.vespa.feed.client.FeedClient$CircuitBreaker)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setCertificate(java.nio.file.Path, java.nio.file.Path)",
@@ -157,7 +160,9 @@
       "public abstract ai.vespa.feed.client.FeedClientBuilder setDryrun(boolean)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setSpeedTest(boolean)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setCaCertificatesFile(java.nio.file.Path)",
+      "public abstract ai.vespa.feed.client.FeedClientBuilder setProxyCaCertificatesFile(java.nio.file.Path)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setCaCertificates(java.util.Collection)",
+      "public abstract ai.vespa.feed.client.FeedClientBuilder setProxyCaCertificates(java.util.Collection)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setEndpointUris(java.util.List)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setProxy(java.net.URI)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setCompression(ai.vespa.feed.client.FeedClientBuilder$Compression)",
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 02dabaf9ef8..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
@@ -69,6 +69,9 @@ public interface FeedClientBuilder {
     /** Sets {@link HostnameVerifier} instance (e.g for disabling default SSL hostname verification). */
     FeedClientBuilder setHostnameVerifier(HostnameVerifier verifier);
 
+    /** Sets {@link HostnameVerifier} instance for proxy (e.g for disabling default SSL hostname verification). */
+    FeedClientBuilder setProxyHostnameVerifier(HostnameVerifier verifier);
+
     /** Turns off benchmarking. Attempting to get {@link FeedClient#stats()} will result in an exception. */
     FeedClientBuilder noBenchmarking();
 
@@ -81,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
@@ -114,9 +126,18 @@ public interface FeedClientBuilder {
      */
     FeedClientBuilder setCaCertificatesFile(Path caCertificatesFile);
 
+    /**
+     * Overrides JVM default SSL truststore for proxy
+     * @param caCertificatesFile Path to PEM encoded file containing trusted certificates
+     */
+    FeedClientBuilder setProxyCaCertificatesFile(Path caCertificatesFile);
+
     /** Overrides JVM default SSL truststore */
     FeedClientBuilder setCaCertificates(Collection<X509Certificate> caCertificates);
 
+    /** Overrides JVM default SSL truststore for proxy */
+    FeedClientBuilder setProxyCaCertificates(Collection<X509Certificate> caCertificates);
+
     /** Overrides endpoint URIs for this client */
     FeedClientBuilder setEndpointUris(List<URI> endpoints);
 
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");
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 30dc1ab0d07..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
@@ -8,10 +8,12 @@ import org.eclipse.jetty.client.HttpClient;
 import org.eclipse.jetty.client.HttpProxy;
 import org.eclipse.jetty.client.MultiplexConnectionPool;
 import org.eclipse.jetty.client.Origin;
+import org.eclipse.jetty.client.WWWAuthenticationProtocolHandler;
 import org.eclipse.jetty.client.api.Authentication;
 import org.eclipse.jetty.client.api.Request;
 import org.eclipse.jetty.client.api.Response;
 import org.eclipse.jetty.client.api.Result;
+import org.eclipse.jetty.client.dynamic.HttpClientTransportDynamic;
 import org.eclipse.jetty.client.util.BufferingResponseListener;
 import org.eclipse.jetty.client.util.BytesRequestContent;
 import org.eclipse.jetty.http.HttpField;
@@ -19,7 +21,8 @@ import org.eclipse.jetty.http.HttpHeader;
 import org.eclipse.jetty.http.HttpMethod;
 import org.eclipse.jetty.http.HttpVersion;
 import org.eclipse.jetty.http2.client.HTTP2Client;
-import org.eclipse.jetty.http2.client.http.HttpClientTransportOverHTTP2;
+import org.eclipse.jetty.http2.client.http.ClientConnectionFactoryOverHTTP2;
+import org.eclipse.jetty.io.ClientConnectionFactory;
 import org.eclipse.jetty.io.ClientConnector;
 import org.eclipse.jetty.util.HttpCookieStore;
 import org.eclipse.jetty.util.Pool;
@@ -82,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) {
@@ -144,7 +145,8 @@ class JettyCluster implements Cluster {
         int initialWindow = Integer.MAX_VALUE;
         h2Client.setInitialSessionRecvWindow(initialWindow);
         h2Client.setInitialStreamRecvWindow(initialWindow);
-        HttpClientTransportOverHTTP2 transport = new HttpClientTransportOverHTTP2(h2Client);
+        ClientConnectionFactory.Info http2 = new ClientConnectionFactoryOverHTTP2.HTTP2(h2Client);
+        HttpClientTransportDynamic transport = new HttpClientTransportDynamic(connector, http2);
         transport.setConnectionPoolFactory(dest -> {
             MultiplexConnectionPool pool = new MultiplexConnectionPool(
                     dest, Pool.StrategyType.RANDOM, b.connectionsPerEndpoint, false, dest, Integer.MAX_VALUE);
@@ -161,6 +163,8 @@ class JettyCluster implements Cluster {
 
         if (b.proxy != null) addProxyConfiguration(b, httpClient);
         try { httpClient.start(); } catch (Exception e) { throw new IOException(e); }
+        // Must be removed after client has started
+        httpClient.getProtocolHandlers().remove(WWWAuthenticationProtocolHandler.NAME);
         return httpClient;
     }
 
@@ -168,9 +172,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)));
@@ -178,23 +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()) {
-            for (URI endpoint : b.endpoints) {
-                httpClient.getAuthenticationStore().addAuthenticationResult(new Authentication.Result() {
-                    @Override public URI getURI() { return URI.create(endpointUri(endpoint)); }
-                    @Override public void apply(Request r) {
-                        r.headers(hs -> proxyHeaders.forEach((k, v) -> hs.add(k, v.get())));
-                    }
-                });
-
-            }
+        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 -> 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();
diff --git a/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/SslContextBuilder.java b/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/SslContextBuilder.java
index 1855b657a75..85144ae3e8c 100644
--- a/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/SslContextBuilder.java
+++ b/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/SslContextBuilder.java
@@ -85,7 +85,8 @@ class SslContextBuilder {
             } else if (hasCaCertificateInstance()) {
                 addCaCertificates(keystore, caCertificates);
             }
-            SSLContext sslContext = SSLContext.getInstance("TLSv1.2"); // Protocol version must match TlsContext.SSL_CONTEXT_VERSION
+            // Protocol version must be equal to TlsContext.SSL_CONTEXT_VERSION or higher
+            SSLContext sslContext = SSLContext.getInstance("TLSv1.3");
             sslContext.init(
                     createKeyManagers(keystore).orElse(null),
                     createTrustManagers(keystore).orElse(null),
diff --git a/vespa-feed-client/src/test/java/ai/vespa/feed/client/impl/SslContextBuilderTest.java b/vespa-feed-client/src/test/java/ai/vespa/feed/client/impl/SslContextBuilderTest.java
index 95952d37c3c..bddb8857dc3 100644
--- a/vespa-feed-client/src/test/java/ai/vespa/feed/client/impl/SslContextBuilderTest.java
+++ b/vespa-feed-client/src/test/java/ai/vespa/feed/client/impl/SslContextBuilderTest.java
@@ -57,13 +57,13 @@ class SslContextBuilderTest {
                         .withCaCertificates(certificateFile)
                         .withCertificateAndKey(certificateFile, privateKeyFile)
                         .build());
-        assertEquals("TLSv1.2", sslContext.getProtocol());
+        assertEquals("TLSv1.3", sslContext.getProtocol());
     }
 
     @Test
     void successfully_constructs_sslcontext_when_no_builder_parameter_given() {
         SSLContext sslContext = Assertions.assertDoesNotThrow(() -> new SslContextBuilder().build());
-        assertEquals("TLSv1.2", sslContext.getProtocol());
+        assertEquals("TLSv1.3", sslContext.getProtocol());
     }
 
     @Test
@@ -72,7 +72,7 @@ class SslContextBuilderTest {
                 new SslContextBuilder()
                         .withCertificateAndKey(certificateFile, privateKeyFile)
                         .build());
-        assertEquals("TLSv1.2", sslContext.getProtocol());
+        assertEquals("TLSv1.3", sslContext.getProtocol());
     }
 
     @Test
@@ -81,7 +81,7 @@ class SslContextBuilderTest {
                 new SslContextBuilder()
                         .withCaCertificates(certificateFile)
                         .build());
-        assertEquals("TLSv1.2", sslContext.getProtocol());
+        assertEquals("TLSv1.3", sslContext.getProtocol());
     }
 
     private static void writePem(Path file, String type, byte[] asn1DerEncodedObject) throws IOException {