Escape comma for CSV header and all queries (opendistro-for-elasticse…

…arch#456)
penghuo · May 4, 2020 · ced1cd5 · ced1cd5
1 parent ba9cd2e
commit ced1cd5
Show file tree

Hide file tree

Showing 4 changed files with 98 additions and 12 deletions.
diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/sql/executor/csv/CSVResult.java b/src/main/java/com/amazon/opendistroforelasticsearch/sql/executor/csv/CSVResult.java
@@ -41,8 +41,13 @@ public CSVResult(List<String> headers, List<String> lines) {
         this.lines = lines;
     }
 
+    /**
+     * Sanitize both headers and data lines by:
+     *  1) First prepend single quote if first char is sensitive (= - + @)
+     *  2) Second double quote entire cell if any comma found
+     */
     public CSVResult(String separator, List<String> headers, List<List<String>> lines) {
-        this.headers = sanitizeHeaders(headers);
+        this.headers = sanitizeHeaders(separator, headers);
         this.lines = sanitizeLines(separator, lines);
     }
 
@@ -63,9 +68,10 @@ public List<String> getLines() {
         return lines;
     }
 
-    private List<String> sanitizeHeaders(List<String> headers) {
+    private List<String> sanitizeHeaders(String separator, List<String> headers) {
         return headers.stream().
                        map(this::sanitizeCell).
+                       map(cell -> quoteIfRequired(separator, cell)).
                        collect(Collectors.toList());
     }
 
@@ -74,6 +80,7 @@ private List<String> sanitizeLines(String separator, List<List<String>> lines) {
         for (List<String> line : lines) {
             result.add(line.stream().
                             map(this::sanitizeCell).
+                            map(cell -> quoteIfRequired(separator, cell)).
                             collect(Collectors.joining(separator)));
         }
         return result;
@@ -86,6 +93,12 @@ private String sanitizeCell(String cell) {
         return cell;
     }
 
+    private String quoteIfRequired(String separator, String cell) {
+        final String quote = "\"";
+        return cell.contains(separator)
+            ? quote + cell.replaceAll("\"", "\"\"") + quote : cell;
+    }
+
     private boolean isStartWithSensitiveChar(String cell) {
         return SENSITIVE_CHAR.stream().
                               anyMatch(cell::startsWith);

diff --git a/...main/java/com/amazon/opendistroforelasticsearch/sql/executor/csv/CSVResultsExtractor.java b/...main/java/com/amazon/opendistroforelasticsearch/sql/executor/csv/CSVResultsExtractor.java
@@ -331,22 +331,15 @@ private String findFieldValue(String header, Map<String, Object> doc, boolean fl
                     return "";
                 }
             }
-            return quoteValueIfRequired(innerDoc.toString(), separator);
+            return innerDoc.toString();
         } else {
             if (doc.containsKey(header)) {
-                return quoteValueIfRequired(String.valueOf(doc.get(header)), separator);
+                return String.valueOf(doc.get(header));
             }
         }
         return "";
     }
 
-    private String quoteValueIfRequired(final String input, final String separator) {
-        final String quote = "\"";
-
-        return input.contains(separator)
-                ? quote + input.replaceAll("\"", "\"\"") + quote : input;
-    }
-
     private void mergeHeaders(Set<String> headers, Map<String, Object> doc, boolean flat) {
         if (!flat) {
             headers.addAll(doc.keySet());

diff --git a/src/test/java/com/amazon/opendistroforelasticsearch/sql/esintgtest/CsvFormatResponseIT.java b/src/test/java/com/amazon/opendistroforelasticsearch/sql/esintgtest/CsvFormatResponseIT.java
@@ -94,7 +94,7 @@ public void specificPercentilesIntAndDouble() throws IOException {
         final String result = executeQueryWithStringOutput(query);
 
         final String[] unexpectedPercentiles = {"1.0", "5.0", "25.0", "50.0", "75.0", "95.0", "99.0"};
-        final String expectedHeaders = "PERCENTILES(age,10,49.0).10.0,PERCENTILES(age,10,49.0).49.0";
+        final String expectedHeaders = "\"PERCENTILES(age,10,49.0).10.0\",\"PERCENTILES(age,10,49.0).49.0\"";
         Assert.assertThat(result, containsString(expectedHeaders));
         for (final String unexpectedPercentile : unexpectedPercentiles) {
             Assert.assertThat(result, not(containsString("PERCENTILES(age,10,49.0)." + unexpectedPercentile)));
@@ -634,6 +634,32 @@ public void sensitiveCharacterSanitizeTest() throws IOException {
         Assert.assertTrue(lines.get(0).contains("'@cmd|' /C notepad'!_xlbgnm.A1"));
     }
 
+    @Test
+    public void sensitiveCharacterSanitizeAndQuotedTest() throws IOException {
+        String requestBody =
+            "{" +
+            "  \"=cmd|' /C notepad'!_xlbgnm.A1,,\": \",+cmd|' /C notepad'!_xlbgnm.A1\",\n" +
+            "  \",@cmd|' /C notepad'!_xlbgnm.A1\": \"+cmd|' /C notepad,,'!_xlbgnm.A1\",\n" +
+            "  \"-cmd|' /C notepad,,'!_xlbgnm.A1\": \",,,@cmd|' /C notepad'!_xlbgnm.A1\"\n" +
+            "}";
+
+        Request request = new Request("PUT", "/userdata2/_doc/1?refresh=true");
+        request.setJsonEntity(requestBody);
+        TestUtils.performRequest(client(), request);
+
+        CSVResult csvResult = executeCsvRequest("SELECT * FROM userdata2", false, false, false, false);
+        String headers = String.join(",", csvResult.getHeaders());
+        Assert.assertTrue(headers.contains("\"'=cmd|' /C notepad'!_xlbgnm.A1,,\""));
+        Assert.assertTrue(headers.contains("\",@cmd|' /C notepad'!_xlbgnm.A1\""));
+        Assert.assertTrue(headers.contains("\"'-cmd|' /C notepad,,'!_xlbgnm.A1\""));
+
+        List<String> lines = csvResult.getLines();
+        Assert.assertEquals(1, lines.size());
+        Assert.assertTrue(lines.get(0).contains("\",+cmd|' /C notepad'!_xlbgnm.A1\""));
+        Assert.assertTrue(lines.get(0).contains("\"'+cmd|' /C notepad,,'!_xlbgnm.A1\""));
+        Assert.assertTrue(lines.get(0).contains("\",,,@cmd|' /C notepad'!_xlbgnm.A1\""));
+    }
+
     private void verifyFieldOrder(final String[] expectedFields) throws IOException {
 
         final String fields = String.join(", ", expectedFields);

diff --git a/src/test/java/com/amazon/opendistroforelasticsearch/sql/executor/csv/CSVResultTest.java b/src/test/java/com/amazon/opendistroforelasticsearch/sql/executor/csv/CSVResultTest.java
@@ -66,6 +66,60 @@ public void getLinesShouldReturnLinesSanitized() {
         );
     }
 
+    @Test
+    public void getHeadersShouldReturnHeadersQuotedIfRequired() {
+        CSVResult csv = csv(headers("na,me", ",,age"), lines(line("John", "30")));
+        assertEquals(
+            headers("\"na,me\"", "\",,age\""),
+            csv.getHeaders()
+        );
+    }
+
+    @Test
+    public void getLinesShouldReturnLinesQuotedIfRequired() {
+        CSVResult csv = csv(headers("name", "age"), lines(line("John,Smith", "30,,,")));
+        assertEquals(
+            line("\"John,Smith\",\"30,,,\""),
+            csv.getLines()
+        );
+    }
+
+    @Test
+    public void getHeadersShouldReturnHeadersBothSanitizedAndQuotedIfRequired() {
+        CSVResult csv = csv(headers("na,+me", ",,,=age", "=city,"), lines(line("John", "30", "Seattle")));
+        assertEquals(
+            headers("\"na,+me\"", "\",,,=age\"", "\"'=city,\""),
+            csv.getHeaders()
+        );
+    }
+
+    @Test
+    public void getLinesShouldReturnLinesBothSanitizedAndQuotedIfRequired() {
+        CSVResult csv = csv(
+            headers("name", "city"),
+            lines(
+                line("John", "Seattle"),
+                line("John", "=Seattle"),
+                line("John", "+Sea,ttle"),
+                line(",-John", "Seattle"),
+                line(",,,@John", "Seattle"),
+                line("John", "Seattle=")
+            )
+        );
+
+        assertEquals(
+            line(
+                "John,Seattle",
+                "John,'=Seattle",
+                "John,\"'+Sea,ttle\"",
+                "\",-John\",Seattle",
+                "\",,,@John\",Seattle",
+                "John,Seattle="
+            ),
+            csv.getLines()
+        );
+    }
+
     private CSVResult csv(List<String> headers, List<List<String>> lines) {
         return new CSVResult(SEPARATOR, headers, lines);
     }