Production hardening: fix race conditions, retry logic, metrics, and CI pipelines

gensecai-dev · gensecai-dev · commit 3e1e7bec3fbc · 2026-02-25T09:18:06.000+05:30
- Add asyncio.Lock to circuit breaker state transitions and indexer initialization
- Narrow retry scope to transient errors only (5xx, connection/timeout)
- Fix circuit breaker monitoring always returning "unknown" state
- Bound Prometheus endpoint labels to prevent cardinality explosion
- Add JSONDecodeError handling on all response.json() call sites
- Replace non-deterministic hash() cache keys with sorted() representation
- Add SSE keepalive loop cancellation on client disconnect
- Add MAX_BATCH_SIZE=100 limit on batch JSON-RPC requests
- Throttle session cleanup to every 60s instead of every request
- Remove premature REQUEST_COUNT increment (was always counting 200)
- Remove || true from release and security CI workflows
- Add fastmcp to pyproject.toml dependencies
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -44,7 +44,7 @@ jobs:
     - name: Run tests
       run: |
         pip install pytest pytest-asyncio
-        pytest tests/ -v || true
+        pytest tests/ -v
     
     - name: Build package
       run: python -m build
diff --git a/.github/workflows/security.yml b/.github/workflows/security.yml
@@ -35,18 +35,20 @@ jobs:
       run: pip install -e .
     
     - name: Run safety check
+      continue-on-error: true
       run: |
-        safety check --json --output safety-report.json || true
+        safety check --json --output safety-report.json
         if [ -f safety-report.json ]; then
           echo "### Safety Report" >> $GITHUB_STEP_SUMMARY
           echo '```json' >> $GITHUB_STEP_SUMMARY
           cat safety-report.json >> $GITHUB_STEP_SUMMARY
           echo '```' >> $GITHUB_STEP_SUMMARY
         fi
-    
+
     - name: Run pip-audit
+      continue-on-error: true
       run: |
-        pip-audit --format json --output pip-audit-report.json || true
+        pip-audit --format json --output pip-audit-report.json
         if [ -f pip-audit-report.json ]; then
           echo "### Pip Audit Report" >> $GITHUB_STEP_SUMMARY
           echo '```json' >> $GITHUB_STEP_SUMMARY
@@ -76,12 +78,13 @@ jobs:
         python-version: '3.13'
     
     - name: Run Bandit
+      continue-on-error: true
       run: |
         pip install bandit
-        bandit -r src/ -f json -o bandit-report.json || true
+        bandit -r src/ -f json -o bandit-report.json
         echo "### Bandit Security Report" >> $GITHUB_STEP_SUMMARY
         echo '```' >> $GITHUB_STEP_SUMMARY
-        bandit -r src/ -f txt || true
+        bandit -r src/ -f txt
         echo '```' >> $GITHUB_STEP_SUMMARY
     
     - name: Run Semgrep
@@ -158,7 +161,7 @@ jobs:
         # Install gitleaks CLI (free version)
         wget -q https://github.com/gitleaks/gitleaks/releases/download/v8.18.4/gitleaks_8.18.4_linux_x64.tar.gz
         tar -xzf gitleaks_8.18.4_linux_x64.tar.gz
-        ./gitleaks detect --source . --verbose --report-path gitleaks-report.json || true
+        ./gitleaks detect --source . --verbose --report-path gitleaks-report.json
         echo "### Gitleaks Report" >> $GITHUB_STEP_SUMMARY
         if [ -f gitleaks-report.json ]; then
           echo '```json' >> $GITHUB_STEP_SUMMARY
diff --git a/pyproject.toml b/pyproject.toml
@@ -46,6 +46,8 @@ dependencies = [
     # Utilities
     "python-dotenv>=1.0.0",
     "aiofiles>=24.0.0",
+    # MCP Framework
+    "fastmcp>=2.14.0",
 ]
 
 [project.optional-dependencies]
diff --git a/src/wazuh_mcp_server/api/wazuh_client.py b/src/wazuh_mcp_server/api/wazuh_client.py
@@ -84,7 +84,10 @@ async def _authenticate(self):
             response = await self.client.post(auth_url, auth=(self.config.wazuh_user, self.config.wazuh_pass))
             response.raise_for_status()
 
-            data = response.json()
+            try:
+                data = response.json()
+            except (json.JSONDecodeError, ValueError):
+                raise ValueError("Invalid JSON in authentication response from Wazuh API")
             if "data" not in data or "token" not in data["data"]:
                 raise ValueError("Invalid authentication response from Wazuh API")
 
@@ -224,7 +227,7 @@ async def _get_cached(self, cache_key: str, endpoint: str, **kwargs) -> Dict[str
     async def get_rules(self, **params) -> Dict[str, Any]:
         """Get Wazuh detection rules (cached for 5 minutes)."""
         # Use caching for rules as they rarely change
-        cache_key = f"rules:{hash(frozenset(params.items()) if params else 'all')}"
+        cache_key = f"rules:{sorted(params.items()) if params else 'all'}"
         return await self._get_cached(cache_key, "/rules", params=params)
 
     async def get_rule_info(self, rule_id: str) -> Dict[str, Any]:
@@ -234,7 +237,7 @@ async def get_rule_info(self, rule_id: str) -> Dict[str, Any]:
     async def get_decoders(self, **params) -> Dict[str, Any]:
         """Get Wazuh log decoders (cached for 5 minutes)."""
         # Use caching for decoders as they rarely change
-        cache_key = f"decoders:{hash(frozenset(params.items()) if params else 'all')}"
+        cache_key = f"decoders:{sorted(params.items()) if params else 'all'}"
         return await self._get_cached(cache_key, "/decoders", params=params)
 
     async def execute_active_response(self, data: Dict[str, Any]) -> Dict[str, Any]:
@@ -370,7 +373,10 @@ async def _execute_request(self, method: str, endpoint: str, **kwargs) -> Dict[s
             response = await self.client.request(method, url, headers=headers, **kwargs)
             response.raise_for_status()
 
-            data = response.json()
+            try:
+                data = response.json()
+            except (json.JSONDecodeError, ValueError):
+                raise ValueError(f"Invalid JSON response from Wazuh API: {endpoint}")
 
             # Validate response structure
             if "data" not in data:
@@ -391,7 +397,10 @@ async def _execute_request(self, method: str, endpoint: str, **kwargs) -> Dict[s
                 headers = {"Authorization": f"Bearer {self.token}"}
                 response = await self.client.request(method, url, headers=headers, **kwargs)
                 response.raise_for_status()
-                return response.json()
+                try:
+                    return response.json()
+                except (json.JSONDecodeError, ValueError):
+                    raise ValueError(f"Invalid JSON response from Wazuh API after re-auth: {endpoint}")
             else:
                 logger.error(f"Wazuh API request failed: {e.response.status_code} - {e.response.text}")
                 raise ValueError(f"Wazuh API request failed: {e.response.status_code} - {e.response.text}")
diff --git a/src/wazuh_mcp_server/api/wazuh_indexer.py b/src/wazuh_mcp_server/api/wazuh_indexer.py
@@ -7,6 +7,8 @@
   - wazuh-states-vulnerabilities-* — vulnerability data (removed from Manager API in 4.8.0)
 """
 
+import asyncio
+import json
 import logging
 from typing import Any, Dict, Optional
 
@@ -45,6 +47,7 @@ def __init__(
         self.timeout = timeout
         self.client: Optional[httpx.AsyncClient] = None
         self._initialized = False
+        self._init_lock = asyncio.Lock()
 
     @staticmethod
     def _normalize_host(host: str) -> str:
@@ -79,9 +82,12 @@ async def close(self):
             self._initialized = False
 
     async def _ensure_initialized(self):
-        """Ensure client is initialized."""
-        if not self._initialized:
-            await self.initialize()
+        """Ensure client is initialized (thread-safe)."""
+        if self._initialized:
+            return
+        async with self._init_lock:
+            if not self._initialized:
+                await self.initialize()
 
     @retry(
         stop=stop_after_attempt(3),
@@ -109,7 +115,10 @@ async def _search(self, index: str, query: Dict[str, Any], size: int = 100) -> D
         try:
             response = await self.client.post(url, json=body, headers={"Content-Type": "application/json"})
             response.raise_for_status()
-            return response.json()
+            try:
+                return response.json()
+            except (json.JSONDecodeError, ValueError):
+                raise ValueError(f"Invalid JSON response from Wazuh Indexer: {index}")
 
         except httpx.HTTPStatusError as e:
             logger.error(f"Indexer search failed: {e.response.status_code} - {e.response.text}")
@@ -185,7 +194,10 @@ async def get_alerts(
         try:
             response = await self.client.post(url, json=body, headers={"Content-Type": "application/json"})
             response.raise_for_status()
-            result = response.json()
+            try:
+                result = response.json()
+            except (json.JSONDecodeError, ValueError):
+                raise ValueError("Invalid JSON response from Wazuh Indexer alerts query")
         except httpx.HTTPStatusError as e:
             logger.error(f"Alerts query failed: {e.response.status_code} - {e.response.text}")
             raise ValueError(f"Alerts query failed: {e.response.status_code}")
diff --git a/src/wazuh_mcp_server/monitoring.py b/src/wazuh_mcp_server/monitoring.py
@@ -74,6 +74,17 @@ def debug(self, message: str, **extra: Any) -> None:
 # Prometheus metrics registry
 REGISTRY = CollectorRegistry()
 
+# Known endpoints for metric label normalization (prevents unbounded cardinality)
+_KNOWN_ENDPOINTS = {"/", "/mcp", "/health", "/metrics"}
+
+
+def _normalize_endpoint(path: str) -> str:
+    """Collapse unknown paths to 'other' to prevent label explosion."""
+    if path in _KNOWN_ENDPOINTS:
+        return path
+    return "other"
+
+
 # Core metrics
 REQUEST_COUNT = Counter(
     "wazuh_mcp_requests_total", "Total number of requests", ["method", "endpoint", "status_code"], registry=REGISTRY
@@ -495,7 +506,7 @@ async def check_circuit_breaker() -> Dict[str, Any]:
         client = await get_wazuh_client()
         if hasattr(client, "_circuit_breaker"):
             cb = client._circuit_breaker
-            state = cb._state if hasattr(cb, "_state") else "unknown"
+            state = cb.state.value if hasattr(cb, "state") else "unknown"
             failure_count = cb.failure_count if hasattr(cb, "failure_count") else 0
 
             # Update Prometheus metric
@@ -608,8 +619,9 @@ async def monitoring_middleware(request: Request, call_next):
             # Record metrics
             duration = time.time() - start_time
 
-            REQUEST_COUNT.labels(method=method, endpoint=path, status_code=status_code).inc()
-            REQUEST_DURATION.labels(method=method, endpoint=path).observe(duration)
+            normalized = _normalize_endpoint(path)
+            REQUEST_COUNT.labels(method=method, endpoint=normalized, status_code=status_code).inc()
+            REQUEST_DURATION.labels(method=method, endpoint=normalized).observe(duration)
 
             # Record slow requests with correlation ID
             if duration > performance_profiler.slow_threshold:
diff --git a/src/wazuh_mcp_server/resilience.py b/src/wazuh_mcp_server/resilience.py
@@ -14,7 +14,7 @@
 
 import httpx
 from fastapi import HTTPException
-from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
+from tenacity import retry, retry_if_exception, retry_if_exception_type, stop_after_attempt, wait_exponential
 
 logger = logging.getLogger(__name__)
 
@@ -53,6 +53,7 @@ def __init__(self, config: CircuitBreakerConfig):
         self.failure_count = 0
         self.last_failure_time: Optional[float] = None
         self.next_retry_time: Optional[float] = None
+        self._lock = asyncio.Lock()
 
     def __call__(self, func: Callable) -> Callable:
         """Decorator to apply circuit breaker to function."""
@@ -66,19 +67,20 @@ async def wrapper(*args, **kwargs):
     async def _call(self, func: Callable, *args, **kwargs) -> Any:
         """Execute function with circuit breaker logic."""
 
-        # Check if circuit is open
-        if self.state == CircuitBreakerState.OPEN:
-            if self._should_attempt_reset():
-                self.state = CircuitBreakerState.HALF_OPEN
-                logger.info(f"Circuit breaker {func.__name__} moved to HALF_OPEN")
-            else:
-                if self.config.fallback_function:
-                    logger.warning(f"Circuit breaker {func.__name__} OPEN, using fallback")
-                    return await self.config.fallback_function(*args, **kwargs)
+        # Check if circuit is open (under lock to prevent race conditions)
+        async with self._lock:
+            if self.state == CircuitBreakerState.OPEN:
+                if self._should_attempt_reset():
+                    self.state = CircuitBreakerState.HALF_OPEN
+                    logger.info(f"Circuit breaker {func.__name__} moved to HALF_OPEN")
                 else:
-                    raise HTTPException(
-                        status_code=503, detail="Service temporarily unavailable - circuit breaker open"
-                    )
+                    if self.config.fallback_function:
+                        logger.warning(f"Circuit breaker {func.__name__} OPEN, using fallback")
+                        return await self.config.fallback_function(*args, **kwargs)
+                    else:
+                        raise HTTPException(
+                            status_code=503, detail="Service temporarily unavailable - circuit breaker open"
+                        )
 
         try:
             result = await func(*args, **kwargs)
@@ -101,23 +103,24 @@ def _should_attempt_reset(self) -> bool:
 
     async def _on_success(self, func_name: str):
         """Handle successful execution."""
-        if self.state == CircuitBreakerState.HALF_OPEN:
-            self.state = CircuitBreakerState.CLOSED
-            logger.info(f"Circuit breaker {func_name} reset to CLOSED")
-
-        self.failure_count = 0
-        self.last_failure_time = None
+        async with self._lock:
+            if self.state == CircuitBreakerState.HALF_OPEN:
+                self.state = CircuitBreakerState.CLOSED
+                logger.info(f"Circuit breaker {func_name} reset to CLOSED")
+            self.failure_count = 0
+            self.last_failure_time = None
 
     async def _on_failure(self, func_name: str, exception: Exception):
         """Handle failed execution."""
-        self.failure_count += 1
-        self.last_failure_time = time.time()
-
-        if self.failure_count >= self.config.failure_threshold:
-            self.state = CircuitBreakerState.OPEN
-            logger.warning(
-                f"Circuit breaker {func_name} opened after {self.failure_count} failures. " f"Last error: {exception}"
-            )
+        async with self._lock:
+            self.failure_count += 1
+            self.last_failure_time = time.time()
+            if self.failure_count >= self.config.failure_threshold:
+                self.state = CircuitBreakerState.OPEN
+                logger.warning(
+                    f"Circuit breaker {func_name} opened after {self.failure_count} failures. "
+                    f"Last error: {exception}"
+                )
 
 
 class TimeoutManager:
@@ -155,13 +158,22 @@ async def wrapper(*args, **kwargs):
         return decorator
 
 
+def _is_retryable(exception):
+    """Only retry on transient errors (5xx, connection, timeout)."""
+    if isinstance(exception, httpx.RequestError):
+        return True  # Connection/timeout errors
+    if isinstance(exception, httpx.HTTPStatusError):
+        return exception.response.status_code >= 500
+    return False
+
+
 class RetryConfig:
     """Retry configuration."""
 
     WAZUH_API_RETRY = retry(
         stop=stop_after_attempt(3),
         wait=wait_exponential(multiplier=1, min=1, max=10),
-        retry=retry_if_exception_type((httpx.RequestError, httpx.HTTPStatusError)),
+        retry=retry_if_exception(_is_retryable),
         reraise=True,
     )
 
diff --git a/src/wazuh_mcp_server/server.py b/src/wazuh_mcp_server/server.py

Original file line number	Diff line number	Diff line change
`@@ -46,6 +46,8 @@ dependencies = [`
`46`	`46`	`# Utilities`
`47`	`47`	`"python-dotenv>=1.0.0",`
`48`	`48`	`"aiofiles>=24.0.0",`
	`49`	`+ # MCP Framework`
	`50`	`+ "fastmcp>=2.14.0",`
`49`	`51`	`]`
`50`	`52`
`51`	`53`	`[project.optional-dependencies]`