From 4a7c1c286fc33203ee3ac858be1f480ec03ad6b5 Mon Sep 17 00:00:00 2001
From: Vihas Splunk <vhiteshbhai@splunk.com>
Date: Tue, 21 Nov 2023 15:55:43 +0530
Subject: [PATCH] fix: fix bug while dealing with raw events, increase coverage

---
 .../com/splunk/hecclient/HecURIBuilder.java   | 17 ++++++--
 .../splunk/hecclient/HecURIBuilderTest.java   | 41 +++++++++++++++++--
 2 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/src/main/java/com/splunk/hecclient/HecURIBuilder.java b/src/main/java/com/splunk/hecclient/HecURIBuilder.java
index d9c60f50..cb9e7e50 100644
--- a/src/main/java/com/splunk/hecclient/HecURIBuilder.java
+++ b/src/main/java/com/splunk/hecclient/HecURIBuilder.java
@@ -1,6 +1,8 @@
 package com.splunk.hecclient;
 
+import org.apache.http.Consts;
 import org.apache.http.client.utils.URIBuilder;
+import org.apache.http.client.utils.URLEncodedUtils;
 
 import java.net.URI;
 import java.net.URISyntaxException;
@@ -18,9 +20,18 @@ public HecURIBuilder(String baseUrl, HecConfig hecConfig) {
 
     public URI getURI(String endpoint) {
         try {
-            URIBuilder uriBuilder = new URIBuilder(baseUrl)
-                    .setPath(endpoint);
-
+            URIBuilder uriBuilder = new URIBuilder(baseUrl);
+            int idx = endpoint.indexOf('?');
+            if (idx == -1) {
+                // json endpoint
+                uriBuilder = uriBuilder.setPath(endpoint);
+            } else {
+                // in case of raw endpoint, the endpoint will be in form "/services/collector/raw?index=xxx&source=xxx"
+                // extract the path and params via a split on '?'
+                uriBuilder = uriBuilder.setPath(endpoint.substring(0, idx));
+                uriBuilder = uriBuilder.setParameters(URLEncodedUtils.parse(endpoint.substring(idx+1), Consts.UTF_8));
+            }
+            
             if (hecConfig.getAutoExtractTimestamp() != null) {
                 uriBuilder.addParameter(AUTO_EXTRACT_TIMESTAMP_PARAMETER, hecConfig.getAutoExtractTimestamp().toString());
             }
diff --git a/src/test/java/com/splunk/hecclient/HecURIBuilderTest.java b/src/test/java/com/splunk/hecclient/HecURIBuilderTest.java
index 643ac636..6c03a456 100644
--- a/src/test/java/com/splunk/hecclient/HecURIBuilderTest.java
+++ b/src/test/java/com/splunk/hecclient/HecURIBuilderTest.java
@@ -9,17 +9,28 @@
 import static com.splunk.hecclient.JsonEventBatch.ENDPOINT;
 
 public class HecURIBuilderTest {
+    private static final String RAW_ENDPOINT = "/services/collector/raw?index=main&source=source";
     private static final String BASE_URL =  "https://localhost:8088";
     private static final String TOKEN =  "mytoken";
 
     @Test
     public void testDefaultValues() {
-        HecConfig hecConfig = new HecConfig(Collections.emptyList(), TOKEN);
-        HecURIBuilder builder = new HecURIBuilder(BASE_URL, hecConfig);
+        {
+            HecConfig hecConfig = new HecConfig(Collections.emptyList(), TOKEN);
+            HecURIBuilder builder = new HecURIBuilder(BASE_URL, hecConfig);
+
+            URI uri = builder.getURI(ENDPOINT);
+
+            Assert.assertEquals("https://localhost:8088/services/collector/event", uri.toString());
+        }
+        {
+            HecConfig hecConfig = new HecConfig(Collections.emptyList(), TOKEN);
+            HecURIBuilder builder = new HecURIBuilder(BASE_URL, hecConfig);
 
-        URI uri = builder.getURI(ENDPOINT);
+            URI uri = builder.getURI(RAW_ENDPOINT);
 
-        Assert.assertEquals("https://localhost:8088/services/collector/event", uri.toString());
+            Assert.assertEquals("https://localhost:8088/services/collector/raw?index=main&source=source", uri.toString());
+        }
     }
 
     @Test
@@ -46,5 +57,27 @@ public void testAutoExtractTimestamp() {
                             HecURIBuilder.AUTO_EXTRACT_TIMESTAMP_PARAMETER + "=false",
                     uri.toString());
         }
+        {
+            HecConfig hecConfig = new HecConfig(Collections.emptyList(), TOKEN)
+                    .setAutoExtractTimestamp(false);
+            HecURIBuilder builder = new HecURIBuilder(BASE_URL, hecConfig);
+
+            URI uri = builder.getURI(RAW_ENDPOINT);
+
+            Assert.assertEquals("https://localhost:8088/services/collector/raw?index=main&source=source&" +
+                            HecURIBuilder.AUTO_EXTRACT_TIMESTAMP_PARAMETER + "=false",
+                    uri.toString());
+        }
+        {
+            HecConfig hecConfig = new HecConfig(Collections.emptyList(), TOKEN)
+                    .setAutoExtractTimestamp(true);
+            HecURIBuilder builder = new HecURIBuilder(BASE_URL, hecConfig);
+
+            URI uri = builder.getURI(RAW_ENDPOINT);
+
+            Assert.assertEquals("https://localhost:8088/services/collector/raw?index=main&source=source&" +
+                            HecURIBuilder.AUTO_EXTRACT_TIMESTAMP_PARAMETER + "=true",
+                    uri.toString());
+        }
     }
 }
\ No newline at end of file