Add support for additional headers to DiagnosticService

src/main/java/co/elastic/support/BaseConfig.java

@@ -17,6 +17,7 @@ public class BaseConfig {
     public int socketTimeout;
     public int maxTotalConn;
     public int maxConnPerRoute;
+    public Map<String, String> extraHeaders;
 
     public String diagReleaseHost = "api.github.com";
     public String diagReleaseDest = "/repos/elastic/support-diagnostics/releases/latest";
@@ -64,6 +65,8 @@ public BaseConfig(Map configuration) {
         maxTotalConn = restConfig.get("maxTotalConn");
         maxConnPerRoute = restConfig.get("maxConnPerRoute");
 
+        extraHeaders = (Map<String, String>) configuration.get("extra-headers");
+
         dockerGlobal = (Map<String, String>) configuration.get("docker-global");
         //kubernates = (Map<String, String>) configuration.get("kuberantes");
 

src/main/java/co/elastic/support/BaseInputs.java

@@ -121,7 +121,7 @@ protected void runOutputDirInteractive(){
                 .withIgnoreCase()
                 .withInputTrimming(true)
                 .withValueChecker(( String val, String propname) -> validateArchiveType(val))
-                .read(SystemProperties.lineSeparator + outputDirDescription);
+                .read(SystemProperties.lineSeparator + archiveTypeDescription);
     }
 
     public List<String> validatePort(int val){

src/main/java/co/elastic/support/diagnostics/DiagnosticService.java

@@ -47,6 +47,7 @@ public File exec(DiagnosticInputs inputs, DiagConfig config) throws DiagnosticEx
                     inputs.pkiKeystore,
                     inputs.pkiKeystorePass,
                     inputs.skipVerification,
+                    config.extraHeaders,
                     config.connectionTimeout,
                     config.connectionRequestTimeout,
                     config.socketTimeout

src/main/java/co/elastic/support/diagnostics/commands/CheckDiagnosticVersion.java

@@ -55,6 +55,7 @@ public void execute(DiagnosticContext context) {
                     "",
                     "",
                     true,
+                    context.diagsConfig.extraHeaders,
                     context.diagsConfig.connectionTimeout,
                     context.diagsConfig.connectionRequestTimeout,
                     context.diagsConfig.socketTimeout)){

src/main/java/co/elastic/support/diagnostics/commands/CheckElasticsearchVersion.java

@@ -56,6 +56,7 @@ public void execute(DiagnosticContext context) throws DiagnosticException {
                     context.diagnosticInputs.pkiKeystore,
                     context.diagnosticInputs.pkiKeystorePass,
                     context.diagnosticInputs.skipVerification,
+                    context.diagsConfig.extraHeaders,
                     context.diagsConfig.connectionTimeout,
                     context.diagsConfig.connectionRequestTimeout,
                     context.diagsConfig.socketTimeout);

src/main/java/co/elastic/support/diagnostics/commands/CheckKibanaVersion.java

@@ -60,6 +60,7 @@ public void execute(DiagnosticContext context) throws DiagnosticException {
                     context.diagnosticInputs.pkiKeystore,
                     context.diagnosticInputs.pkiKeystorePass,
                     context.diagnosticInputs.skipVerification,
+                    context.diagsConfig.extraHeaders,
                     context.diagsConfig.connectionTimeout,
                     context.diagsConfig.connectionRequestTimeout,
                     context.diagsConfig.socketTimeout);

src/main/java/co/elastic/support/monitoring/MonitoringExportService.java

@@ -75,6 +75,7 @@ public void execExtract(MonitoringExportInputs inputs) throws DiagnosticExceptio
                     inputs.pkiKeystore,
                     inputs.pkiKeystorePass,
                     inputs.skipVerification,
+                    config.extraHeaders,
                     config.connectionTimeout,
                     config.connectionRequestTimeout,
                     config.socketTimeout

src/main/java/co/elastic/support/monitoring/MonitoringImportService.java

@@ -88,6 +88,7 @@ private RestClient getClient(MonitoringImportInputs inputs, MonitoringImportConf
                 inputs.pkiKeystore,
                 inputs.pkiKeystorePass,
                 inputs.skipVerification,
+                config.extraHeaders,
                 config.connectionTimeout,
                 config.connectionRequestTimeout,
                 config.socketTimeout

src/main/java/co/elastic/support/rest/RestClient.java

@@ -40,6 +40,7 @@
 import java.io.FileInputStream;
 import java.io.UnsupportedEncodingException;
 import java.security.KeyStore;
+import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
@@ -48,15 +49,17 @@ public class RestClient implements Closeable {
     private static final Logger logger = LogManager.getLogger(RestClient.class);
     private static final int maxTotal = 100, defaultMaxPerRoute = 10;
 
-    private static ConcurrentMap<String, RestClient> cachedClients = new ConcurrentHashMap<>();
     private CloseableHttpClient client;
     private HttpHost httpHost;
     private HttpClientContext httpContext;
 
-    public RestClient(CloseableHttpClient client, HttpHost httpHost, HttpClientContext context) {
+    private Map<String, String> extraHeaders;
+
+    public RestClient(CloseableHttpClient client, HttpHost httpHost, HttpClientContext context, Map<String, String> extraHeaders) {
         this.client = client;
         this.httpHost = httpHost;
         this.httpContext = context;
+        this.extraHeaders = extraHeaders;
     }
 
     public RestResult execQuery(String url) {
@@ -74,6 +77,11 @@ public HttpResponse execGet(String query) {
     }
 
     private HttpResponse execRequest(HttpRequestBase httpRequest) {
+        if (extraHeaders != null) {
+            for (Map.Entry<String, String> entry : extraHeaders.entrySet()) {
+                httpRequest.addHeader(entry.getKey(), entry.getValue());
+            }
+        }
         try {
             return client.execute(httpHost, httpRequest, httpContext);
         } catch (HttpHostConnectException e) {
@@ -130,6 +138,7 @@ public static RestClient getClient(
             String pkiKeystore,
             String pkiKeystorePass,
             boolean bypassVerify,
+            Map<String, String> extraHeaders,
             int connectionTimeout,
             int connectionRequestTimeout,
             int socketTimeout){
@@ -216,7 +225,7 @@ else if (StringUtils.isNotEmpty(proxyUser)){
             clientBuilder.setConnectionManager(mgr);
 
             CloseableHttpClient httpClient = clientBuilder.build();
-            RestClient restClient = new RestClient(httpClient, httpHost, context);
+            RestClient restClient = new RestClient(httpClient, httpHost, context, extraHeaders);
 
             return restClient;
         }

src/test/java/co/elastic/support/diagnostics/TestDiagnosticService.java

@@ -0,0 +1,146 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package co.elastic.support.diagnostics;
+
+import co.elastic.support.Constants;
+import co.elastic.support.util.JsonYamlUtils;
+import co.elastic.support.util.ResourceCache;
+import com.google.common.io.Files;
+import org.junit.jupiter.api.*;
+import org.mockserver.configuration.ConfigurationProperties;
+import org.mockserver.integration.ClientAndServer;
+import org.mockserver.model.Header;
+import org.mockserver.model.HttpRequest;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.*;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipFile;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockserver.integration.ClientAndServer.startClientAndServer;
+import static org.mockserver.model.HttpRequest.request;
+import static org.mockserver.model.HttpResponse.response;
+
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+class TestDiagnosticService {
+    private ClientAndServer mockServer;
+
+    static private String headerKey1 = "k1";
+    static private String headerVal1 = "v1";
+    static private String headerKey2 = "k2";
+    static private String headerVal2 = "v2";
+
+    @BeforeAll
+    public void globalSetup() {
+        mockServer = startClientAndServer(9880);
+        // mockserver by default is in verbose mode (useful when creating new test), move it to warning.
+        ConfigurationProperties.disableSystemOut(true);
+        ConfigurationProperties.logLevel("WARN");
+        ResourceCache.terminal.dispose();
+    }
+
+    private DiagConfig newDiagConfig() {
+        Map diagMap = Collections.emptyMap();
+        try {
+            diagMap = JsonYamlUtils.readYamlFromClasspath(Constants.DIAG_CONFIG, true);
+        } catch (DiagnosticException e) {
+            fail(e);
+        }
+        return new DiagConfig(diagMap);
+    }
+
+    private DiagnosticInputs newDiagnosticInputs() {
+        DiagnosticInputs diagnosticInputs = new DiagnosticInputs();
+        diagnosticInputs.port = 9880;
+        diagnosticInputs.outputDir = Files.createTempDir().toString();
+        return diagnosticInputs;
+    }
+
+    private HttpRequest myRequest(Boolean withHeaders) {
+        if (withHeaders) {
+            return request().withHeaders(
+                    new Header(headerKey1, headerVal1),
+                    new Header(headerKey2, headerVal2)
+            );
+        } else {
+            return request();
+        }
+    }
+
+    private void setupResponse(Boolean withHeaders) {
+        mockServer
+                .when(
+                        myRequest(withHeaders)
+                                .withPath("/")
+                )
+                .respond(
+                        response()
+                                .withBody("{\"version\": {\"number\": \"7.14.0\"}}")
+                );
+        mockServer
+                .when(
+                        myRequest(withHeaders)
+                )
+                .respond(
+                        response()
+                                .withBody("some_response_body")
+                );
+    }
+
+    @AfterAll
+    public void globalTeardown() {
+        mockServer.stop();
+    }
+
+    @Test
+    public void testWithExtraHeaders() {
+        setupResponse(true);
+
+        Map extraHeaders = new HashMap<String, String>();
+        extraHeaders.put(headerKey1, headerVal1);
+        extraHeaders.put(headerKey2, headerVal2);
+        DiagConfig diagConfig = newDiagConfig();
+        diagConfig.extraHeaders = extraHeaders;
+        DiagnosticService diag = new DiagnosticService();
+
+        try {
+            File result = diag.exec(newDiagnosticInputs(), diagConfig);
+            assertTrue(result.toString().matches(".*\\.zip$"), result.toString());
+            try {
+                ZipFile zipFile = new ZipFile(result, ZipFile.OPEN_READ);
+                Set<String> filenames = new HashSet<>();
+
+                Enumeration<? extends ZipEntry> entries = zipFile.entries();
+
+                while(entries.hasMoreElements()){
+                    ZipEntry entry = entries.nextElement();
+                    // Add file path without leading directory
+                    filenames.add(entry.getName().replaceFirst("/[^/]*/", ""));
+                }
+                assertTrue(filenames.contains("manifest.json"));
+            } catch (IOException e) {
+                fail("Error processing result zip file", e);
+            }
+        } catch (DiagnosticException e) {
+            fail(e);
+        }
+    }
+
+    @Test
+    public void testWithoutExtraHeaders() {
+        setupResponse(false);
+
+        DiagnosticService diag = new DiagnosticService();
+
+        try {
+            File result = diag.exec(newDiagnosticInputs(), newDiagConfig());
+        } catch (DiagnosticException e) {
+            fail(e);
+        }
+    }
+}

src/test/java/co/elastic/support/diagnostics/commands/TestCheckKibanaVersion.java

@@ -14,6 +14,8 @@
 import co.elastic.support.rest.RestClient;
 import com.vdurmont.semver4j.Semver;
 
+import java.util.Collections;
+
 import static org.junit.jupiter.api.Assertions.*;
 import static org.mockserver.integration.ClientAndServer.startClientAndServer;
 import static org.mockserver.model.HttpRequest.request;
@@ -52,6 +54,7 @@ public void setup() {
             "",
             "",
             true,
+            Collections.emptyMap(),
            3000,
            3000,
            3000);

src/test/java/co/elastic/support/diagnostics/commands/TestKibanaGetDetails.java

@@ -16,6 +16,7 @@
 import org.mockserver.integration.ClientAndServer;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Map;
 import org.junit.jupiter.api.*;
 import co.elastic.support.rest.RestClient;
@@ -61,6 +62,7 @@ public void setup() {
             "",
             "",
             true,
+            Collections.emptyMap(),
            3000,
            3000,
            3000);

src/test/java/co/elastic/support/diagnostics/commands/TestRunKibanaQueries.java

@@ -19,6 +19,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import java.io.File;
+import java.util.Collections;
 import java.util.Map;
 import org.junit.jupiter.api.*;
 import co.elastic.support.rest.RestClient;
@@ -67,6 +68,7 @@ public void setup() {
             "",
             "",
             true,
+            Collections.emptyMap(),
            3000,
            3000,
            3000);

src/test/java/co/elastic/support/rest/TestRestExecCalls.java

@@ -17,9 +17,11 @@
 import org.mockserver.model.Parameter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+
 import java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -35,8 +37,6 @@ public class TestRestExecCalls {
 
     private final Logger logger = LoggerFactory.getLogger(TestRestExecCalls.class);
     private ClientAndServer mockServer;
-    private Map config;
-    private PoolingHttpClientConnectionManager connectionManager;
     private RestClient httpRestClient, httpsRestClient;
     private String temp = SystemProperties.userDir + SystemProperties.fileSeparator + "temp";
     private File tempDir = new File(temp);
@@ -52,7 +52,7 @@ public void globalSetup() {
     }
 
     @AfterAll
-    public void globalTeardoown() {
+    public void globalTeardown() {
         mockServer.stop();
     }
 
@@ -73,6 +73,7 @@ public void setup() {
             "",
             "",
             true,
+            Collections.emptyMap(),
            3000,
            3000,
            3000);
@@ -89,6 +90,7 @@ public void setup() {
             "",
             "",
             true,
+            Collections.emptyMap(),
            3000,
            3000,
            3000);