Allow hyphen as a valid value for sanitization checks

trinodb · Jan 13, 2023 · 9b5a8c5 · 9b5a8c5
1 parent d334c38
commit 9b5a8c5
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 8 deletions.
diff --git a/...-base-jdbc/src/main/java/io/trino/plugin/jdbc/logging/FormatBasedRemoteQueryModifier.java b/...-base-jdbc/src/main/java/io/trino/plugin/jdbc/logging/FormatBasedRemoteQueryModifier.java
@@ -81,7 +81,7 @@ String value(ConnectorSession session)
     private static class SanitizedValuesProvider
             implements Function<ConnectorSession, String>
     {
-        private static final Predicate<String> VALIDATION_MATCHER = Pattern.compile("^[\\w_]*$").asMatchPredicate();
+        private static final Predicate<String> VALIDATION_MATCHER = Pattern.compile("^[\\w_-]*$").asMatchPredicate();
         private final Function<ConnectorSession, String> valueProvider;
         private final String name;
 
@@ -98,7 +98,7 @@ public String apply(ConnectorSession session)
             if (VALIDATION_MATCHER.test(value)) {
                 return value;
             }
-            throw new TrinoException(JDBC_NON_TRANSIENT_ERROR, format("Passed value %s as %s does not meet security criteria. It can contain only letters, digits and underscores", value, name));
+            throw new TrinoException(JDBC_NON_TRANSIENT_ERROR, format("Passed value %s as %s does not meet security criteria. It can contain only letters, digits, underscores and hyphens", value, name));
         }
     }
 }
diff --git a/...e-jdbc/src/test/java/io/trino/plugin/jdbc/logging/TestFormatBasedRemoteQueryModifier.java b/...e-jdbc/src/test/java/io/trino/plugin/jdbc/logging/TestFormatBasedRemoteQueryModifier.java
@@ -16,6 +16,7 @@
 import io.trino.spi.TrinoException;
 import io.trino.spi.security.ConnectorIdentity;
 import io.trino.testing.TestingConnectorSession;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -79,7 +80,7 @@ public void testForSQLInjectionsByTraceToken()
 
         assertThatThrownBy(() -> modifier.apply(connectorSession, "SELECT * from USERS"))
                 .isInstanceOf(TrinoException.class)
-                .hasMessage("Passed value */; DROP TABLE TABLE_A; /* as $TRACE_TOKEN does not meet security criteria. It can contain only letters, digits and underscores");
+                .hasMessage("Passed value */; DROP TABLE TABLE_A; /* as $TRACE_TOKEN does not meet security criteria. It can contain only letters, digits, underscores and hyphens");
     }
 
     @Test
@@ -95,23 +96,41 @@ public void testForSQLInjectionsBySource()
 
         assertThatThrownBy(() -> modifier.apply(connectorSession, "SELECT * from USERS"))
                 .isInstanceOf(TrinoException.class)
-                .hasMessage("Passed value */; DROP TABLE TABLE_A; /* as $SOURCE does not meet security criteria. It can contain only letters, digits and underscores");
+                .hasMessage("Passed value */; DROP TABLE TABLE_A; /* as $SOURCE does not meet security criteria. It can contain only letters, digits, underscores and hyphens");
     }
 
-    @Test
-    public void testFormatWithEmptyValues()
+    @Test(dataProvider = "validValues")
+    public void testFormatWithValidValues(String value)
     {
         TestingConnectorSession connectorSession = TestingConnectorSession.builder()
                 .setIdentity(ConnectorIdentity.ofUser("Alice"))
-                .setSource("")
+                .setSource(value)
+                .setTraceToken(value)
                 .build();
 
         FormatBasedRemoteQueryModifier modifier = createRemoteQueryModifier("source=$SOURCE ttoken=$TRACE_TOKEN");
 
         String modifiedQuery = modifier.apply(connectorSession, "SELECT * FROM USERS");
 
         assertThat(modifiedQuery)
-                .isEqualTo("SELECT * FROM USERS /*source= ttoken=*/");
+                .isEqualTo("SELECT * FROM USERS /*source=%1$s ttoken=%1$s*/".formatted(value));
+    }
+
+    @DataProvider
+    public Object[][] validValues()
+    {
+        return new Object[][] {
+                {"trino"},
+                {"123"},
+                {"1t2r3i4n0"},
+                {"trino-cli"},
+                {"trino_cli"},
+                {"trino-cli_123"},
+                {"123_trino-cli"},
+                {"123-trino_cli"},
+                {"-trino-cli"},
+                {"_trino_cli"}
+        };
     }
 
     private static FormatBasedRemoteQueryModifier createRemoteQueryModifier(String commentFormat)