This is an automated email from the ASF dual-hosted git repository. FreeOnePlus pushed a commit to branch agent/sec-015-shared-auth-state in repository https://gitbox.apache.org/repos/asf/doris-mcp-server.git
commit e8ec7876491a5ff0a0f55bf80337a30d76e84e10 Author: FreeOnePlus <[email protected]> AuthorDate: Wed Jul 29 23:35:07 2026 +0800 fix: share static token state across workers --- .env.example | 3 + .gitignore | 2 + README.md | 36 ++- doris_mcp_server/auth/token_manager.py | 534 ++++++++++++++++++++++--------- pyproject.toml | 1 + test/security/test_shared_token_state.py | 296 +++++++++++++++++ tokens.json | 2 + uv.lock | 2 + 8 files changed, 716 insertions(+), 160 deletions(-) diff --git a/.env.example b/.env.example index 900155f..61e2508 100644 --- a/.env.example +++ b/.env.example @@ -87,6 +87,9 @@ TOKEN_HASH_ALGORITHM=sha256 # Managed tokens.json writes persist only a self-describing token digest and # return plaintext only at creation time. Legacy plaintext files are migrated # to digest-only version 2.0 on first successful load. +# Local worker processes serialize updates through tokens.json.lock and check +# shared state before authentication. Revocation digests also disable matching +# TOKEN_<ID> environment credentials across every worker. # =================================================================== # Token Management Security Configuration (NEW in v0.6.0) - CRITICAL SECURITY SETTINGS diff --git a/.gitignore b/.gitignore index 92aa8a1..e25cc12 100644 --- a/.gitignore +++ b/.gitignore @@ -20,4 +20,6 @@ env.bak/ venv.bak/ .idea/ .coverage +.coverage.* coverage.xml +tokens.json.lock diff --git a/README.md b/README.md index 003ec03..d06e8a0 100644 --- a/README.md +++ b/README.md @@ -174,6 +174,8 @@ field and do not persist or propagate the token in URLs. - Complete audit logging - **Digest-only persistence** to `tokens.json`; plaintext is returned once when a token is created + - **Multi-worker consistency** through a process-shared lock, atomic + read-modify-write updates, and digest-only revocation records > **🔐 Security Note**: The interface is designed for localhost administration > only. It cannot be accessed remotely, ensuring maximum security for token > management operations. @@ -1860,16 +1862,27 @@ Use `WORKERS=1` with Doris-backed OAuth. `WORKERS=0` expands to CPU count and fa **A:** The hot reload system is designed for enterprise production environments with comprehensive safety measures: **How It Works:** -- **File Monitoring**: Checks tokens.json every 10 seconds for modifications -- **Immediate Validation**: New tokens are validated including database connectivity -- **Atomic Updates**: All-or-nothing configuration updates -- **Rollback Protection**: Automatic rollback if any token validation fails +- **Request-time synchronization**: Every token lookup compares the shared + file signature, so another local worker's create or revoke is observed on + the next authenticated request rather than waiting for the polling interval +- **Background monitoring**: A 10-second monitor still refreshes idle workers +- **Serialized updates**: `tokens.json.lock` protects every managed + read-modify-write operation across local worker processes +- **Atomic updates**: A same-directory temporary file is flushed and replaced + atomically with owner-only permissions +- **Rollback protection**: Invalid externally edited state does not partially + replace a worker's current in-memory view +- **Shared revocation**: `revoked_tokens` stores only token digests and also + disables matching `TOKEN_<ID>` environment credentials in every worker **Safety Features:** -- **Backup and Restore**: Current configuration backed up before changes -- **Connection Testing**: Database connections tested before applying changes -- **Error Isolation**: Invalid tokens don't affect existing valid tokens -- **Audit Logging**: Complete trail of all configuration changes +- **No lost updates**: Concurrent create/revoke operations reload the latest + document while holding the process-shared lock +- **No bearer-token plaintext persistence**: Live and revoked bearer values + are represented only by self-describing digests +- **Owner-only state files**: Managed state and lock files use mode `0600` +- **Error isolation**: Invalid state is rejected before it can replace the + complete local token map **Best Practices:** ```bash @@ -1907,6 +1920,13 @@ environment secret or generate a bearer/digest pair offline as shown above. Manual `tokens.json` entries must use `token_digest`; plaintext `token` entries exist only for one-way migration from version 1. +The file backend coordinates multiple worker processes on one host. Every +worker must use the same `TOKEN_FILE_PATH`, and the underlying filesystem must +provide reliable file locking and atomic rename semantics. Multiple hosts or +containers without a shared locking filesystem require an external +transactional state backend; copying separate `tokens.json` files does not +provide cluster-wide revocation. + **Administrative Endpoints (Secure, Local Access Only):** 🛡️ **SECURITY**: These endpoints are protected by comprehensive security controls and are **disabled by default**. diff --git a/doris_mcp_server/auth/token_manager.py b/doris_mcp_server/auth/token_manager.py index 32f55ec..2416377 100644 --- a/doris_mcp_server/auth/token_manager.py +++ b/doris_mcp_server/auth/token_manager.py @@ -27,10 +27,13 @@ import json import os import secrets import tempfile +from contextlib import contextmanager from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone -from typing import Dict, List, Optional, Any from pathlib import Path +from typing import Any, Dict, Iterator, List, Optional + +from filelock import FileLock from ..utils.logger import get_logger from ..utils.secret_policy import ( @@ -97,9 +100,19 @@ class TokenManager: self._tokens: Dict[str, TokenInfo] = {} # token_hash -> TokenInfo self._token_ids: Dict[str, str] = {} # token_id -> token_hash self._digest_algorithms: set[str] = set() + self._revoked_tokens: Dict[str, Dict[str, str]] = {} # Configuration self.token_file_path = getattr(config.security, 'token_file_path', 'tokens.json') + self._token_file = Path(self.token_file_path) + self._token_lock_file = self._token_file.with_name( + f"{self._token_file.name}.lock" + ) + self._token_file_lock = FileLock( + str(self._token_lock_file), + timeout=30, + mode=0o600, + ) self.enable_token_expiry = getattr(config.security, 'enable_token_expiry', True) self.default_token_expiry_hours = getattr(config.security, 'default_token_expiry_hours', 24 * 30) # 30 days self.token_hash_algorithm = normalize_token_hash_algorithm( @@ -110,6 +123,7 @@ class TokenManager: self.enable_hot_reload = True self.hot_reload_interval = 10 # Check every 10 seconds self._file_last_modified = 0.0 + self._file_signature: Optional[tuple[int, int, int, int]] = None self._hot_reload_task: Optional[asyncio.Task[None]] = None # Load tokens from configuration @@ -142,6 +156,70 @@ class TokenManager: f"'{RESERVED_DORIS_OAUTH_TOKEN_PREFIX}'" ) + @contextmanager + def _shared_state_lock(self) -> Iterator[None]: + """Serialize read-modify-write operations across worker processes.""" + with self._token_file_lock: + os.chmod(self._token_lock_file, 0o600) + yield + + def _current_file_signature(self) -> Optional[tuple[int, int, int, int]]: + """Return a signature that also changes after an atomic replacement.""" + try: + file_stat = self._token_file.stat() + except FileNotFoundError: + return None + return ( + file_stat.st_dev, + file_stat.st_ino, + file_stat.st_size, + file_stat.st_mtime_ns, + ) + + def _reload_tokens_from_sources(self) -> None: + """Replace local state from the environment and shared token file.""" + previous_state = ( + self._tokens, + self._token_ids, + self._digest_algorithms, + self._revoked_tokens, + self._file_signature, + self._file_last_modified, + ) + self._tokens = {} + self._token_ids = {} + self._digest_algorithms = set() + self._revoked_tokens = {} + try: + self._load_tokens_from_env() + if self._token_file.exists(): + self._load_tokens_from_file() + self._update_file_modified_time() + except Exception: + ( + self._tokens, + self._token_ids, + self._digest_algorithms, + self._revoked_tokens, + self._file_signature, + self._file_last_modified, + ) = previous_state + raise + + def _synchronize_shared_state(self) -> bool: + """Reload immediately when another worker atomically changes the file.""" + if self._current_file_signature() == self._file_signature: + return False + with self._shared_state_lock(): + if self._current_file_signature() == self._file_signature: + return False + self._reload_tokens_from_sources() + self.logger.info( + "Synchronized %s static tokens from shared state", + len(self._tokens), + ) + return True + @staticmethod def _parse_datetime(value: Any, *, setting: str) -> Optional[datetime]: """Parse an RFC 3339 timestamp into the manager's naive UTC form.""" @@ -247,6 +325,15 @@ class TokenManager: token_hash = self._hash_token(raw_token) self._digest_algorithms.add(token_hash.partition(":")[0]) + # A later source wins cleanly when the same token ID is rotated. + previous_hash = self._token_ids.get(token_info.token_id) + if ( + previous_hash + and previous_hash != token_hash + and previous_hash in self._tokens + ): + del self._tokens[previous_hash] + # Store token self._tokens[token_hash] = token_info self._token_ids[token_info.token_id] = token_hash @@ -261,13 +348,12 @@ class TokenManager: def _load_tokens(self) -> None: """Load tokens from configuration sources""" - # 1. Load from environment variables - self._load_tokens_from_env() - - # 2. Load from token file if exists - if os.path.exists(self.token_file_path): - self._load_tokens_from_file() - + if self._token_file.exists(): + with self._shared_state_lock(): + self._reload_tokens_from_sources() + else: + self._reload_tokens_from_sources() + self.logger.info(f"Token loading completed, total tokens: {len(self._tokens)}") def _load_tokens_from_env(self) -> None: @@ -333,8 +419,18 @@ class TokenManager: return if not isinstance(tokens_list, list): raise ValueError("Static token file must contain a tokens array") + revoked_tokens = ( + tokens_data.get("revoked_tokens", []) + if isinstance(tokens_data, dict) + else [] + ) + if not isinstance(revoked_tokens, list): + raise ValueError( + "Static token file revoked_tokens must be an array" + ) persisted_tokens = [] + persisted_revocations = [] needs_migration = ( not isinstance(tokens_data, dict) or tokens_data.get("version") != "2.0" @@ -350,8 +446,28 @@ class TokenManager: ) needs_migration = needs_migration or had_raw_token + for revoked_token in revoked_tokens: + sanitized = self._sanitize_revoked_token(revoked_token) + token_hash = sanitized["token_digest"] + self._revoked_tokens[token_hash] = sanitized + self._digest_algorithms.add(token_hash.partition(":")[0]) + persisted_revocations.append(sanitized) + + revoked_token_info = self._tokens.get(token_hash) + if revoked_token_info is not None: + del self._tokens[token_hash] + if ( + revoked_token_info is not None + and self._token_ids.get(revoked_token_info.token_id) + == token_hash + ): + self._token_ids.pop(revoked_token_info.token_id, None) + if needs_migration: - migrated_data = self._token_file_document(persisted_tokens) + migrated_data = self._token_file_document( + persisted_tokens, + persisted_revocations, + ) try: self._atomic_write_token_file( Path(self.token_file_path), @@ -376,13 +492,29 @@ class TokenManager: def _hash_token(self, token: str, algorithm: Optional[str] = None) -> str: """Hash token for secure storage""" - return build_token_digest(token, algorithm or self.token_hash_algorithm) + return str( + build_token_digest( + token, + algorithm or self.token_hash_algorithm, + ) + ) + + def _candidate_token_digests(self, token: str) -> List[str]: + algorithms = self._digest_algorithms | {self.token_hash_algorithm} + return [ + self._hash_token(token, algorithm) + for algorithm in sorted(algorithms) + ] def _lookup_token(self, token: str) -> tuple[str, Optional[TokenInfo]]: """Look up a raw token against every digest algorithm in the store.""" - algorithms = self._digest_algorithms or {self.token_hash_algorithm} - for algorithm in sorted(algorithms): - token_hash = self._hash_token(token, algorithm) + candidate_digests = self._candidate_token_digests(token) + if any( + token_hash in self._revoked_tokens + for token_hash in candidate_digests + ): + return "", None + for token_hash in candidate_digests: token_info = self._tokens.get(token_hash) if token_info is not None: return token_hash, token_info @@ -391,6 +523,8 @@ class TokenManager: async def validate_token(self, token: str) -> TokenValidationResult: """Validate token and return user information""" try: + self._synchronize_shared_state() + # Find token info _token_hash, token_info = self._lookup_token(token) if not token_info: @@ -442,10 +576,6 @@ class TokenManager: ) -> str: """Create a new token""" try: - # Check if token_id already exists - if token_id in self._token_ids: - raise ValueError(f"Token ID '{token_id}' already exists") - # Generate or use provided token if custom_token: raw_token = custom_token @@ -471,27 +601,35 @@ class TokenManager: description=description, database_config=database_config ) - - # Hash and store token - token_hash = self._hash_token(raw_token) - if token_hash in self._tokens: - raise ValueError("Token value already exists") - self._tokens[token_hash] = token_info - self._token_ids[token_id] = token_hash - self._digest_algorithms.add(self.token_hash_algorithm) - - self.logger.info(f"Created new token '{token_id}'") - - # Save token to file - try: + + with self._shared_state_lock(): + self._reload_tokens_from_sources() + if token_id in self._token_ids: + raise ValueError( + f"Token ID '{token_id}' already exists" + ) + + candidate_digests = self._candidate_token_digests(raw_token) + if any( + digest in self._revoked_tokens + for digest in candidate_digests + ): + raise ValueError( + "A revoked token value cannot be reused" + ) + if any( + digest in self._tokens + for digest in candidate_digests + ): + raise ValueError("Token value already exists") + + token_hash = self._hash_token(raw_token) self._save_token_to_file(token_id, token_hash, token_info) - except Exception: - self._tokens.pop(token_hash, None) - self._token_ids.pop(token_id, None) - self._digest_algorithms = { - digest.partition(":")[0] for digest in self._tokens - } - raise + self._tokens[token_hash] = token_info + self._token_ids[token_id] = token_hash + self._digest_algorithms.add(self.token_hash_algorithm) + + self.logger.info(f"Created new token '{token_id}'") return raw_token @@ -502,18 +640,25 @@ class TokenManager: async def revoke_token(self, token_id: str) -> bool: """Revoke a token by token ID""" try: - if token_id not in self._token_ids: - self.logger.warning(f"Token ID '{token_id}' not found") - return False - - # Persist the revocation before changing live state. Otherwise a - # failed file write would let the token reappear after restart. - token_hash = self._token_ids[token_id] - self._remove_token_from_file(token_id) - - if token_hash in self._tokens: - del self._tokens[token_hash] - del self._token_ids[token_id] + with self._shared_state_lock(): + self._reload_tokens_from_sources() + if token_id not in self._token_ids: + self.logger.warning( + f"Token ID '{token_id}' not found" + ) + return False + + # Persist before changing live state. The shared revocation + # record also disables environment-provisioned credentials. + token_hash = self._token_ids[token_id] + revocation = self._remove_token_from_file( + token_id, + token_hash, + ) + self._tokens.pop(token_hash, None) + self._token_ids.pop(token_id, None) + self._revoked_tokens[token_hash] = revocation + self._digest_algorithms.add(token_hash.partition(":")[0]) self.logger.info(f"Revoked token '{token_id}'") return True @@ -525,15 +670,20 @@ class TokenManager: def _save_tokens_to_file(self) -> None: """Save current tokens to JSON file""" try: - tokens_list = [ - self._token_info_to_config(token_hash, token_info) - for token_hash, token_info in self._tokens.items() - ] - file_content = self._token_file_document(tokens_list) - self._atomic_write_token_file( - Path(self.token_file_path), - file_content, - ) + with self._shared_state_lock(): + self._reload_tokens_from_sources() + tokens_list = [ + self._token_info_to_config(token_hash, token_info) + for token_hash, token_info in self._tokens.items() + ] + file_content = self._token_file_document( + tokens_list, + list(self._revoked_tokens.values()), + ) + self._atomic_write_token_file( + Path(self.token_file_path), + file_content, + ) self.logger.info(f"Saved {len(tokens_list)} tokens to file: {self.token_file_path}") except Exception as e: @@ -548,7 +698,7 @@ class TokenManager: """Save one new token as a digest-only record.""" try: # Load existing file - existing_data: Dict[str, Any] = {"tokens": []} + existing_data: Any = {"tokens": []} if os.path.exists(self.token_file_path): try: with open(self.token_file_path, 'r', encoding='utf-8') as f: @@ -557,10 +707,19 @@ class TokenManager: raise ValueError( f"Could not load existing token file: {e}" ) from e + if isinstance(existing_data, list): + existing_data = {"tokens": existing_data} + if not isinstance(existing_data, dict): + raise ValueError("Static token file must contain an object") # Ensure tokens list exists if 'tokens' not in existing_data or not isinstance(existing_data['tokens'], list): existing_data['tokens'] = [] + revoked_tokens = existing_data.get("revoked_tokens", []) + if not isinstance(revoked_tokens, list): + raise ValueError( + "Static token file revoked_tokens must be an array" + ) # Check if token already exists in file token_exists = False @@ -586,9 +745,16 @@ class TokenManager: self._sanitize_persisted_token(token_config) for token_config in existing_data['tokens'] ] + revoked_tokens = [ + self._sanitize_revoked_token(revoked_token) + for revoked_token in revoked_tokens + ] # Update metadata - existing_data = self._token_file_document(existing_data['tokens']) + existing_data = self._token_file_document( + existing_data['tokens'], + revoked_tokens, + ) self._atomic_write_token_file( Path(self.token_file_path), existing_data, @@ -599,16 +765,30 @@ class TokenManager: self.logger.error(f"Failed to save token '{token_id}' to file: {e}") raise - def _token_file_document(self, tokens: List[Dict[str, Any]]) -> Dict[str, Any]: + def _token_file_document( + self, + tokens: List[Dict[str, Any]], + revoked_tokens: Optional[List[Dict[str, str]]] = None, + ) -> Dict[str, Any]: """Build the versioned digest-only token file document.""" return { "version": "2.0", - "description": "Doris MCP Server digest-only token configuration file", + "description": ( + "Doris MCP Server digest-only shared token state" + ), "updated_at": datetime.utcnow().isoformat() + "Z", "tokens": tokens, + "revoked_tokens": sorted( + revoked_tokens or [], + key=lambda entry: ( + entry["revoked_at"], + entry["token_id"], + ), + ), "notes": [ "Bearer token plaintext is returned only when a token is created.", "token_digest is self-describing and may use sha256 or sha512.", + "revoked_tokens is shared by every worker and prevents token reuse.", "Do not replace token_digest with a plaintext token.", ], } @@ -691,6 +871,39 @@ class TokenManager: ) return sanitized + def _sanitize_revoked_token( + self, + revoked_token: Dict[str, Any], + ) -> Dict[str, str]: + """Validate and whitelist a shared revocation record.""" + if not isinstance(revoked_token, dict): + raise ValueError("Static token revocation entries must be objects") + token_id = str(revoked_token.get("token_id") or "").strip() + if not token_id: + raise ValueError("Static token revocation requires token_id") + raw_digest = revoked_token.get("token_digest") + if not isinstance(raw_digest, str): + raise ValueError( + f"revoked static token '{token_id}' token_digest is required" + ) + token_digest = normalize_token_digest( + raw_digest, + setting=f"revoked static token '{token_id}' token_digest", + ) + revoked_at = self._parse_datetime( + revoked_token.get("revoked_at"), + setting=f"revoked static token '{token_id}' revoked_at", + ) + if revoked_at is None: + raise ValueError( + f"revoked static token '{token_id}' revoked_at is required" + ) + return { + "token_id": token_id, + "token_digest": token_digest, + "revoked_at": revoked_at.isoformat() + "Z", + } + def _token_info_to_config( self, token_hash: str, @@ -732,35 +945,72 @@ class TokenManager: return token_config - def _remove_token_from_file(self, token_id: str) -> None: - """Remove a token from the JSON file""" + def _remove_token_from_file( + self, + token_id: str, + token_hash: str, + ) -> Dict[str, str]: + """Remove a live record and persist a digest-only revocation.""" try: - if not os.path.exists(self.token_file_path): - return - - # Load existing file - with open(self.token_file_path, 'r', encoding='utf-8') as f: - existing_data = json.load(f) - - if 'tokens' not in existing_data or not isinstance(existing_data['tokens'], list): - return - - # Remove the token - original_count = len(existing_data['tokens']) - existing_data['tokens'] = [ - self._sanitize_persisted_token(token) - for token in existing_data['tokens'] - if token.get('token_id') != token_id - ] - - if len(existing_data['tokens']) < original_count: - # Update metadata - existing_data = self._token_file_document(existing_data['tokens']) - self._atomic_write_token_file( - Path(self.token_file_path), - existing_data, + existing_data: Any = {"tokens": [], "revoked_tokens": []} + if self._token_file.exists(): + with self._token_file.open("r", encoding="utf-8") as token_file: + existing_data = json.load(token_file) + if isinstance(existing_data, list): + existing_data = {"tokens": existing_data} + if not isinstance(existing_data, dict): + raise ValueError("Static token file must contain an object") + + tokens = existing_data.get("tokens", []) + revoked_tokens = existing_data.get("revoked_tokens", []) + if not isinstance(tokens, list): + raise ValueError( + "Static token file must contain a tokens array" ) - self.logger.info(f"Removed token '{token_id}' from file: {self.token_file_path}") + if not isinstance(revoked_tokens, list): + raise ValueError( + "Static token file revoked_tokens must be an array" + ) + + sanitized_tokens = [] + for token in tokens: + sanitized = self._sanitize_persisted_token(token) + if ( + sanitized["token_id"] != token_id + and sanitized["token_digest"] != token_hash + ): + sanitized_tokens.append(sanitized) + + sanitized_revocations = [ + self._sanitize_revoked_token(revoked_token) + for revoked_token in revoked_tokens + ] + revocation = self._sanitize_revoked_token( + { + "token_id": token_id, + "token_digest": token_hash, + "revoked_at": datetime.utcnow().isoformat() + "Z", + } + ) + sanitized_revocations = [ + item + for item in sanitized_revocations + if item["token_digest"] != token_hash + ] + sanitized_revocations.append(revocation) + + self._atomic_write_token_file( + Path(self.token_file_path), + self._token_file_document( + sanitized_tokens, + sanitized_revocations, + ), + ) + self.logger.info( + f"Persisted revocation for '{token_id}' in " + f"{self.token_file_path}" + ) + return revocation except Exception as e: self.logger.error(f"Failed to remove token '{token_id}' from file: {e}") @@ -768,6 +1018,7 @@ class TokenManager: async def list_tokens(self) -> List[Dict[str, Any]]: """List all tokens (without sensitive data)""" + self._synchronize_shared_state() tokens = [] for token_hash, token_info in self._tokens.items(): @@ -804,20 +1055,24 @@ class TokenManager: """Remove expired tokens and return count""" if not self.enable_token_expiry: return 0 - - now = datetime.utcnow() - expired_tokens = [] - - # Find expired tokens - for token_hash, token_info in self._tokens.items(): - if token_info.expires_at and now > token_info.expires_at: - expired_tokens.append((token_hash, token_info.token_id)) - - # Remove expired tokens - for token_hash, token_id in expired_tokens: - del self._tokens[token_hash] - if token_id in self._token_ids: - del self._token_ids[token_id] + + with self._shared_state_lock(): + self._reload_tokens_from_sources() + now = datetime.utcnow() + expired_tokens = [ + (token_hash, token_info.token_id) + for token_hash, token_info in self._tokens.items() + if token_info.expires_at and now > token_info.expires_at + ] + + for token_hash, token_id in expired_tokens: + revocation = self._remove_token_from_file( + token_id, + token_hash, + ) + self._tokens.pop(token_hash, None) + self._token_ids.pop(token_id, None) + self._revoked_tokens[token_hash] = revocation if expired_tokens: self.logger.info(f"Cleaned up {len(expired_tokens)} expired tokens") @@ -828,14 +1083,19 @@ class TokenManager: """Save current tokens to JSON file""" try: target_path = Path(file_path or self.token_file_path) - tokens_list = [ - self._token_info_to_config(token_hash, token_info) - for token_hash, token_info in self._tokens.items() - ] - self._atomic_write_token_file( - target_path, - self._token_file_document(tokens_list), - ) + with self._shared_state_lock(): + self._reload_tokens_from_sources() + tokens_list = [ + self._token_info_to_config(token_hash, token_info) + for token_hash, token_info in self._tokens.items() + ] + self._atomic_write_token_file( + target_path, + self._token_file_document( + tokens_list, + list(self._revoked_tokens.values()), + ), + ) self.logger.info(f"Saved {len(tokens_list)} tokens to file: {target_path}") return True @@ -853,6 +1113,7 @@ class TokenManager: DatabaseConfig if token exists and has database binding, None otherwise """ try: + self._synchronize_shared_state() _token_hash, token_info = self._lookup_token(token) if not token_info or not token_info.is_active: @@ -870,6 +1131,7 @@ class TokenManager: def get_token_stats(self) -> Dict[str, Any]: """Get token statistics""" + self._synchronize_shared_state() now = datetime.utcnow() total_tokens = len(self._tokens) active_tokens = sum(1 for info in self._tokens.values() if info.is_active) @@ -893,10 +1155,9 @@ class TokenManager: """Start hot reload monitoring task""" if self._hot_reload_task: return # Already running - - # Update initial file modification time - self._update_file_modified_time() - + + # _load_tokens() already captured a signature for the exact state held + # in memory. Re-statting here could hide a change made in between. # Start monitoring task self._hot_reload_task = asyncio.create_task(self._hot_reload_monitor()) self.logger.info(f"Started hot reload monitoring for {self.token_file_path}") @@ -911,8 +1172,11 @@ class TokenManager: def _update_file_modified_time(self) -> None: """Update the last modified time of tokens file""" try: - if os.path.exists(self.token_file_path): - self._file_last_modified = os.path.getmtime(self.token_file_path) + self._file_signature = self._current_file_signature() + if self._file_signature is None: + self._file_last_modified = 0.0 + else: + self._file_last_modified = self._token_file.stat().st_mtime except Exception as e: self.logger.debug(f"Failed to get file modification time: {e}") @@ -921,41 +1185,7 @@ class TokenManager: while True: try: await asyncio.sleep(self.hot_reload_interval) - - if not os.path.exists(self.token_file_path): - continue - - # Check if file was modified - current_mtime = os.path.getmtime(self.token_file_path) - if current_mtime > self._file_last_modified: - self.logger.info(f"Detected changes in {self.token_file_path}, reloading tokens...") - - try: - # Backup current tokens - old_tokens = self._tokens.copy() - old_token_ids = self._token_ids.copy() - old_digest_algorithms = self._digest_algorithms.copy() - - # Clear and reload - self._tokens.clear() - self._token_ids.clear() - self._digest_algorithms.clear() - - # Environment credentials remain effective across file reloads. - self._load_tokens_from_env() - self._load_tokens_from_file() - - # Update modification time - self._update_file_modified_time() - - self.logger.info(f"Hot reload completed, {len(self._tokens)} tokens loaded") - - except Exception as reload_error: - # Restore backup on failure - self.logger.error(f"Hot reload failed, restoring previous tokens: {reload_error}") - self._tokens = old_tokens - self._token_ids = old_token_ids - self._digest_algorithms = old_digest_algorithms + self._synchronize_shared_state() except asyncio.CancelledError: self.logger.info("Hot reload monitor stopped") diff --git a/pyproject.toml b/pyproject.toml index fbe49a8..646c82f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,7 @@ dependencies = [ "aiofiles>=23.0.0", "aiohttp>=3.9.0", "aioredis>=2.0.0", + "filelock>=3.16.1,<4.0.0", # Data processing "pandas>=2.0.0", "numpy>=1.24.0", diff --git a/test/security/test_shared_token_state.py b/test/security/test_shared_token_state.py new file mode 100644 index 0000000..cdf780e --- /dev/null +++ b/test/security/test_shared_token_state.py @@ -0,0 +1,296 @@ +# 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. + +import asyncio +import json +import multiprocessing +import os +import secrets +import stat +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path + +import pytest + +from doris_mcp_server.auth.token_manager import TokenManager +from doris_mcp_server.utils.config import DorisConfig +from doris_mcp_server.utils.secret_policy import ( + is_static_token_environment_variable, +) + + +def _config( + token_path: str | Path, + *, + algorithm: str = "sha256", +) -> DorisConfig: + config = DorisConfig() + config.security.token_file_path = str(token_path) + config.security.token_hash_algorithm = algorithm + return config + + +def _clear_static_token_environment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + for name in list(os.environ): + if is_static_token_environment_variable(name): + monkeypatch.delenv(name, raising=False) + + +async def _create_token_in_worker( + token_path: str, + token_id: str, + raw_token: str, +) -> str: + manager = TokenManager(_config(token_path)) + try: + await manager.create_token( + token_id, + custom_token=raw_token, + ) + return "created" + except ValueError as exc: + return str(exc) + finally: + manager.stop_hot_reload() + + +def _process_create_token( + token_path: str, + token_id: str, + raw_token: str, +) -> str: + return asyncio.run(_create_token_in_worker(token_path, token_id, raw_token)) + + +async def _revoke_token_in_worker( + token_path: str, + token_id: str, +) -> bool: + manager = TokenManager(_config(token_path)) + try: + return await manager.revoke_token(token_id) + finally: + manager.stop_hot_reload() + + +def _process_revoke_token(token_path: str, token_id: str) -> bool: + return asyncio.run(_revoke_token_in_worker(token_path, token_id)) + + [email protected] +async def test_workers_observe_create_and_revoke_on_the_next_access( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _clear_static_token_environment(monkeypatch) + token_path = tmp_path / "tokens.json" + first_worker = TokenManager(_config(token_path, algorithm="sha512")) + second_worker = TokenManager(_config(token_path)) + first_worker.hot_reload_interval = 3600 + second_worker.hot_reload_interval = 3600 + try: + raw_token = await first_worker.create_token("shared-service") + + observed = await second_worker.validate_token(raw_token) + assert observed.is_valid is True + assert observed.token_info is not None + assert observed.token_info.token_id == "shared-service" + + assert await second_worker.revoke_token("shared-service") is True + rejected = await first_worker.validate_token(raw_token) + assert rejected.is_valid is False + assert rejected.error_message == "Invalid token" + + shared_state = json.loads(token_path.read_text(encoding="utf-8")) + assert shared_state["tokens"] == [] + assert len(shared_state["revoked_tokens"]) == 1 + assert raw_token not in token_path.read_text(encoding="utf-8") + assert stat.S_IMODE(token_path.stat().st_mode) == 0o600 + assert stat.S_IMODE(Path(f"{token_path}.lock").stat().st_mode) == 0o600 + finally: + first_worker.stop_hot_reload() + second_worker.stop_hot_reload() + + [email protected] +async def test_environment_token_revocation_is_shared_and_persistent( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _clear_static_token_environment(monkeypatch) + raw_token = secrets.token_urlsafe(40) + monkeypatch.setenv("TOKEN_SHARED_ENV", raw_token) + token_path = tmp_path / "tokens.json" + first_worker = TokenManager(_config(token_path, algorithm="sha512")) + second_worker = TokenManager(_config(token_path)) + try: + assert (await second_worker.validate_token(raw_token)).is_valid + assert await first_worker.revoke_token("shared_env") is True + assert (await second_worker.validate_token(raw_token)).is_valid is False + + restarted_worker = TokenManager(_config(token_path)) + try: + assert (await restarted_worker.validate_token(raw_token)).is_valid is False + with pytest.raises( + ValueError, + match="revoked token value", + ): + await restarted_worker.create_token( + "replayed", + custom_token=raw_token, + ) + finally: + restarted_worker.stop_hot_reload() + + stored_text = token_path.read_text(encoding="utf-8") + stored = json.loads(stored_text) + assert stored["tokens"] == [] + assert stored["revoked_tokens"][0]["token_id"] == "shared_env" + assert raw_token not in stored_text + finally: + first_worker.stop_hot_reload() + second_worker.stop_hot_reload() + + [email protected] +async def test_concurrent_worker_creates_do_not_lose_updates( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _clear_static_token_environment(monkeypatch) + token_path = tmp_path / "tokens.json" + entries = [(f"worker-{index}", secrets.token_urlsafe(40)) for index in range(8)] + context = multiprocessing.get_context("spawn") + + with ProcessPoolExecutor( + max_workers=4, + mp_context=context, + ) as executor: + futures = [ + executor.submit( + _process_create_token, + str(token_path), + token_id, + raw_token, + ) + for token_id, raw_token in entries + ] + assert [future.result(timeout=30) for future in futures] == ["created"] * len( + entries + ) + + manager = TokenManager(_config(token_path)) + try: + listed = await manager.list_tokens() + assert {item["token_id"] for item in listed} == { + token_id for token_id, _raw_token in entries + } + for _token_id, raw_token in entries: + assert (await manager.validate_token(raw_token)).is_valid + finally: + manager.stop_hot_reload() + + stored_text = token_path.read_text(encoding="utf-8") + stored = json.loads(stored_text) + assert len(stored["tokens"]) == len(entries) + assert all(raw_token not in stored_text for _, raw_token in entries) + + [email protected] +async def test_concurrent_create_and_revoke_preserve_both_updates( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _clear_static_token_environment(monkeypatch) + token_path = tmp_path / "tokens.json" + old_token = secrets.token_urlsafe(40) + new_token = secrets.token_urlsafe(40) + manager = TokenManager(_config(token_path)) + try: + await manager.create_token( + "old-service", + custom_token=old_token, + ) + finally: + manager.stop_hot_reload() + + context = multiprocessing.get_context("spawn") + with ProcessPoolExecutor( + max_workers=2, + mp_context=context, + ) as executor: + create_future = executor.submit( + _process_create_token, + str(token_path), + "new-service", + new_token, + ) + revoke_future = executor.submit( + _process_revoke_token, + str(token_path), + "old-service", + ) + assert create_future.result(timeout=30) == "created" + assert revoke_future.result(timeout=30) is True + + verifier = TokenManager(_config(token_path)) + try: + assert (await verifier.validate_token(old_token)).is_valid is False + assert (await verifier.validate_token(new_token)).is_valid is True + finally: + verifier.stop_hot_reload() + + stored = json.loads(token_path.read_text(encoding="utf-8")) + assert [token["token_id"] for token in stored["tokens"]] == ["new-service"] + assert [revoked["token_id"] for revoked in stored["revoked_tokens"]] == [ + "old-service" + ] + + [email protected] +async def test_concurrent_duplicate_token_id_has_one_winner( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _clear_static_token_environment(monkeypatch) + token_path = tmp_path / "tokens.json" + entries = [ + ("same-id", secrets.token_urlsafe(40)), + ("same-id", secrets.token_urlsafe(40)), + ] + context = multiprocessing.get_context("spawn") + + with ProcessPoolExecutor( + max_workers=2, + mp_context=context, + ) as executor: + futures = [ + executor.submit( + _process_create_token, + str(token_path), + token_id, + raw_token, + ) + for token_id, raw_token in entries + ] + results = [future.result(timeout=30) for future in futures] + + assert results.count("created") == 1 + assert sum("already exists" in result for result in results) == 1 + stored = json.loads(token_path.read_text(encoding="utf-8")) + assert [token["token_id"] for token in stored["tokens"]] == ["same-id"] diff --git a/tokens.json b/tokens.json index 1ae66dd..da03620 100644 --- a/tokens.json +++ b/tokens.json @@ -2,10 +2,12 @@ "version": "2.0", "description": "Credential-free Doris MCP Server digest-only token configuration template", "tokens": [], + "revoked_tokens": [], "notes": [ "No usable credential is shipped with the server.", "Create tokens through the protected management API or store a self-describing token_digest generated from a high-entropy bearer token.", "Bearer token plaintext is returned only when a token is created and must not be written into this file.", + "Managed revocations are stored as digests and shared by all local worker processes.", "Do not commit deployment credentials or usable token digests to source control." ] } diff --git a/uv.lock b/uv.lock index 8df5bde..c0b093f 100644 --- a/uv.lock +++ b/uv.lock @@ -605,6 +605,7 @@ dependencies = [ { name = "cryptography" }, { name = "fastapi", version = "0.115.12", source = { registry = "https://pypi.org/simple"; }, marker = "python_full_version < '3.14'" }, { name = "fastapi", version = "0.140.13", source = { registry = "https://pypi.org/simple"; }, marker = "python_full_version >= '3.14'" }, + { name = "filelock" }, { name = "httpx" }, { name = "mcp" }, { name = "numpy" }, @@ -698,6 +699,7 @@ requires-dist = [ { name = "click", specifier = ">=8.1.0" }, { name = "cryptography", specifier = ">=41.0.0" }, { name = "fastapi", specifier = ">=0.108.0" }, + { name = "filelock", specifier = ">=3.16.1,<4.0.0" }, { name = "flake8", marker = "extra == 'dev'", specifier = ">=7.0.0" }, { name = "grafana-client", marker = "extra == 'monitoring'", specifier = ">=3.5.0" }, { name = "httpx", specifier = ">=0.26.0" }, --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
