Simplify how SQL injection is detected when using query.comment-format

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/logging/SessionInterpolatedValues.java

@@ -17,18 +17,14 @@
 import io.trino.spi.connector.ConnectorSession;
 
 import java.util.function.Function;
-import java.util.function.Predicate;
-import java.util.regex.Pattern;
-
-import static java.util.Objects.requireNonNull;
 
 public enum SessionInterpolatedValues
         implements InterpolatedValue<ConnectorSession>
 {
     QUERY_ID(ConnectorSession::getQueryId),
-    SOURCE(new SanitizedValuesProvider(session -> session.getSource().orElse(""), "$SOURCE")),
+    SOURCE(session -> session.getSource().orElse("")),
     USER(ConnectorSession::getUser),
-    TRACE_TOKEN(new SanitizedValuesProvider(session -> session.getTraceToken().orElse(""), "$TRACE_TOKEN"));
+    TRACE_TOKEN(session -> session.getTraceToken().orElse(""));
 
     private final Function<ConnectorSession, String> valueProvider;
 
@@ -42,28 +38,4 @@ public String value(ConnectorSession session)
     {
         return valueProvider.apply(session);
     }
-
-    static class SanitizedValuesProvider
-            implements Function<ConnectorSession, String>
-    {
-        private static final Predicate<String> VALIDATION_MATCHER = Pattern.compile("^[\\w_-]*$").asMatchPredicate();
-        private final Function<ConnectorSession, String> valueProvider;
-        private final String name;
-
-        private SanitizedValuesProvider(Function<ConnectorSession, String> valueProvider, String name)
-        {
-            this.valueProvider = requireNonNull(valueProvider, "valueProvider is null");
-            this.name = requireNonNull(name, "name is null");
-        }
-
-        @Override
-        public String apply(ConnectorSession session)
-        {
-            String value = valueProvider.apply(session);
-            if (VALIDATION_MATCHER.test(value)) {
-                return value;
-            }
-            throw new SecurityException("Passed value %s as %s does not meet security criteria. It can contain only letters, digits, underscores and hyphens".formatted(value, name));
-        }
-    }
 }

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/logging/FormatBasedRemoteQueryModifier.java

@@ -39,11 +39,14 @@ public FormatBasedRemoteQueryModifier(FormatBasedRemoteQueryModifierConfig confi
     @Override
     public String apply(ConnectorSession session, String query)
     {
-        try {
-            return query + " /*" + interpolator.interpolate(session) + "*/";
-        }
-        catch (SecurityException e) {
-            throw new TrinoException(JDBC_NON_TRANSIENT_ERROR, e.getMessage());
+        return query + " /*" + checkForSqlInjection(interpolator.interpolate(session)) + "*/";
+    }
+
+    private String checkForSqlInjection(String sql)
+    {
+        if (sql.contains("*/")) {
+            throw new TrinoException(JDBC_NON_TRANSIENT_ERROR, "Rendering metadata using 'query.comment-format' does not meet security criteria: " + sql);
         }
+        return sql;
     }
 }

plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/logging/TestFormatBasedRemoteQueryModifier.java

@@ -79,23 +79,18 @@ 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, underscores and hyphens");
+                .hasMessageMatching("Rendering metadata using 'query.comment-format' does not meet security criteria: Query=.* Execution for user=Alice with source=source ttoken=\\*/; DROP TABLE TABLE_A; /\\*");
     }
 
     @Test
     public void testForSQLInjectionsBySource()
     {
-        TestingConnectorSession connectorSession = TestingConnectorSession.builder()
-                .setTraceToken("trace_token")
-                .setSource("*/; DROP TABLE TABLE_A; /*")
-                .setIdentity(ConnectorIdentity.ofUser("Alice"))
-                .build();
+        testForSQLInjectionsBySource("*/; DROP TABLE TABLE_A; /*");
+        testForSQLInjectionsBySource("Prefix */; DROP TABLE TABLE_A; /*");
+        testForSQLInjectionsBySource("""
 
-        FormatBasedRemoteQueryModifier modifier = createRemoteQueryModifier("Query=$QUERY_ID Execution for user=$USER with source=$SOURCE ttoken=$TRACE_TOKEN");
 
-        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, underscores and hyphens");
+                                     Multiline */; DROP TABLE TABLE_A; /*""");
     }
 
     @Test
@@ -180,6 +175,21 @@ private void testFormatWithValidValues(String value)
                 .isEqualTo("SELECT * FROM USERS /*source=%1$s ttoken=%1$s*/".formatted(value));
     }
 
+    private void testForSQLInjectionsBySource(String sqlInjection)
+    {
+        TestingConnectorSession connectorSession = TestingConnectorSession.builder()
+                .setTraceToken("trace_token")
+                .setSource(sqlInjection)
+                .setIdentity(ConnectorIdentity.ofUser("Alice"))
+                .build();
+
+        FormatBasedRemoteQueryModifier modifier = createRemoteQueryModifier("Query=$QUERY_ID Execution for user=$USER with source=$SOURCE ttoken=$TRACE_TOKEN");
+
+        assertThatThrownBy(() -> modifier.apply(connectorSession, "SELECT * from USERS"))
+                .isInstanceOf(TrinoException.class)
+                .hasMessageContaining("Rendering metadata using 'query.comment-format' does not meet security criteria: Query=");
+    }
+
     private static FormatBasedRemoteQueryModifier createRemoteQueryModifier(String commentFormat)
     {
         return new FormatBasedRemoteQueryModifier(new FormatBasedRemoteQueryModifierConfig().setFormat(commentFormat));