fix(security): avoid blocked keyword false positives in SQL literals and comments

doris_mcp_server/utils/security.py

@@ -36,7 +36,7 @@
 
 import sqlparse
 from sqlparse.sql import Statement
-from sqlparse.tokens import Keyword, Name
+from sqlparse.tokens import Comment, Keyword, Name
 
 from .auth_credentials import (
     EMPTY_CREDENTIAL,
@@ -1306,32 +1306,35 @@ def _has_suspicious_quotes_or_comments(self, sql: str) -> bool:
 
     async def _check_blocked_keywords(self, parsed: Statement) -> ValidationResult:
         """Check blocked keywords"""
-        blocked_operations = []
+        blocked_operations: set[str] = set()
+        first_code_token = True
 
-        # Check all tokens in the parsed statement
         for token in parsed.flatten():
-            # Check if token is a keyword (including DML/DDL) or name that matches blocked operations
-            if (
-                token.ttype is Keyword
-                or token.ttype is Name
-                or (token.ttype and str(token.ttype).startswith("Token.Keyword"))
-            ):
-                token_value = token.value.upper().strip()
+            if token.is_whitespace or token.ttype in Comment:
+                continue
+            if token.ttype in Keyword:
+                # sqlparse combines some keyword tokens, such as CREATE OR
+                # REPLACE. Honor both configured compound and individual words.
+                if token.normalized in self.blocked_keywords:
+                    blocked_operations.add(token.normalized)
+                blocked_operations.update(
+                    self.blocked_keywords.intersection(token.normalized.split())
+                )
+            elif first_code_token and token.ttype in Name:
+                # KILL/SHUTDOWN are parsed as names. Only an unquoted command
+                # at the start of a statement is an operation, not a column or
+                # table identifier elsewhere in a read-only query.
+                token_value = token.value.upper()
                 if token_value in self.blocked_keywords:
-                    blocked_operations.append(token_value)
-            # Also check for DDL/DML keywords in token values
-            elif hasattr(token, "value") and token.value:
-                token_value = token.value.upper().strip()
-                for blocked_keyword in self.blocked_keywords:
-                    if blocked_keyword in token_value:
-                        blocked_operations.append(blocked_keyword)
+                    blocked_operations.add(token_value)
+            first_code_token = False
 
         if blocked_operations:
             return ValidationResult(
                 is_valid=False,
-                error_message=f"Contains blocked operations: {', '.join(set(blocked_operations))}",
+                error_message=f"Contains blocked operations: {', '.join(sorted(blocked_operations))}",
                 risk_level="high",
-                blocked_operations=list(set(blocked_operations)),
+                blocked_operations=sorted(blocked_operations),
             )
 
         return ValidationResult(is_valid=True)

test/security/test_sql_validation.py

@@ -109,6 +109,62 @@ async def test_legitimate_comment_is_allowed(self, sql_validator, analyst_contex
 
         assert result.is_valid is True
 
+    @pytest.mark.parametrize(
+        "sql",
+        [
+            "SELECT 'ios_created_time_parse_failure' AS metric",
+            "SELECT 'CREATE DROP DELETE INSERT UPDATE EXECUTE' AS label",
+            'SELECT created_time AS "CREATE" FROM users',
+            "SELECT `CREATE`, `DROP`, `KILL` FROM `UPDATE`",
+            "SELECT created_time, updated_time FROM users",
+            "SELECT 1 /* CREATE TABLE example */",
+            "SELECT 1 /* note: DROP INSERT CREATE UPDATE */",
+            "SELECT 1 -- CREATE TABLE example\n",
+            "SELECT 1 # CREATE TABLE example\n",
+            "WITH p AS (SELECT created_time FROM users) "
+            "SELECT 'ios_created_time_parse_failure' AS metric FROM p",
+        ],
+    )
+    @pytest.mark.asyncio
+    async def test_blocked_words_in_non_executable_tokens_are_allowed(
+        self, sql_validator, analyst_context, sql
+    ):
+        result = await sql_validator.validate(sql, analyst_context)
+
+        assert result.is_valid is True, result.error_message
+        assert result.blocked_operations == []
+
+    @pytest.mark.parametrize(
+        ("sql", "operation"),
+        [
+            ("CREATE OR REPLACE VIEW v AS SELECT 1", "CREATE"),
+            ("/* note */ cReAtE TABLE example (id INT)", "CREATE"),
+            ("/* note */ DROP TABLE example", "DROP"),
+            ("/* note */ UPDATE example SET id = 1", "UPDATE"),
+            ("/* note */ KILL 1", "KILL"),
+            ("/* note */ SHUTDOWN", "SHUTDOWN"),
+        ],
+    )
+    @pytest.mark.asyncio
+    async def test_actual_operations_still_report_blocked_keywords(
+        self, sql_validator, analyst_context, sql, operation
+    ):
+        result = await sql_validator.validate(sql, analyst_context)
+
+        assert result.is_valid is False
+        assert operation in result.blocked_operations
+
+    @pytest.mark.asyncio
+    async def test_configured_combined_keyword_still_blocked(
+        self, sql_validator, analyst_context
+    ):
+        sql_validator.blocked_keywords = {"UNION ALL"}
+
+        result = await sql_validator.validate("SELECT 1 UNION ALL SELECT 2", analyst_context)
+
+        assert result.is_valid is False
+        assert result.blocked_operations == ["UNION ALL"]
+
     @pytest.mark.asyncio
     async def test_legitimate_union_is_allowed(self, sql_validator, analyst_context):
         """Test ordinary UNION queries are not rejected as sensitive extraction."""