Apply sanitization checks only for entity specified in comment-format

trinodb · Jan 13, 2023 · d9776a7 · d9776a7
1 parent 160017e
commit d9776a7
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 1 deletion.
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
@@ -43,7 +43,9 @@ public String apply(ConnectorSession session, String query)
     {
         String message = commentFormat;
         for (PredefinedValue predefinedValue : PredefinedValue.values()) {
-            message = message.replaceAll(predefinedValue.getMatchCase(), predefinedValue.value(session));
+            if (message.contains(predefinedValue.getPredefinedValueCode())) {
+                message = message.replaceAll(predefinedValue.getMatchCase(), predefinedValue.value(session));
+            }
         }
         return query + " /*" + message + "*/";
     }

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
@@ -99,6 +99,57 @@ public void testForSQLInjectionsBySource()
                 .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 testFormatQueryModifierWithUser()
+    {
+        TestingConnectorSession connectorSession = TestingConnectorSession.builder()
+                .setIdentity(ConnectorIdentity.ofUser("Alice"))
+                .setSource("$invalid@value")
+                .setTraceToken("#invalid&value")
+                .build();
+
+        FormatBasedRemoteQueryModifier modifier = createRemoteQueryModifier("user=$USER");
+
+        assertThat(modifier.apply(connectorSession, "SELECT * FROM USERS"))
+                .isEqualTo("SELECT * FROM USERS /*user=Alice*/");
+    }
+
+    @Test
+    public void testFormatQueryModifierWithSource()
+    {
+        String validValue = "valid-value";
+        String invalidValue = "$invalid@value";
+
+        TestingConnectorSession connectorSession = TestingConnectorSession.builder()
+                .setIdentity(ConnectorIdentity.ofUser("Alice"))
+                .setSource(validValue)
+                .setTraceToken(invalidValue)
+                .build();
+
+        FormatBasedRemoteQueryModifier modifier = createRemoteQueryModifier("source=$SOURCE");
+
+        assertThat(modifier.apply(connectorSession, "SELECT * FROM USERS"))
+                .isEqualTo("SELECT * FROM USERS /*source=valid-value*/");
+    }
+
+    @Test
+    public void testFormatQueryModifierWithTraceToken()
+    {
+        String validValue = "valid-value";
+        String invalidValue = "$invalid@value";
+
+        TestingConnectorSession connectorSession = TestingConnectorSession.builder()
+                .setIdentity(ConnectorIdentity.ofUser("Alice"))
+                .setSource(invalidValue)
+                .setTraceToken(validValue)
+                .build();
+
+        FormatBasedRemoteQueryModifier modifier = createRemoteQueryModifier("ttoken=$TRACE_TOKEN");
+
+        assertThat(modifier.apply(connectorSession, "SELECT * FROM USERS"))
+                .isEqualTo("SELECT * FROM USERS /*ttoken=valid-value*/");
+    }
+
     @Test(dataProvider = "validValues")
     public void testFormatWithValidValues(String value)
     {