Merge branch 'persist-cookies' into revert-v2-changes

Author: ashah-splunk

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

@@ -316,12 +316,12 @@ public void removeAllCookies() {
     }
 
     /**
-     * Returns true if the cookeStore has any cookies, false otherwise
+     * Returns true if the cookieStore has any Splunk Authorization cookies, false otherwise
      *
      * @return True if there are cookies, false otherwise
      */
-    public Boolean hasCookies() {
-        return !cookieStore.isEmpty();
+    public Boolean hasSplunkAuthCookies() {
+        return cookieStore.hasSplunkAuthCookie();
     }
 
     /**

splunk/src/main/java/com/splunk/Receiver.java

@@ -98,9 +98,13 @@ public Socket attach(String indexName, Args args) throws IOException {
         headers.add("Accept-Encoding: identity");
         headers.add("X-Splunk-Input-Mode: Streaming");
 
-        if (service.hasCookies()) {
+        if (service.hasSplunkAuthCookies()) {
             headers.add(String.format("Cookie: %s", service.stringifyCookies()));
         } else {
+            // to persist the cookies other than Splunk such as from Load Balancer
+            if(!service.cookieStore.isEmpty()){
+                headers.add(String.format("Cookie: %s", service.stringifyCookies()));
+            }
             headers.add(String.format("Authorization: %s", service.getToken()));
         }
         headers.add("");

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

@@ -1112,7 +1112,7 @@ public UserCollection getUsers(Args args) {
      * @return The current {@code Service} instance.
      */
     public Service login() {
-        if (!this.cookieStore.isEmpty() && (this.username == null || this.password == null)) {
+        if (this.cookieStore.hasSplunkAuthCookie() && (this.username == null || this.password == null)) {
             return this;
         }
         else if (this.username == null || this.password == null) {
@@ -1312,7 +1312,7 @@ public Job search(String query, Map<String, Object> args) {
      */
     @Override public ResponseMessage send(String path, RequestMessage request) {
         // cookieStore is a protected member of HttpService
-        if (token != null && cookieStore.isEmpty()) {
+        if (token != null && !cookieStore.hasSplunkAuthCookie() ) {
             request.getHeader().put("Authorization", token);
         }
         return super.send(fullpath(path), request);

splunk/src/main/java/com/splunk/SimpleCookieStore.java

@@ -28,6 +28,8 @@
  */
 class SimpleCookieStore {
 
+    public static final String SPLUNK_AUTH_COOKIE = "splunkd_";
+
     private Map<String, String> cookieJar = new HashMap<String, String>();
     /**
      * Adds cookies from a "Set-Cookie" header to the cookie store.
@@ -69,6 +71,18 @@ public Boolean isEmpty() {
         return cookieJar.isEmpty();
     }
 
+    public boolean hasSplunkAuthCookie(){
+        if(cookieJar.isEmpty()){
+            return false;
+        }
+        for(String cookie : cookieJar.keySet()){
+            if(cookie.startsWith(SPLUNK_AUTH_COOKIE)){
+                return true;
+            }
+        }
+        return false;
+    }
+
     /**
      * Removes all cookies from SimpleCookieStore
      */

splunk/src/test/java/com/splunk/CookieTest.java

@@ -16,13 +16,11 @@
 
 package com.splunk;
 
+import java.net.HttpCookie;
 import java.util.HashMap;
 import java.util.Map;
 
-import org.junit.Assert;
-import org.junit.Assume;
-import org.junit.Before;
-import org.junit.Test;
+import org.junit.*;
 
 public class CookieTest extends SDKTestCase {
 
@@ -124,6 +122,32 @@ public void testLoginWithMultipleCookies() {
         s.getSettings().refresh();
     }
 
+    @Test
+    public void testLoginWithOtherCookies() {
+        String otherCookies = "load=balancer;";
+        service.logout();
+        service.cookieStore.removeAll();
+        service.cookieStore.add(otherCookies);
+        service.login();
+        service.getApplications();
+        service.cookieStore.removeAll();
+    }
+
+    @Test
+    public void testUsingAuthTokenAndOtherCookie(){
+        String validToken = service.getToken();
+        Assert.assertTrue(validToken.startsWith("Splunk "));
+        String otherCookies = "load=balancer;";
+        Map<String, Object> args = new HashMap<>();
+        args.put("cookie", otherCookies);
+        args.put("host",service.getHost());
+        args.put("port", service.getPort());
+        Service s = new Service(args);
+        s.setToken(validToken);
+        s.getApplications();
+        Assert.assertEquals(otherCookies.trim(),s.cookieStore.getCookies().trim());
+    }
+
     @Test
     public void testLoginWithMultipleInvalidCookies() {
         String validCookie = service.stringifyCookies();

splunk/src/test/java/com/splunk/IndexTest.java

@@ -77,8 +77,8 @@ public void testAttachWithCookieHeader() throws IOException {
             // Cookies not implemented before version 6.2
             return;
         }
-        // Check that their are cookies at all
-        Assert.assertTrue(service.hasCookies());
+        // Check that their are Splunk Auth cookies at all
+        Assert.assertTrue(service.hasSplunkAuthCookies());
 
         // Make a service that only has that cookie
         String validCookie = service.stringifyCookies();

splunk/src/test/java/com/splunk/ReceiverTest.java

@@ -24,6 +24,8 @@
 import java.io.IOException;
 import java.io.OutputStream;
 import java.net.Socket;
+import java.util.HashMap;
+import java.util.Map;
 
 public class ReceiverTest extends SDKTestCase {
 
@@ -36,9 +38,28 @@ public void testReceiverWithoutCookie() {
     @Test
     public void testReceiverWithCookie() {
         Assume.assumeTrue(service.versionIsAtLeast("6.2"));
-        Assert.assertTrue(service.hasCookies());
+        Assert.assertTrue(service.hasSplunkAuthCookies());
         testReceiver(service);
     }
+
+    @Test
+    public void testReceiverWithoutSplunkCookie() {
+        String validToken = service.getToken();
+        Assert.assertTrue(validToken.startsWith("Splunk "));
+        String otherCookies = "load=balancer;";
+        Map<String, Object> args = new HashMap<>();
+        args.put("cookie", otherCookies);
+        args.put("host",service.getHost());
+        args.put("port", service.getPort());
+        Service s = new Service(args);
+        s.setToken(validToken);
+        s.version = s.getInfo().getVersion();
+        System.out.println(s.version);
+        Assume.assumeTrue(s.versionIsAtLeast("6.2"));
+        Assert.assertTrue(!s.cookieStore.isEmpty());
+        testReceiver(s);
+    }
+
     // Make a few simple requests and make sure the results look ok.
     public void testReceiver(Service passedService) {
         Receiver receiver = passedService.getReceiver();