From 3e39d5a50e933e3ef54cdd9fb75f7b4dd4fb0e31 Mon Sep 17 00:00:00 2001
From: praveenkrishna <praveenkrishna@tutanota.com>
Date: Thu, 12 Jan 2023 15:20:03 +0100
Subject: [PATCH] Apply sanitization checks only for entity specified in
 comment-format

---
 .../FormatBasedRemoteQueryModifier.java       |  4 +-
 .../TestFormatBasedRemoteQueryModifier.java   | 51 +++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/logging/FormatBasedRemoteQueryModifier.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/logging/FormatBasedRemoteQueryModifier.java
index 54a5746982df..e1ecb6fb2e9b 100644
--- a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/logging/FormatBasedRemoteQueryModifier.java
+++ b/plugin/trino-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/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/logging/TestFormatBasedRemoteQueryModifier.java b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/logging/TestFormatBasedRemoteQueryModifier.java
index a0e9474608a0..faa64733a79e 100644
--- a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/logging/TestFormatBasedRemoteQueryModifier.java
+++ b/plugin/trino-base-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)
     {