test/security/test_sql_injection.py - doris-mcp-server - Git at Google

 #!/usr/bin/env python3
 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
 # regarding copyright ownership.  The ASF licenses this file
 # to you under the Apache License, Version 2.0 (the
 # "License"); you may not use this file except in compliance
 # with the License.  You may obtain a copy of the License at
 #
 #   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.

 """
 SQL Security Test Suite for Apache Doris MCP Server

 Tests for:
 1. SQL injection prevention via identifier validation
 2. Multi-statement SQL parsing in security validator
 3. auth_context enforcement
 """


 import pytest


 class TestSQLSecurityUtils:
     """Test cases for sql_security_utils module"""

     def test_validate_identifier_accepts_valid_names(self):
         """Test that valid identifiers are accepted"""
         from doris_mcp_server.utils.sql_security_utils import validate_identifier

         valid_names = [
             "users",
             "my_table",
             "Table123",
             "_private_table",
             "CamelCaseTable",
             "table_with_numbers_123",
         ]

         for name in valid_names:
             result = validate_identifier(name, "table")
             assert result == name, f"Valid identifier '{name}' should be accepted"

     def test_validate_identifier_rejects_sql_injection(self):
         """Test that SQL injection attempts are rejected"""
         from doris_mcp_server.utils.sql_security_utils import (
             SQLSecurityError,
             validate_identifier,
         )

         injection_attempts = [
             # Basic SQL injection
             "'; DROP TABLE users; --",
             "table' OR '1'='1",
             "table'; DELETE FROM users; --",

             # Union-based injection
             "table' UNION SELECT * FROM passwords --",

             # Comment injection
             "table/**/OR/**/1=1",
             "table--comment",

             # Special characters
             "table`; DROP TABLE users;",
             'table"; DROP TABLE users;',
             "table\"; DELETE FROM",

             # Backtick escape attempt
             "analytics`; SELECT * FROM sensitive_table;--",

             # Whitespace injection
             "table name with spaces",
             "table\ttab",
             "table\nnewline",
         ]

         for injection in injection_attempts:
             with pytest.raises(SQLSecurityError):
                 validate_identifier(injection, "table")

     def test_validate_identifier_rejects_empty(self):
         """Test that empty identifiers are rejected"""
         from doris_mcp_server.utils.sql_security_utils import (
             SQLSecurityError,
             validate_identifier,
         )

         with pytest.raises(SQLSecurityError):
             validate_identifier("", "table")

         with pytest.raises(SQLSecurityError):
             validate_identifier(None, "table")

     def test_validate_identifier_rejects_too_long(self):
         """Test that identifiers exceeding max length are rejected"""
         from doris_mcp_server.utils.sql_security_utils import (
             SQLSecurityError,
             validate_identifier,
         )

         # Doris identifier max length is typically 64 characters
         long_name = "a" * 100
         with pytest.raises(SQLSecurityError):
             validate_identifier(long_name, "table")

     def test_quote_identifier_adds_backticks(self):
         """Test that quote_identifier properly escapes identifiers"""
         from doris_mcp_server.utils.sql_security_utils import quote_identifier

         assert quote_identifier("my_table", "table") == "`my_table`"
         assert quote_identifier("users", "table") == "`users`"
         assert quote_identifier("Table123", "table") == "`Table123`"

     def test_quote_identifier_validates_first(self):
         """Test that quote_identifier validates before quoting"""
         from doris_mcp_server.utils.sql_security_utils import (
             SQLSecurityError,
             quote_identifier,
         )

         with pytest.raises(SQLSecurityError):
             quote_identifier("'; DROP TABLE users; --", "table")


 class TestSQLSecurityValidator:
     """Test cases for SQLSecurityValidator multi-statement parsing"""

     @pytest.fixture
     def dict_config(self):
         """Create dictionary configuration"""
         return {
             "blocked_keywords": [
                 "DROP", "CREATE", "ALTER", "TRUNCATE",
                 "DELETE", "INSERT", "UPDATE",
                 "GRANT", "REVOKE", "EXEC", "EXECUTE"
             ],
             "max_query_complexity": 100,
             "enable_security_check": True
         }

     @pytest.fixture
     def mock_auth_context(self):
         """Create mock auth context"""
         from doris_mcp_server.utils.security import AuthContext, SecurityLevel
         return AuthContext(
             user_id="test_user",
             roles=["user"],
             security_level=SecurityLevel.INTERNAL
         )

     @pytest.mark.asyncio
     async def test_validates_all_statements(self, dict_config, mock_auth_context):
         """Test that validator checks ALL SQL statements, not just the first"""
         from doris_mcp_server.utils.security import SQLSecurityValidator

         validator = SQLSecurityValidator(dict_config)

         # Multi-statement with injection in second statement
         # This should be BLOCKED
         malicious_sql = "SELECT 1; DROP TABLE users; SELECT 2"

         result = await validator.validate(malicious_sql, mock_auth_context)

         assert not result.is_valid, "Multi-statement injection should be blocked"
         # Check for either DROP keyword detection or SQL injection detection
         error_upper = result.error_message.upper()
         assert ("DROP" in error_upper or
                 "INJECTION" in error_upper or
                 "BLOCKED" in error_upper), f"Expected DROP/injection/blocked in: {result.error_message}"

     @pytest.mark.asyncio
     async def test_blocks_hidden_dangerous_statement(self, dict_config, mock_auth_context):
         """Test that dangerous statements hidden after safe ones are blocked"""
         from doris_mcp_server.utils.security import SQLSecurityValidator

         validator = SQLSecurityValidator(dict_config)

         # Safe statement followed by dangerous one
         malicious_sql = """
         SELECT * FROM users WHERE id = 1;
         DELETE FROM audit_log;
         SELECT 1;
         """

         result = await validator.validate(malicious_sql, mock_auth_context)

         assert not result.is_valid, "Hidden DELETE statement should be blocked"

     @pytest.mark.asyncio
     async def test_allows_safe_multi_statement(self, dict_config, mock_auth_context):
         """Test that multiple safe SELECT statements are allowed"""
         from doris_mcp_server.utils.security import SQLSecurityValidator

         validator = SQLSecurityValidator(dict_config)

         safe_sql = """
         SELECT * FROM users;
         SELECT COUNT(*) FROM orders;
         SELECT id, name FROM products;
         """

         result = await validator.validate(safe_sql, mock_auth_context)

         assert result.is_valid, f"Multiple safe SELECT statements should be allowed, got: {result.error_message}"

     @pytest.mark.asyncio
     async def test_context_switch_injection_blocked(self, dict_config, mock_auth_context):
         """Test that context switch SQL injection is blocked"""
         from doris_mcp_server.utils.security import SQLSecurityValidator

         validator = SQLSecurityValidator(dict_config)

         # Simulating the exec_query_for_mcp attack vector
         injected_sql = """
         USE `analytics`; SELECT * FROM sensitive_table;-- `;
         SELECT * FROM public_table;
         """

         await validator.validate(injected_sql, mock_auth_context)

         # The validator should process all statements
         # Even if USE is allowed, subsequent unauthorized access should be caught
         # by table access checks (if configured)


 class TestExecQueryForMCP:
     """Test cases for exec_query_for_mcp function"""

     @pytest.mark.asyncio
     async def test_rejects_malicious_db_name(self):
         """Test that malicious db_name is rejected"""
         from doris_mcp_server.utils.sql_security_utils import (
             SQLSecurityError,
             validate_identifier,
         )

         # The attack vector from security report
         malicious_db_name = "analytics`; SELECT * FROM sensitive_table;--"

         with pytest.raises(SQLSecurityError):
             validate_identifier(malicious_db_name, "database name")

     @pytest.mark.asyncio
     async def test_rejects_malicious_catalog_name(self):
         """Test that malicious catalog_name is rejected"""
         from doris_mcp_server.utils.sql_security_utils import (
             SQLSecurityError,
             validate_identifier,
         )

         malicious_catalog_name = "internal'; DROP DATABASE production;--"

         with pytest.raises(SQLSecurityError):
             validate_identifier(malicious_catalog_name, "catalog name")


 class TestDependencyAnalysisTools:
     """Test cases for dependency_analysis_tools security fixes"""

     @pytest.mark.asyncio
     async def test_get_tables_metadata_rejects_injection(self):
         """Test that _get_tables_metadata rejects SQL injection"""
         from doris_mcp_server.utils.sql_security_utils import (
             SQLSecurityError,
             validate_identifier,
         )

         # The attack vector from security report
         injection_db_name = "test_db' OR '1'='1' --"

         with pytest.raises(SQLSecurityError):
             validate_identifier(injection_db_name, "database name")


 class TestAuthContextEnforcement:
     """Test cases for auth_context enforcement"""

     def test_execute_requires_auth_context_for_security(self):
         """Test that security checks require auth_context"""
         # This test documents the expected behavior:
         # When auth_context is None, security checks are skipped
         # When auth_context is provided, security checks are performed

         # The fix ensures all execute() calls pass auth_context
         pass

     @pytest.mark.asyncio
     async def test_get_auth_context_returns_context(self):
         """Test that get_auth_context retrieves context from ContextVar"""
         from doris_mcp_server.utils.sql_security_utils import get_auth_context

         # When no context is set, should return None
         result = get_auth_context()
         # This is expected - context is set by HTTP middleware
         assert result is None or hasattr(result, 'user_id')


 class TestIntegrationScenarios:
     """Integration test scenarios for security fixes"""

     def test_attack_scenario_1_permission_bypass(self):
         """
         Attack Scenario 1: Permission Bypass in Multi-Tenant Environment

         Expected: User can only query their own database (db_name="tenant_a_db")
         Attack: Inject "tenant_a_db' OR '1'='1' --" to query ALL databases
         Result: Should be BLOCKED by validate_identifier()
         """
         from doris_mcp_server.utils.sql_security_utils import (
             SQLSecurityError,
             validate_identifier,
         )

         with pytest.raises(SQLSecurityError):
             validate_identifier("tenant_a_db' OR '1'='1' --", "database name")

     def test_attack_scenario_2_union_injection(self):
         """
         Attack Scenario 2: UNION-based Information Disclosure

         Attack: Inject UNION SELECT to extract sensitive data
         Result: Should be BLOCKED
         """
         from doris_mcp_server.utils.sql_security_utils import (
             SQLSecurityError,
             validate_identifier,
         )

         with pytest.raises(SQLSecurityError):
             validate_identifier(
                 "test' UNION SELECT password FROM users --",
                 "database name"
             )

     def test_attack_scenario_3_backtick_escape(self):
         """
         Attack Scenario 3: Backtick Escape Attempt

         Attack: Use backticks to break out of quoted identifier
         Result: Should be BLOCKED
         """
         from doris_mcp_server.utils.sql_security_utils import (
             SQLSecurityError,
             validate_identifier,
         )

         with pytest.raises(SQLSecurityError):
             validate_identifier(
                 "analytics`; SELECT * FROM sensitive_table;--",
                 "database name"
             )


 # Run tests with: pytest tests/test_sql_security.py -v
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "--tb=short"])
	#!/usr/bin/env python3
	# Licensed to the Apache Software Foundation (ASF) under one
	# or more contributor license agreements. See the NOTICE file
	# distributed with this work for additional information
	# regarding copyright ownership. The ASF licenses this file
	# to you under the Apache License, Version 2.0 (the
	# "License"); you may not use this file except in compliance
	# with the License. You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing,
	# software distributed under the License is distributed on an
	# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	# KIND, either express or implied. See the License for the
	# specific language governing permissions and limitations
	# under the License.

	"""
	SQL Security Test Suite for Apache Doris MCP Server

	Tests for:
	1. SQL injection prevention via identifier validation
	2. Multi-statement SQL parsing in security validator
	3. auth_context enforcement
	"""


	import pytest


	class TestSQLSecurityUtils:
	"""Test cases for sql_security_utils module"""

	def test_validate_identifier_accepts_valid_names(self):
	"""Test that valid identifiers are accepted"""
	from doris_mcp_server.utils.sql_security_utils import validate_identifier

	valid_names = [
	"users",
	"my_table",
	"Table123",
	"_private_table",
	"CamelCaseTable",
	"table_with_numbers_123",
	]

	for name in valid_names:
	result = validate_identifier(name, "table")
	assert result == name, f"Valid identifier '{name}' should be accepted"

	def test_validate_identifier_rejects_sql_injection(self):
	"""Test that SQL injection attempts are rejected"""
	from doris_mcp_server.utils.sql_security_utils import (
	SQLSecurityError,
	validate_identifier,
	)

	injection_attempts = [
	# Basic SQL injection
	"'; DROP TABLE users; --",
	"table' OR '1'='1",
	"table'; DELETE FROM users; --",

	# Union-based injection
	"table' UNION SELECT * FROM passwords --",

	# Comment injection
	"table//OR//1=1",
	"table--comment",

	# Special characters
	"table`; DROP TABLE users;",
	'table"; DROP TABLE users;',
	"table\"; DELETE FROM",

	# Backtick escape attempt
	"analytics`; SELECT * FROM sensitive_table;--",

	# Whitespace injection
	"table name with spaces",
	"table\ttab",
	"table\nnewline",
	]

	for injection in injection_attempts:
	with pytest.raises(SQLSecurityError):
	validate_identifier(injection, "table")

	def test_validate_identifier_rejects_empty(self):
	"""Test that empty identifiers are rejected"""
	from doris_mcp_server.utils.sql_security_utils import (
	SQLSecurityError,
	validate_identifier,
	)

	with pytest.raises(SQLSecurityError):
	validate_identifier("", "table")

	with pytest.raises(SQLSecurityError):
	validate_identifier(None, "table")

	def test_validate_identifier_rejects_too_long(self):
	"""Test that identifiers exceeding max length are rejected"""
	from doris_mcp_server.utils.sql_security_utils import (
	SQLSecurityError,
	validate_identifier,
	)

	# Doris identifier max length is typically 64 characters
	long_name = "a" * 100
	with pytest.raises(SQLSecurityError):
	validate_identifier(long_name, "table")

	def test_quote_identifier_adds_backticks(self):
	"""Test that quote_identifier properly escapes identifiers"""
	from doris_mcp_server.utils.sql_security_utils import quote_identifier

	assert quote_identifier("my_table", "table") == "`my_table`"
	assert quote_identifier("users", "table") == "`users`"
	assert quote_identifier("Table123", "table") == "`Table123`"

	def test_quote_identifier_validates_first(self):
	"""Test that quote_identifier validates before quoting"""
	from doris_mcp_server.utils.sql_security_utils import (
	SQLSecurityError,
	quote_identifier,
	)

	with pytest.raises(SQLSecurityError):
	quote_identifier("'; DROP TABLE users; --", "table")


	class TestSQLSecurityValidator:
	"""Test cases for SQLSecurityValidator multi-statement parsing"""

	@pytest.fixture
	def dict_config(self):
	"""Create dictionary configuration"""
	return {
	"blocked_keywords": [
	"DROP", "CREATE", "ALTER", "TRUNCATE",
	"DELETE", "INSERT", "UPDATE",
	"GRANT", "REVOKE", "EXEC", "EXECUTE"
	],
	"max_query_complexity": 100,
	"enable_security_check": True
	}

	@pytest.fixture
	def mock_auth_context(self):
	"""Create mock auth context"""
	from doris_mcp_server.utils.security import AuthContext, SecurityLevel
	return AuthContext(
	user_id="test_user",
	roles=["user"],
	security_level=SecurityLevel.INTERNAL
	)

	@pytest.mark.asyncio
	async def test_validates_all_statements(self, dict_config, mock_auth_context):
	"""Test that validator checks ALL SQL statements, not just the first"""
	from doris_mcp_server.utils.security import SQLSecurityValidator

	validator = SQLSecurityValidator(dict_config)

	# Multi-statement with injection in second statement
	# This should be BLOCKED
	malicious_sql = "SELECT 1; DROP TABLE users; SELECT 2"

	result = await validator.validate(malicious_sql, mock_auth_context)

	assert not result.is_valid, "Multi-statement injection should be blocked"
	# Check for either DROP keyword detection or SQL injection detection
	error_upper = result.error_message.upper()
	assert ("DROP" in error_upper or
	"INJECTION" in error_upper or
	"BLOCKED" in error_upper), f"Expected DROP/injection/blocked in: {result.error_message}"

	@pytest.mark.asyncio
	async def test_blocks_hidden_dangerous_statement(self, dict_config, mock_auth_context):
	"""Test that dangerous statements hidden after safe ones are blocked"""
	from doris_mcp_server.utils.security import SQLSecurityValidator

	validator = SQLSecurityValidator(dict_config)

	# Safe statement followed by dangerous one
	malicious_sql = """
	SELECT * FROM users WHERE id = 1;
	DELETE FROM audit_log;
	SELECT 1;
	"""

	result = await validator.validate(malicious_sql, mock_auth_context)

	assert not result.is_valid, "Hidden DELETE statement should be blocked"

	@pytest.mark.asyncio
	async def test_allows_safe_multi_statement(self, dict_config, mock_auth_context):
	"""Test that multiple safe SELECT statements are allowed"""
	from doris_mcp_server.utils.security import SQLSecurityValidator

	validator = SQLSecurityValidator(dict_config)

	safe_sql = """
	SELECT * FROM users;
	SELECT COUNT(*) FROM orders;
	SELECT id, name FROM products;
	"""

	result = await validator.validate(safe_sql, mock_auth_context)

	assert result.is_valid, f"Multiple safe SELECT statements should be allowed, got: {result.error_message}"

	@pytest.mark.asyncio
	async def test_context_switch_injection_blocked(self, dict_config, mock_auth_context):
	"""Test that context switch SQL injection is blocked"""
	from doris_mcp_server.utils.security import SQLSecurityValidator

	validator = SQLSecurityValidator(dict_config)

	# Simulating the exec_query_for_mcp attack vector
	injected_sql = """
	USE `analytics`; SELECT * FROM sensitive_table;-- `;
	SELECT * FROM public_table;
	"""

	await validator.validate(injected_sql, mock_auth_context)

	# The validator should process all statements
	# Even if USE is allowed, subsequent unauthorized access should be caught
	# by table access checks (if configured)


	class TestExecQueryForMCP:
	"""Test cases for exec_query_for_mcp function"""

	@pytest.mark.asyncio
	async def test_rejects_malicious_db_name(self):
	"""Test that malicious db_name is rejected"""
	from doris_mcp_server.utils.sql_security_utils import (
	SQLSecurityError,
	validate_identifier,
	)

	# The attack vector from security report
	malicious_db_name = "analytics`; SELECT * FROM sensitive_table;--"

	with pytest.raises(SQLSecurityError):
	validate_identifier(malicious_db_name, "database name")

	@pytest.mark.asyncio
	async def test_rejects_malicious_catalog_name(self):
	"""Test that malicious catalog_name is rejected"""
	from doris_mcp_server.utils.sql_security_utils import (
	SQLSecurityError,
	validate_identifier,
	)

	malicious_catalog_name = "internal'; DROP DATABASE production;--"

	with pytest.raises(SQLSecurityError):
	validate_identifier(malicious_catalog_name, "catalog name")


	class TestDependencyAnalysisTools:
	"""Test cases for dependency_analysis_tools security fixes"""

	@pytest.mark.asyncio
	async def test_get_tables_metadata_rejects_injection(self):
	"""Test that _get_tables_metadata rejects SQL injection"""
	from doris_mcp_server.utils.sql_security_utils import (
	SQLSecurityError,
	validate_identifier,
	)

	# The attack vector from security report
	injection_db_name = "test_db' OR '1'='1' --"

	with pytest.raises(SQLSecurityError):
	validate_identifier(injection_db_name, "database name")


	class TestAuthContextEnforcement:
	"""Test cases for auth_context enforcement"""

	def test_execute_requires_auth_context_for_security(self):
	"""Test that security checks require auth_context"""
	# This test documents the expected behavior:
	# When auth_context is None, security checks are skipped
	# When auth_context is provided, security checks are performed

	# The fix ensures all execute() calls pass auth_context
	pass

	@pytest.mark.asyncio
	async def test_get_auth_context_returns_context(self):
	"""Test that get_auth_context retrieves context from ContextVar"""
	from doris_mcp_server.utils.sql_security_utils import get_auth_context

	# When no context is set, should return None
	result = get_auth_context()
	# This is expected - context is set by HTTP middleware
	assert result is None or hasattr(result, 'user_id')


	class TestIntegrationScenarios:
	"""Integration test scenarios for security fixes"""

	def test_attack_scenario_1_permission_bypass(self):
	"""
	Attack Scenario 1: Permission Bypass in Multi-Tenant Environment

	Expected: User can only query their own database (db_name="tenant_a_db")
	Attack: Inject "tenant_a_db' OR '1'='1' --" to query ALL databases
	Result: Should be BLOCKED by validate_identifier()
	"""
	from doris_mcp_server.utils.sql_security_utils import (
	SQLSecurityError,
	validate_identifier,
	)

	with pytest.raises(SQLSecurityError):
	validate_identifier("tenant_a_db' OR '1'='1' --", "database name")

	def test_attack_scenario_2_union_injection(self):
	"""
	Attack Scenario 2: UNION-based Information Disclosure

	Attack: Inject UNION SELECT to extract sensitive data
	Result: Should be BLOCKED
	"""
	from doris_mcp_server.utils.sql_security_utils import (
	SQLSecurityError,
	validate_identifier,
	)

	with pytest.raises(SQLSecurityError):
	validate_identifier(
	"test' UNION SELECT password FROM users --",
	"database name"
	)

	def test_attack_scenario_3_backtick_escape(self):
	"""
	Attack Scenario 3: Backtick Escape Attempt

	Attack: Use backticks to break out of quoted identifier
	Result: Should be BLOCKED
	"""
	from doris_mcp_server.utils.sql_security_utils import (
	SQLSecurityError,
	validate_identifier,
	)

	with pytest.raises(SQLSecurityError):
	validate_identifier(
	"analytics`; SELECT * FROM sensitive_table;--",
	"database name"
	)


	# Run tests with: pytest tests/test_sql_security.py -v
	if __name__ == "__main__":
	pytest.main([__file__, "-v", "--tb=short"])