Allow setting http headers through command line interface

2 files changed, 45 insertions, 1 deletions

diff --git a/vespa-http-client/src/main/java/com/yahoo/vespa/http/client/runner/CommandLineArguments.java b/vespa-http-client/src/main/java/com/yahoo/vespa/http/client/runner/CommandLineArguments.java
index 671038c852a..a159d277e1a 100644
--- a/vespa-http-client/src/main/java/com/yahoo/vespa/http/client/runner/CommandLineArguments.java
+++ b/vespa-http-client/src/main/java/com/yahoo/vespa/http/client/runner/CommandLineArguments.java
@@ -15,6 +15,8 @@ import org.apache.http.conn.ssl.NoopHostnameVerifier;
 import org.apache.http.conn.ssl.SSLConnectionSocketFactory;
 
 import javax.inject.Inject;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -180,6 +182,10 @@ public class CommandLineArguments {
             description = "Skip hostname verification when using TLS")
     private boolean insecure = false;
 
+    @Option(name = {"--header"},
+            description = "Add http header to every request. Header must have the format '<Name>: <Value>'. Use this parameter multiple times for multiple headers")
+    private List<String> headers = new ArrayList<>();
+
     int getWhenVerboseEnabledPrintMessageForEveryXDocuments() {
         return whenVerboseEnabledPrintMessageForEveryXDocuments;
     }
@@ -192,6 +198,11 @@ public class CommandLineArguments {
 
     SessionParams createSessionParams(boolean useJson) {
         final int minThrottleValue = useDynamicThrottlingArg ? 10 : 0;
+        ConnectionParams.Builder connectionParamsBuilder = new ConnectionParams.Builder();
+        for (String header : headers) {
+            String[] nameAndValue = header.split(":");
+            connectionParamsBuilder.addHeader(nameAndValue[0].trim(), nameAndValue[1].trim());
+        }
         SessionParams.Builder builder = new SessionParams.Builder()
                 .setFeedParams(
                         new FeedParams.Builder()
@@ -208,7 +219,7 @@ public class CommandLineArguments {
                                 .build()
                 )
                 .setConnectionParams(
-                        new ConnectionParams.Builder()
+                        connectionParamsBuilder
                                 .setHostnameVerifier(insecure ? NoopHostnameVerifier.INSTANCE :
                                         SSLConnectionSocketFactory.getDefaultHostnameVerifier())
                                 .setNumPersistentConnectionsPerEndpoint(16)
diff --git a/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/runner/CommandLineArgumentsTest.java b/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/runner/CommandLineArgumentsTest.java
index e0d93a7fa18..84a69520a84 100644
--- a/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/runner/CommandLineArgumentsTest.java
+++ b/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/runner/CommandLineArgumentsTest.java
@@ -7,7 +7,13 @@ import com.yahoo.vespa.http.client.config.SessionParams;
 import org.junit.Test;
 
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
@@ -109,6 +115,7 @@ public class CommandLineArgumentsTest {
         add("debugport", "7890");
         args.add("--verbose");
         args.add("--useTls");
+        add("header", "Header-Name: Header-Value");
         CommandLineArguments arguments = CommandLineArguments.build(asArray());
         SessionParams params = arguments.createSessionParams(true /* use json */);
         assertThat(params.getClientQueueSize(), is(3456));
@@ -116,6 +123,7 @@ public class CommandLineArgumentsTest {
         assertThat(params.getClusters().get(0).getEndpoints().get(0).getPort(), is(1234));
         assertThat(params.getClusters().get(0).getEndpoints().get(0).isUseSsl(), is(true));
         assertThat(params.getConnectionParams().getUseCompression(), is(true));
+        assertThat(params.getConnectionParams().getHeaders().size(), is(1));
         assertThat(params.getFeedParams().getRoute(), is("routeValue"));
         assertThat(params.getFeedParams().getDataFormat(), is(FeedParams.DataFormat.JSON_UTF8));
         assertThat(params.getFeedParams().getLocalQueueTimeOut(), is(2345000L));
@@ -124,6 +132,31 @@ public class CommandLineArgumentsTest {
     }
 
     @Test
+    public void testAddingMultipleHttpHeaders() {
+        add("host", "hostValue");
+        String header1Name = "Header-Name-1";
+        String header1Value = "Header-Value";
+        add("header", header1Name + ": " + header1Value);
+        String header2Name = "Header-Name-2";
+        String header2Value = "Another-Header-Value";
+        add("header", header2Name + ": " + header2Value);
+
+        CommandLineArguments arguments = CommandLineArguments.build(asArray());
+        SessionParams params = arguments.createSessionParams(true /* use json */);
+
+        List<Map.Entry<String, String>> headers = new ArrayList<>(params.getConnectionParams().getHeaders());
+        headers.sort(Comparator.comparing(Map.Entry::getKey));
+
+        assertThat(headers.size(), is(2));
+        Map.Entry<String, String> actualHeader1 = headers.get(0);
+        assertThat(actualHeader1.getKey(), is(header1Name));
+        assertThat(actualHeader1.getValue(), is(header1Value));
+        Map.Entry<String, String> actualHeader2 = headers.get(1);
+        assertThat(actualHeader2.getKey(), is(header2Name));
+        assertThat(actualHeader2.getValue(), is(header2Value));
+    }
+
+    @Test
     public void testMultiHost() {
         add("file", "fileValue.json");
         add("port", "1234");