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 index fd05f6d52d..dce1ec7229 100644 --- 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 headers, List 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 headers, List> lines) { - this.headers = sanitizeHeaders(headers); + this.headers = sanitizeHeaders(separator, headers); this.lines = sanitizeLines(separator, lines); } @@ -63,9 +68,10 @@ public List getLines() { return lines; } - private List sanitizeHeaders(List headers) { + private List sanitizeHeaders(String separator, List headers) { return headers.stream(). map(this::sanitizeCell). + map(cell -> quoteIfRequired(separator, cell)). collect(Collectors.toList()); } @@ -74,6 +80,7 @@ private List sanitizeLines(String separator, List> lines) { for (List 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/src/main/java/com/amazon/opendistroforelasticsearch/sql/executor/csv/CSVResultsExtractor.java b/src/main/java/com/amazon/opendistroforelasticsearch/sql/executor/csv/CSVResultsExtractor.java index 859a8bdd77..8eadc524d2 100644 --- a/src/main/java/com/amazon/opendistroforelasticsearch/sql/executor/csv/CSVResultsExtractor.java +++ b/src/main/java/com/amazon/opendistroforelasticsearch/sql/executor/csv/CSVResultsExtractor.java @@ -331,22 +331,15 @@ private String findFieldValue(String header, Map 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 headers, Map 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 index c2914ad6a7..f10a0013c9 100644 --- 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 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 index f096c7d01e..5d83ee6981 100644 --- 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 headers, List> lines) { return new CSVResult(SEPARATOR, headers, lines); }