From 216f45d5debed006c9e69d6e979cc21c7392fbd5 Mon Sep 17 00:00:00 2001
From: Bjørn Christian Seime <bjorncs@verizonmedia.com>
Date: Mon, 28 Mar 2022 14:37:15 +0200
Subject: Add proxy option to vespa-feed-client

---
 .../ai/vespa/feed/client/impl/ApacheCluster.java   | 24 +++++++++++++---------
 .../feed/client/impl/FeedClientBuilderImpl.java    |  5 +++--
 2 files changed, 17 insertions(+), 12 deletions(-)

(limited to 'vespa-feed-client')

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() {
-- 
cgit v1.2.3