diff --git a/doris_mcp_server/utils/security.py b/doris_mcp_server/utils/security.py index 1f47578..23a8ea4 100644 --- a/doris_mcp_server/utils/security.py +++ b/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) diff --git a/test/security/test_sql_validation.py b/test/security/test_sql_validation.py index 79f4b0b..9806e8d 100644 --- a/test/security/test_sql_validation.py +++ b/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."""