diff --git a/.env.example b/.env.example index 1b9ff10..0f731a3 100644 --- a/.env.example +++ b/.env.example @@ -100,6 +100,9 @@ TOKEN_FILE_PATH=tokens.json ENABLE_TOKEN_EXPIRY=true DEFAULT_TOKEN_EXPIRY_HOURS=720 TOKEN_HASH_ALGORITHM=sha256 +# Cache successful token-bound Doris route checks so MCP pings and discovery +# requests do not reconnect on every request. Set 0 to validate every request. +TOKEN_DB_VALIDATION_TTL_SECONDS=30 # No static bearer token is shipped. Generate one outside source control: # python -c 'import secrets; print(secrets.token_urlsafe(32))' # TOKEN_ADMIN= diff --git a/CHANGELOG.md b/CHANGELOG.md index 1e9fa59..18b3470 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -86,6 +86,8 @@ under **Unreleased** until a new version is selected and published. - Preserved service availability after malformed requests, unknown methods, header mismatches, unsupported versions, and missing capabilities. +- Kept static-token Doris pools usable after repeated query timeouts, and + stopped per-request token validation from mutating the shared global pool. - Stopped list operations from converting Doris, permission, or internal failures into successful empty resource, tool, or prompt collections. - Corrected resource and prompt error semantics and tool `isError` results. diff --git a/README.md b/README.md index 6f735a7..c3891b9 100644 --- a/README.md +++ b/README.md @@ -797,7 +797,7 @@ The Doris MCP Server includes a comprehensive enterprise-grade security framewor * **🔐 Multi-Authentication System**: Complete Token, JWT, and OAuth authentication with independent control switches * **🔗 Token-Bound Database Configuration**: Revolutionary approach allowing tokens to carry their own database connection parameters * **🔄 Hot Reload Security**: Zero-downtime security configuration updates with intelligent token revalidation -* **⚡ Immediate Validation**: Real-time database and authentication validation at connection time +* **⚡ Route-Safe Validation**: Token-bound Doris routes are validated through their dedicated pools, with a short success cache so pings do not reconnect on every request * **🛡️ Role-Based Authorization**: Advanced RBAC with four-tier security classification * **🚫 Enhanced SQL Security**: Advanced SQL injection protection with improved pattern detection * **🎭 Intelligent Data Masking**: Automatic sensitive data masking with user-based permissions @@ -819,6 +819,7 @@ ENABLE_OAUTH_AUTH=false # Enable OAuth authentication # Token Management (New in v0.6.0) TOKEN_FILE_PATH=tokens.json # Token configuration file TOKEN_HOT_RELOAD=true # Enable hot reloading +TOKEN_DB_VALIDATION_TTL_SECONDS=30 # Cache successful Doris route checks # Legacy Configuration (Deprecated) # AUTH_TYPE=token # Use individual switches instead diff --git a/doris_mcp_server/utils/config.py b/doris_mcp_server/utils/config.py index 6102c81..ff68dc3 100644 --- a/doris_mcp_server/utils/config.py +++ b/doris_mcp_server/utils/config.py @@ -463,6 +463,7 @@ class SecurityConfig: enable_token_expiry: bool = True # Enable token expiration default_token_expiry_hours: int = 24 * 30 # Default expiry: 30 days token_hash_algorithm: str = "sha256" # Token hashing algorithm: sha256, sha512 + token_db_validation_ttl_seconds: int = 30 # Token Management Security (New in v0.6.0) enable_http_token_management: bool = False # Enable HTTP token management endpoints (default: disabled for security) @@ -1288,6 +1289,12 @@ def from_env(cls, env_file: str | None = None) -> "DorisConfig": os.getenv("DEFAULT_TOKEN_EXPIRY_HOURS", str(config.security.default_token_expiry_hours)) ) config.security.token_hash_algorithm = os.getenv("TOKEN_HASH_ALGORITHM", config.security.token_hash_algorithm) + config.security.token_db_validation_ttl_seconds = int( + os.getenv( + "TOKEN_DB_VALIDATION_TTL_SECONDS", + str(config.security.token_db_validation_ttl_seconds), + ) + ) # Token Management Security Configuration (New in v0.6.0) config.security.enable_http_token_management = ( @@ -1780,6 +1787,11 @@ def validate(self) -> list[str]: if self.security.token_expiry <= 0: errors.append("Token expiry time must be greater than 0") + if not 0 <= self.security.token_db_validation_ttl_seconds <= 3600: + errors.append( + "Token database validation TTL must be in the range 0-3600 seconds" + ) + if self.security.max_query_complexity <= 0: errors.append("Maximum query complexity must be greater than 0") @@ -2052,6 +2064,10 @@ def normalize_effective_auth_config( ) except ValueError as exc: raise AuthConfigError(str(exc)) from exc + if not 0 <= config.security.token_db_validation_ttl_seconds <= 3600: + raise AuthConfigError( + "TOKEN_DB_VALIDATION_TTL_SECONDS must be in the range 0-3600" + ) modern_auth_explicit = any( [ diff --git a/doris_mcp_server/utils/db.py b/doris_mcp_server/utils/db.py index 71760c6..b596525 100644 --- a/doris_mcp_server/utils/db.py +++ b/doris_mcp_server/utils/db.py @@ -1263,7 +1263,11 @@ async def get_connection_for_token( except Exception as e: self.logger.error( - f"Session {session_id}: Failed to acquire connection from token pool: {e}" + "Session %s: Failed to acquire connection from token pool " + "(%s): %s", + session_id, + type(e).__name__, + str(e) or "", ) raise @@ -1368,7 +1372,7 @@ async def close_all_token_pools(self) -> None: self._token_pool_generations.clear() async def configure_for_token(self, token: str) -> tuple[bool, str]: - """Configure connection manager for token with new priority logic + """Validate the database route selected for a static token. Priority: Token-bound DB config > .env config > error @@ -1381,61 +1385,45 @@ async def configure_for_token(self, token: str) -> tuple[bool, str]: Raises: RuntimeError: If no valid database configuration is available """ - try: - # Priority 1: Try token-bound database config first - if self.token_manager: - db_config = self.token_manager.get_database_config_by_token(token) - if db_config: - # Convert DatabaseConfig to dictionary - token_db_config: DatabasePoolConfig = { - "host": db_config.host, - "port": db_config.port, - "user": db_config.user, - "password": db_config.password, - "database": db_config.database, - "charset": db_config.charset, - } - - # Check if token-bound config is valid - if not self._is_config_empty( - token_db_config["host"] - ) and not self._is_config_empty(token_db_config["user"]): - self.logger.info( - f"Using token-bound database configuration for host: {token_db_config['host']}" - ) - self.active_db_config = token_db_config - self._update_db_params_from_config(self.active_db_config) - - # Create/recreate connection pool with token-bound config - await self._ensure_pool_with_current_config() - - return True, "token-bound" + current_token_config = self._get_current_token_db_config(token) + uses_token_config = bool( + current_token_config + and not self._is_config_empty(current_token_config.get("host")) + and not self._is_config_empty(current_token_config.get("user")) + ) + config_source = "token-bound" if uses_token_config else "global-env" - # Priority 2: Use global .env config if available - if self._has_valid_global_config(): - self.logger.info("Using global .env database configuration") - self.active_db_config = self.original_db_config.copy() - self._update_db_params_from_config(self.active_db_config) - - # Create/recreate connection pool with global config - await self._ensure_pool_with_current_config() - - return True, "global-env" - - # Priority 3: No valid configuration available - error_msg = ( - "No valid database configuration available for this token. " - "Please contact administrator to:\n" - "1. Add database configuration to tokens.json for this token, OR\n" - "2. Configure valid global database settings in .env file\n" - "Required fields: DB_HOST, DB_USER" + connection: DorisConnection | None = None + try: + # Validate the same dedicated route that query execution will use. + # Never mutate active_db_config or rebuild the shared global pool + # while authenticating a tenant token. + connection = await self.get_connection_for_token( + token, + f"token_validation_{self._get_token_hash(token)[:8]}", ) - self.logger.error(error_msg) - raise RuntimeError(error_msg) - + result = await asyncio.wait_for( + connection.execute( + "SELECT 1 AS connection_check", + mask_result=False, + max_rows=1, + max_bytes=256, + ), + timeout=self.connect_timeout, + ) + if result.data != [{"connection_check": 1}]: + raise RuntimeError("Token database validation query returned no result") + return True, config_source except Exception as e: - self.logger.error(f"Failed to configure database for token: {e}") + self.logger.error( + "Failed to validate database route for token (%s): %s", + type(e).__name__, + str(e) or "", + ) raise + finally: + if connection is not None: + await self.release_connection_for_token(token, connection) async def _ensure_pool_with_current_config(self) -> None: """Ensure connection pool exists with current configuration""" diff --git a/doris_mcp_server/utils/secret_policy.py b/doris_mcp_server/utils/secret_policy.py index 60a3b76..8eef510 100644 --- a/doris_mcp_server/utils/secret_policy.py +++ b/doris_mcp_server/utils/secret_policy.py @@ -43,6 +43,7 @@ "TOKEN_FILE_PATH", "TOKEN_HASH_ALGORITHM", "TOKEN_HOT_RELOAD", + "TOKEN_DB_VALIDATION_TTL_SECONDS", "TOKEN_SECRET", } ) diff --git a/doris_mcp_server/utils/security.py b/doris_mcp_server/utils/security.py index d721028..97f64a5 100644 --- a/doris_mcp_server/utils/security.py +++ b/doris_mcp_server/utils/security.py @@ -22,7 +22,10 @@ from __future__ import annotations +import asyncio +import hashlib import re +import time from collections.abc import Callable, Mapping from contextvars import ContextVar from contextvars import Token as ContextToken @@ -191,6 +194,8 @@ def __init__( # Track initialization state self._initialized = False + self._token_database_validation_cache: dict[str, float] = {} + self._token_database_validation_locks: dict[str, asyncio.Lock] = {} async def initialize(self) -> None: """Initialize security manager components""" @@ -215,6 +220,8 @@ async def shutdown(self) -> None: """Shutdown security manager components""" try: await self.auth_provider.shutdown() + self._token_database_validation_cache.clear() + self._token_database_validation_locks.clear() self._initialized = False self.logger.info("DorisSecurityManager shutdown completed") @@ -484,10 +491,12 @@ async def _validate_token_database_config( token: str, token_info: TokenInfo, ) -> None: - """Validate database configuration for token immediately during authentication + """Validate a token database route with a short successful-result cache. - This ensures database connectivity issues are caught at authentication time, - not during query execution, providing better user experience. + Authentication runs for every MCP request, including pings and capability + discovery. Reconnecting to Doris on each request creates avoidable pool + pressure, so successful validations are cached by token and database + configuration fingerprint. Failures are never cached. Args: token: Raw authentication token @@ -496,29 +505,79 @@ async def _validate_token_database_config( Raises: ValueError: If database configuration is invalid or connection fails """ - try: - if not self.connection_manager: - self.logger.warning( - "Connection manager not available for immediate database validation" - ) - return + if not self.connection_manager: + self.logger.warning( + "Connection manager not available for immediate database validation" + ) + return - # Configure and test database connection for this token - success, config_source = await self.connection_manager.configure_for_token( - token + cache_key = self._token_database_validation_cache_key(token, token_info) + ttl_seconds = max( + 0, + int(self.config.security.token_db_validation_ttl_seconds), + ) + now = time.monotonic() + if self._token_database_validation_cache.get(cache_key, 0.0) > now: + self.logger.debug( + "Using cached database validation for token %s", + token_info.token_id, ) + return - if success: + lock = self._token_database_validation_locks.setdefault( + cache_key, + asyncio.Lock(), + ) + async with lock: + now = time.monotonic() + if self._token_database_validation_cache.get(cache_key, 0.0) > now: + return + try: + success, config_source = ( + await self.connection_manager.configure_for_token(token) + ) + if not success: + raise ValueError("Database configuration validation failed") + if ttl_seconds: + self._token_database_validation_cache[cache_key] = ( + time.monotonic() + ttl_seconds + ) self.logger.info( - f"Database configuration validated successfully for token {token_info.token_id} (source: {config_source})" + "Database configuration validated successfully for token %s " + "(source: %s)", + token_info.token_id, + config_source, ) - else: - raise ValueError("Database configuration validation failed") + except Exception as e: + self._token_database_validation_cache.pop(cache_key, None) + error_msg = ( + "Database configuration validation failed for token " + f"{token_info.token_id}: {str(e) or type(e).__name__}" + ) + self.logger.error(error_msg) + raise ValueError(error_msg) from e - except Exception as e: - error_msg = f"Database configuration validation failed for token {token_info.token_id}: {str(e)}" - self.logger.error(error_msg) - raise ValueError(error_msg) + @staticmethod + def _token_database_validation_cache_key( + token: str, + token_info: TokenInfo, + ) -> str: + """Return a non-secret cache key that changes with bound DB credentials.""" + database_config = token_info.database_config + route = ( + getattr(database_config, "host", ""), + getattr(database_config, "port", ""), + getattr(database_config, "user", ""), + getattr(database_config, "password", ""), + getattr(database_config, "database", ""), + getattr(database_config, "charset", ""), + ) + digest = hashlib.sha256() + digest.update(token.encode("utf-8")) + for value in route: + digest.update(b"\0") + digest.update(str(value).encode("utf-8")) + return digest.hexdigest() class AuthenticationProvider: diff --git a/test/integration/test_real_doris_transports.py b/test/integration/test_real_doris_transports.py index cc6ea74..48e1236 100644 --- a/test/integration/test_real_doris_transports.py +++ b/test/integration/test_real_doris_transports.py @@ -45,10 +45,12 @@ from typing import Any import httpx +import httpx2 import pymysql import pytest from mcp import Client, MCPError, StdioServerParameters from mcp.client.stdio import stdio_client +from mcp.client.streamable_http import streamable_http_client from mcp.types import ( SubscriptionFilter, SubscriptionsListenRequest, @@ -57,6 +59,7 @@ ) from doris_mcp_server import __version__ +from doris_mcp_server.utils.secret_policy import build_token_digest pytestmark = [ pytest.mark.integration, @@ -313,6 +316,28 @@ async def _http_client( yield client +@asynccontextmanager +async def _authenticated_http_client( + environment: dict[str, str], + *, + bearer_token: str, + read_timeout_seconds: int = 15, +) -> AsyncIterator[Client]: + async with _http_process(environment) as base_url: + async with httpx2.AsyncClient( + headers={"Authorization": f"Bearer {bearer_token}"}, + ) as http_client: + transport = streamable_http_client( + f"{base_url}/mcp", + http_client=http_client, + ) + async with Client( + transport, + read_timeout_seconds=read_timeout_seconds, + ) as client: + yield client + + @asynccontextmanager async def _stdio_client( environment: dict[str, str], @@ -537,6 +562,89 @@ async def test_real_doris_result_boundaries_and_cancellation( assert recovered_payload["data"] == [{"recovered": 1}] +async def test_real_doris_token_pool_recovers_after_repeated_query_timeouts( + doris_sandbox: DorisSandbox, + tmp_path: Path, +) -> None: + bearer_token = secrets.token_urlsafe(48) + token_file = tmp_path / "tokens.json" + token_file.write_text( + json.dumps( + { + "version": "2.0", + "tokens": [ + { + "token_id": "real-timeout-regression", + "token_digest": build_token_digest( + bearer_token, + "sha256", + ), + "created_at": "2026-07-30T00:00:00Z", + "expires_at": None, + "last_used": None, + "description": "real token pool timeout regression", + "is_active": True, + "database_config": { + "host": doris_sandbox.settings.host, + "port": doris_sandbox.settings.port, + "user": doris_sandbox.readonly_user, + "password": doris_sandbox.readonly_password, + "database": doris_sandbox.settings.database, + "charset": "UTF8", + }, + } + ], + } + ), + encoding="utf-8", + ) + token_file.chmod(0o600) + environment = _server_environment( + doris_sandbox.settings, + user=doris_sandbox.settings.user, + password=doris_sandbox.settings.password, + ) + environment.update( + { + "ENABLE_TOKEN_AUTH": "true", + "TOKEN_FILE_PATH": str(token_file), + "TOKEN_DB_VALIDATION_TTL_SECONDS": "60", + "DORIS_MAX_CONNECTIONS": "1", + "QUERY_TIMEOUT": "5", + } + ) + + async with _authenticated_http_client( + environment, + bearer_token=bearer_token, + read_timeout_seconds=10, + ) as client: + for attempt in range(3): + timed_out, timed_out_payload = await _exec_query( + client, + "SELECT SLEEP(3) AS slept", + max_rows=1, + max_bytes=256, + timeout=1, + ) + assert timed_out.is_error is True + assert timed_out_payload["error_type"] == "timeout" + + recovery_started = time.monotonic() + recovered, recovered_payload = await _exec_query( + client, + "SELECT CURRENT_USER() AS current_user", + max_rows=1, + max_bytes=256, + timeout=5, + ) + assert time.monotonic() - recovery_started < 3 + assert recovered.is_error is False, f"recovery {attempt + 1} failed" + assert doris_sandbox.readonly_user in recovered_payload["data"][0][ + "current_user" + ] + + async def _collect_list_pages( list_method: Any, *, diff --git a/test/security/test_token_database_validation.py b/test/security/test_token_database_validation.py new file mode 100644 index 0000000..f719454 --- /dev/null +++ b/test/security/test_token_database_validation.py @@ -0,0 +1,131 @@ +#!/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. +"""Regression tests for static-token database route validation.""" + +from unittest.mock import AsyncMock + +import pytest + +from doris_mcp_server.auth.token_manager import DatabaseConfig, TokenInfo +from doris_mcp_server.utils.config import ( + AuthConfigError, + DorisConfig, + normalize_effective_auth_config, +) +from doris_mcp_server.utils.secret_policy import ( + is_static_token_environment_variable, +) +from doris_mcp_server.utils.security import DorisSecurityManager + + +def _token_info(password: str = "tenant-password") -> TokenInfo: + return TokenInfo( + token_id="tenant-alpha", + database_config=DatabaseConfig( + host="tenant-fe", + port=9030, + user="tenant-reader", + password=password, + database="tenant-db", + charset="UTF8", + ), + ) + + +@pytest.mark.asyncio +async def test_successful_token_database_validation_is_cached(test_config): + connection_manager = AsyncMock() + connection_manager.configure_for_token.return_value = (True, "token-bound") + test_config.security.token_db_validation_ttl_seconds = 30 + manager = DorisSecurityManager(test_config, connection_manager) + token_info = _token_info() + + await manager._validate_token_database_config("high-entropy-token", token_info) + await manager._validate_token_database_config("high-entropy-token", token_info) + + connection_manager.configure_for_token.assert_awaited_once_with( + "high-entropy-token" + ) + + +@pytest.mark.asyncio +async def test_database_credential_rotation_invalidates_validation_cache(test_config): + connection_manager = AsyncMock() + connection_manager.configure_for_token.return_value = (True, "token-bound") + test_config.security.token_db_validation_ttl_seconds = 30 + manager = DorisSecurityManager(test_config, connection_manager) + + await manager._validate_token_database_config( + "high-entropy-token", + _token_info("password-v1"), + ) + await manager._validate_token_database_config( + "high-entropy-token", + _token_info("password-v2"), + ) + + assert connection_manager.configure_for_token.await_count == 2 + + +@pytest.mark.asyncio +async def test_failed_token_database_validation_is_not_cached(test_config): + connection_manager = AsyncMock() + connection_manager.configure_for_token.side_effect = [ + RuntimeError("Doris unavailable"), + (True, "token-bound"), + ] + test_config.security.token_db_validation_ttl_seconds = 30 + manager = DorisSecurityManager(test_config, connection_manager) + token_info = _token_info() + + with pytest.raises(ValueError, match="Doris unavailable"): + await manager._validate_token_database_config( + "high-entropy-token", + token_info, + ) + await manager._validate_token_database_config("high-entropy-token", token_info) + + assert connection_manager.configure_for_token.await_count == 2 + + +def test_token_database_validation_ttl_loads_from_environment( + monkeypatch, + tmp_path, +): + monkeypatch.setenv("TOKEN_DB_VALIDATION_TTL_SECONDS", "45") + + config = DorisConfig.from_env(str(tmp_path / "missing.env")) + + assert config.security.token_db_validation_ttl_seconds == 45 + + +def test_token_database_validation_ttl_is_not_treated_as_a_bearer_token(): + assert ( + is_static_token_environment_variable( + "TOKEN_DB_VALIDATION_TTL_SECONDS" + ) + is False + ) + + +def test_token_database_validation_ttl_rejects_out_of_range_values(): + config = DorisConfig() + config.security.token_db_validation_ttl_seconds = 3601 + + with pytest.raises(AuthConfigError, match="TOKEN_DB_VALIDATION_TTL_SECONDS"): + normalize_effective_auth_config(config) diff --git a/test/utils/test_doris_user_pool_manager.py b/test/utils/test_doris_user_pool_manager.py index 8540382..bbbde42 100644 --- a/test/utils/test_doris_user_pool_manager.py +++ b/test/utils/test_doris_user_pool_manager.py @@ -413,3 +413,61 @@ async def test_release_routed_connection_releases_closed_raw_to_captured_owner( assert old_owner.release_calls == [raw_connection] assert current_pool.release_calls == [] assert raw_connection.ensure_closed_calls == 0 + + +@pytest.mark.asyncio +async def test_configure_for_token_validates_dedicated_route_without_global_mutation( + manager, + monkeypatch, +): + token = "tenant-token" + manager.token_manager = SimpleNamespace( + get_database_config_by_token=lambda raw: ( + SimpleNamespace( + host="tenant-fe", + port=19030, + user="tenant_reader", + password="tenant_pw", + database="tenant_db", + charset="utf8", + ) + if raw == token + else None + ) + ) + connection = SimpleNamespace( + execute=AsyncMock( + return_value=QueryResult( + data=[{"connection_check": 1}], + metadata={}, + execution_time=0.0, + row_count=1, + sql="SELECT 1 AS connection_check", + ) + ) + ) + get_connection = AsyncMock(return_value=connection) + release_connection = AsyncMock() + monkeypatch.setattr(manager, "get_connection_for_token", get_connection) + monkeypatch.setattr( + manager, + "release_connection_for_token", + release_connection, + ) + original_config = manager.original_db_config.copy() + active_config = manager.active_db_config.copy() + + success, source = await manager.configure_for_token(token) + + assert success is True + assert source == "token-bound" + get_connection.assert_awaited_once() + connection.execute.assert_awaited_once_with( + "SELECT 1 AS connection_check", + mask_result=False, + max_rows=1, + max_bytes=256, + ) + release_connection.assert_awaited_once_with(token, connection) + assert manager.original_db_config == original_config + assert manager.active_db_config == active_config