Add proxy option to vespa-feed-client

8 files changed, 46 insertions, 15 deletions

diff --git a/vespa-feed-client-api/abi-spec.json b/vespa-feed-client-api/abi-spec.json
index 16e532a2c9a..5bd0acf82d3 100644
--- a/vespa-feed-client-api/abi-spec.json
+++ b/vespa-feed-client-api/abi-spec.json
@@ -140,6 +140,7 @@
       "public abstract ai.vespa.feed.client.FeedClientBuilder setCaCertificatesFile(java.nio.file.Path)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setCaCertificates(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.FeedClient build()"
     ],
     "fields": [
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 95c9b2c95fe..7ec5fbb02b7 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
@@ -109,6 +109,9 @@ public interface FeedClientBuilder {
     /** Overrides endpoint URIs for this client */
     FeedClientBuilder setEndpointUris(List<URI> endpoints);
 
+    /** Specify HTTP(S) proxy for all endpoints */
+    FeedClientBuilder setProxy(URI uri);
+
     /** Constructs instance of {@link FeedClient} from builder configuration */
     FeedClient build();
 
diff --git a/vespa-feed-client-cli/src/main/java/ai/vespa/feed/client/impl/CliArguments.java b/vespa-feed-client-cli/src/main/java/ai/vespa/feed/client/impl/CliArguments.java
index ac859359bfa..e024f961e26 100644
--- a/vespa-feed-client-cli/src/main/java/ai/vespa/feed/client/impl/CliArguments.java
+++ b/vespa-feed-client-cli/src/main/java/ai/vespa/feed/client/impl/CliArguments.java
@@ -55,6 +55,7 @@ class CliArguments {
     private static final String VERSION_OPTION = "version";
     private static final String STDIN_OPTION = "stdin";
     private static final String DOOM_OPTION = "max-failure-seconds";
+    private static final String PROXY_OPTION = "proxy";
 
     private final CommandLine arguments;
 
@@ -165,6 +166,16 @@ class CliArguments {
 
     boolean dryrunEnabled() { return has(DRYRUN_OPTION); }
 
+    Optional<URI> proxy() throws CliArgumentsException {
+        try {
+            URL url = (URL) arguments.getParsedOptionValue(PROXY_OPTION);
+            if (url == null) return Optional.empty();
+            return Optional.of(url.toURI());
+        } catch (ParseException | URISyntaxException e) {
+            throw new CliArgumentsException("Invalid proxy: " + e.getMessage(), e);
+        }
+    }
+
     private OptionalInt intValue(String option) throws CliArgumentsException {
         try {
             Number number = (Number) arguments.getParsedOptionValue(option);
@@ -310,6 +321,12 @@ class CliArguments {
                 .addOption(Option.builder()
                         .longOpt(SHOW_ALL_OPTION)
                         .desc("Print the result of every feed operation")
+                        .build())
+                .addOption(Option.builder()
+                        .longOpt(PROXY_OPTION)
+                        .desc("URI to proxy endpoint")
+                        .hasArg()
+                        .type(URL.class)
                         .build());
     }
 
diff --git a/vespa-feed-client-cli/src/main/java/ai/vespa/feed/client/impl/CliClient.java b/vespa-feed-client-cli/src/main/java/ai/vespa/feed/client/impl/CliClient.java
index 5e904b37588..68b9cf6af0e 100644
--- a/vespa-feed-client-cli/src/main/java/ai/vespa/feed/client/impl/CliClient.java
+++ b/vespa-feed-client-cli/src/main/java/ai/vespa/feed/client/impl/CliClient.java
@@ -138,6 +138,7 @@ public class CliClient {
         builder.setDryrun(cliArgs.dryrunEnabled());
         cliArgs.doomSeconds().ifPresent(doom -> builder.setCircuitBreaker(new GracePeriodCircuitBreaker(Duration.ofSeconds(10),
                                                                                                         Duration.ofSeconds(doom))));
+        cliArgs.proxy().ifPresent(builder::setProxy);
         return builder.build();
     }
 
diff --git a/vespa-feed-client-cli/src/test/java/ai/vespa/feed/client/impl/CliArgumentsTest.java b/vespa-feed-client-cli/src/test/java/ai/vespa/feed/client/impl/CliArgumentsTest.java
index fe3dc465814..201a85ed09d 100644
--- a/vespa-feed-client-cli/src/test/java/ai/vespa/feed/client/impl/CliArgumentsTest.java
+++ b/vespa-feed-client-cli/src/test/java/ai/vespa/feed/client/impl/CliArgumentsTest.java
@@ -1,7 +1,6 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.feed.client.impl;
 
-import ai.vespa.feed.client.impl.CliArguments;
 import org.junit.jupiter.api.Test;
 
 import java.io.ByteArrayOutputStream;
@@ -11,7 +10,10 @@ import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.time.Duration;
 
-import static org.junit.jupiter.api.Assertions.*;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 /**
  * @author bjorncs
@@ -26,7 +28,7 @@ class CliArgumentsTest {
                 "--ca-certificates=ca-certs.pem", "--disable-ssl-hostname-verification",
                 "--header=\"My-Header: my-value\"", "--header", "Another-Header: another-value", "--benchmark",
                 "--route=myroute", "--timeout=0.125", "--trace=9", "--verbose", "--silent",
-                "--show-errors", "--show-all", "--max-failure-seconds=30"});
+                "--show-errors", "--show-all", "--max-failure-seconds=30", "--proxy", "https://myproxy:1234"});
         assertEquals(URI.create("https://vespa.ai:4443/"), args.endpoint());
         assertEquals(Paths.get("feed.json"), args.inputFile().get());
         assertEquals(10, args.connections().getAsInt());
@@ -49,6 +51,7 @@ class CliArgumentsTest {
         assertTrue(args.showErrors());
         assertTrue(args.showSuccesses());
         assertFalse(args.showProgress());
+        assertEquals(URI.create("https://myproxy:1234"), args.proxy().orElse(null));
     }
 
     @Test
diff --git a/vespa-feed-client-cli/src/test/resources/help.txt b/vespa-feed-client-cli/src/test/resources/help.txt
index 323206ab128..66d7c3521c2 100644
--- a/vespa-feed-client-cli/src/test/resources/help.txt
+++ b/vespa-feed-client-cli/src/test/resources/help.txt
@@ -25,6 +25,7 @@ Vespa feed client
                                           streams per HTTP/2 connection
     --private-key <arg>                   Path to PEM/PKCS#8 encoded
                                           private key file
+    --proxy <arg>                         URI to proxy endpoint
     --route <arg>                         Target Vespa route for feed
                                           operations
     --show-all                            Print the result of every feed
diff --git a/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/ApacheCluster.java b/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/ApacheCluster.java
index b51210d22ea..62cd56f21ce 100644
--- a/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/ApacheCluster.java
+++ b/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/ApacheCluster.java
@@ -10,15 +10,13 @@ import org.apache.hc.client5.http.impl.async.HttpAsyncClients;
 import org.apache.hc.client5.http.ssl.ClientTlsStrategyBuilder;
 import org.apache.hc.core5.concurrent.FutureCallback;
 import org.apache.hc.core5.http.ContentType;
+import org.apache.hc.core5.http.HttpHost;
 import org.apache.hc.core5.http.message.BasicHeader;
 import org.apache.hc.core5.http2.config.H2Config;
 import org.apache.hc.core5.net.URIAuthority;
 import org.apache.hc.core5.reactor.IOReactorConfig;
-import org.apache.hc.core5.util.Timeout;
-
-import org.apache.hc.core5.function.Factory;
 import org.apache.hc.core5.reactor.ssl.TlsDetails;
-import javax.net.ssl.SSLEngine;
+import org.apache.hc.core5.util.Timeout;
 
 import javax.net.ssl.SSLContext;
 import java.io.IOException;
@@ -45,11 +43,7 @@ class ApacheCluster implements Cluster {
     private final List<Endpoint> endpoints = new ArrayList<>();
     private final List<BasicHeader> defaultHeaders = Arrays.asList(new BasicHeader("User-Agent", String.format("vespa-feed-client/%s", Vespa.VERSION)),
                                                                    new BasicHeader("Vespa-Client-Version", Vespa.VERSION));
-    private final RequestConfig defaultConfig = RequestConfig.custom()
-                                                             .setConnectTimeout(Timeout.ofSeconds(10))
-                                                             .setConnectionRequestTimeout(Timeout.DISABLED)
-                                                             .setResponseTimeout(Timeout.ofSeconds(190))
-                                                             .build();
+    private final RequestConfig requestConfig;
 
     private final ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor(t -> new Thread(t, "request-timeout-thread"));
 
@@ -57,6 +51,7 @@ class ApacheCluster implements Cluster {
         for (URI endpoint : builder.endpoints)
             for (int i = 0; i < builder.connectionsPerEndpoint; i++)
                 endpoints.add(new Endpoint(createHttpClient(builder), endpoint));
+        this.requestConfig = createRequestConfig(builder);
     }
 
     @Override
@@ -75,7 +70,7 @@ class ApacheCluster implements Cluster {
             SimpleHttpRequest request = new SimpleHttpRequest(wrapped.method(), wrapped.path());
             request.setScheme(endpoint.url.getScheme());
             request.setAuthority(new URIAuthority(endpoint.url.getHost(), portOf(endpoint.url)));
-            request.setConfig(defaultConfig);
+            request.setConfig(requestConfig);
             defaultHeaders.forEach(request::setHeader);
             wrapped.headers().forEach((name, value) -> request.setHeader(name, value.get()));
             if (wrapped.body() != null)
@@ -162,6 +157,15 @@ class ApacheCluster implements Cluster {
                                    : url.getPort();
     }
 
+    private static RequestConfig createRequestConfig(FeedClientBuilderImpl b) {
+        RequestConfig.Builder builder = RequestConfig.custom()
+                .setConnectTimeout(Timeout.ofSeconds(10))
+                .setConnectionRequestTimeout(Timeout.DISABLED)
+                .setResponseTimeout(Timeout.ofSeconds(190));
+        if (b.proxy != null) builder.setProxy(new HttpHost(b.proxy.getScheme(), b.proxy.getHost(), b.proxy.getPort()));
+        return builder.build();
+    }
+
     private static class ApacheHttpResponse implements HttpResponse {
 
         private final SimpleHttpResponse wrapped;
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 7dafeb0b541..134ad464618 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
@@ -19,7 +19,6 @@ import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.function.Supplier;
 
 import static java.util.Objects.requireNonNull;
@@ -50,7 +49,7 @@ public class FeedClientBuilderImpl implements FeedClientBuilder {
     Collection<X509Certificate> caCertificates;
     boolean benchmark = true;
     boolean dryrun = false;
-
+    URI proxy;
 
 
     public FeedClientBuilderImpl() {
@@ -188,6 +187,8 @@ public class FeedClientBuilderImpl implements FeedClientBuilder {
         return this;
     }
 
+    @Override public FeedClientBuilder setProxy(URI uri) { this.proxy = uri; return this; }
+
     /** Constructs instance of {@link ai.vespa.feed.client.FeedClient} from builder configuration */
     @Override
     public FeedClient build() {