Added feature to add custom headers to Service

Author: BalasubramanyamEvani

splunk/src/main/java/com/splunk/HttpService.java

@@ -23,10 +23,16 @@
 import java.io.OutputStreamWriter;
 import java.net.*;
 import java.security.cert.X509Certificate;
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import java.util.Set;
 
 /**
  * The {@code HttpService} class represents a generic HTTP service at a given
@@ -79,11 +85,20 @@ public boolean verify(String s, SSLSession sslSession) {
 
     private String prefix = null;
 
+    private static final Set<String> filtertHttpHeaderKeys = Collections.unmodifiableSet(
+    		new HashSet<String>(Arrays.asList(
+    				"User-Agent",
+    				"Accept"
+    		))
+    );
+    
     static Map<String, String> defaultHeader = new HashMap<String, String>() {{
         put("User-Agent", "splunk-sdk-java/1.7.1");
         put("Accept", "*/*");
     }};
-
+    
+    protected Map<String, String> customHeaders = new HashMap<>();
+    
     protected SimpleCookieStore cookieStore = new SimpleCookieStore();
 
     /**
@@ -193,6 +208,21 @@ public String getHost() {
     public int getPort() {
         return this.port;
     }
+    
+    /**
+     * Sets Custom Headers of this service
+     * 
+     * @param headers
+     */
+    public void setCustomHeaders(Map<String, String> headers) {
+    	if (Objects.nonNull(headers) && !headers.isEmpty()) {
+    		Map<String, String> fitleredCustomHeaders = headers.entrySet()
+    				.stream()
+    				.filter(e -> !filtertHttpHeaderKeys.contains(e.getKey()))
+    				.collect(Collectors.toMap(map -> map.getKey(), map -> map.getValue()));
+    		customHeaders.putAll(fitleredCustomHeaders);
+    	}
+    }
 
     /**
      * Returns the SSL security protocol of this service.
@@ -258,6 +288,15 @@ public URL getUrl(String path) {
             throw new RuntimeException(e.getMessage(), e);
         }
     }
+    
+    /**
+     * Returns all the stored custom headers
+     * 
+     * @return customHeaders The custom headers
+     */
+    public Map<String, String> getCustomHeaders() {
+    	return customHeaders;
+    }
 
     /**
      * Returns all the stored cookies
@@ -441,6 +480,13 @@ public ResponseMessage send(String path, RequestMessage request) {
             if (header.containsKey(key)) continue;
             cn.setRequestProperty(key, entry.getValue());
         }
+        // Add Custom Headers
+        for (Entry<String, String> entry: customHeaders.entrySet()) {
+        	String key = entry.getKey();
+        	if (!header.containsKey(key)) {
+        		cn.setRequestProperty(key, entry.getValue());
+        	}
+        }
 
         // Add cookies to header
         cn.setRequestProperty("Cookie", cookieStore.getCookies());

splunk/src/main/java/com/splunk/Service.java

@@ -152,6 +152,7 @@ public Service(ServiceArgs args) {
         this.httpsHandler = Args.<URLStreamHandler>get(args, "httpsHandler", null);
         this.setSslSecurityProtocol(Args.get(args, "SSLSecurityProtocol", Service.getSslSecurityProtocol()));
         this.addCookie((String)args.get("cookie"));
+        this.setCustomHeaders((Map<String, String>) args.get("customHeaders"));
     }
 
     /**

splunk/src/main/java/com/splunk/ServiceArgs.java

@@ -17,6 +17,7 @@
 package com.splunk;
 
 import java.net.URLStreamHandler;
+import java.util.Map;
 
 /**
  * The {@code ServiceArgs} class contains a collection of arguments that are
@@ -164,4 +165,12 @@ public void setUsername(String username) {
     public void setCookie(String cookie) {
         this.put("cookie", cookie);
     }
+    
+    /**
+     * @param httpHeaders
+     *     A map of customHeaders.
+     */
+    public void setHttpHeaders(Map<String, String> httpHeaders) {
+    	this.put("customHeaders", httpHeaders);
+    }
 }

splunk/src/test/java/com/splunk/ServiceTest.java

@@ -146,7 +146,23 @@ public void testLogin() {
         service.logout();
         checkNotLoggedIn(service);
     }
-
+    
+    @Test
+    public void testServiceWithCustomHeaders() {
+    	ServiceArgs args = new ServiceArgs();
+    	args.setHost((String) command.opts.get("host"));
+    	args.setPort((Integer) command.opts.get("port"));
+    	args.setScheme((String) command.opts.get("scheme"));
+    	args.setUsername((String) command.opts.get("username"));
+    	args.setPassword((String) command.opts.get("password"));
+    	args.setHttpHeaders(new HashMap<String, String>() {{
+    		put("some header key", "some value");
+    	}});
+        Service service = new Service(args);
+        Map<String, String> customHeaders = service.getCustomHeaders();
+        Assert.assertEquals(customHeaders.get("some header key"), "some value");
+    }
+    
     @Test
     public void testLoginWithoutArguments() {
         ServiceArgs args = new ServiceArgs();