fix: security hardening and stability improvements for MCP tools

mambo-wang · mambo-wang · commit 30c31df66c6a · 2026-06-22T08:07:38.000+08:00
- Fix shell injection in view_repo_file: replace shell=True subprocess
  with pathlib iteration
- Add path traversal guards in view_repo_file, write_doc_file,
  edit_doc_file (reject paths escaping repo/output dir)
- Add threading.Lock to SessionStore for concurrent access safety
- Cap max sessions to 10, evict oldest when full
- Cap read_code_components to 50 IDs per call
- Cap edit history to 20 entries per file
- Store edit history as native dict instead of JSON string
- Fix undo to run Mermaid validation after reverting content
- Fix Mermaid validation to report "skipped" instead of false success
- Fix Mermaid timeout to warn instead of silent pass
- Fix pagination hint to point to list_components instead of analyze_repo
- Clamp offset to non-negative in _build_component_index
- Add smoke test covering all critical paths (25 assertions)
diff --git a/codewiki/mcp/session.py b/codewiki/mcp/session.py
@@ -8,10 +8,10 @@
 
 from __future__ import annotations
 
+import threading
 import time
 import uuid
 from dataclasses import dataclass, field
-from pathlib import Path
 from typing import Any, Dict, List, Optional
 
 from codewiki.src.be.dependency_analyzer.models.core import Node
@@ -20,6 +20,9 @@
 # Sessions auto-expire after this many seconds of inactivity.
 _SESSION_TTL_SECONDS = 2 * 60 * 60  # 2 hours
 
+# Maximum concurrent sessions to prevent unbounded memory growth.
+_MAX_SESSIONS = 10
+
 
 @dataclass
 class SessionState:
@@ -45,10 +48,11 @@ def is_expired(self) -> bool:
 
 
 class SessionStore:
-    """In-memory store for all active MCP sessions."""
+    """In-memory store for all active MCP sessions (thread-safe)."""
 
     def __init__(self) -> None:
         self._sessions: Dict[str, SessionState] = {}
+        self._lock = threading.Lock()
 
     def create(
         self,
@@ -58,35 +62,48 @@ def create(
         leaf_nodes: List[str],
     ) -> SessionState:
         """Create a new session and return it."""
-        session_id = uuid.uuid4().hex[:12]
-        state = SessionState(
-            session_id=session_id,
-            repo_path=repo_path,
-            output_dir=output_dir,
-            components=components,
-            leaf_nodes=leaf_nodes,
-        )
-        self._sessions[session_id] = state
-        self._purge_expired()
-        return state
+        with self._lock:
+            self._purge_expired_locked()
+            # Evict oldest if at capacity
+            if len(self._sessions) >= _MAX_SESSIONS:
+                oldest_id = min(
+                    self._sessions,
+                    key=lambda sid: self._sessions[sid].last_accessed,
+                )
+                del self._sessions[oldest_id]
+            session_id = uuid.uuid4().hex[:12]
+            # Ensure no collision
+            while session_id in self._sessions:
+                session_id = uuid.uuid4().hex[:12]
+            state = SessionState(
+                session_id=session_id,
+                repo_path=repo_path,
+                output_dir=output_dir,
+                components=components,
+                leaf_nodes=leaf_nodes,
+            )
+            self._sessions[session_id] = state
+            return state
 
     def get(self, session_id: str) -> Optional[SessionState]:
         """Return the session or ``None`` if not found / expired."""
-        state = self._sessions.get(session_id)
-        if state is None:
-            return None
-        if state.is_expired:
-            del self._sessions[session_id]
-            return None
-        state.touch()
-        return state
+        with self._lock:
+            state = self._sessions.get(session_id)
+            if state is None:
+                return None
+            if state.is_expired:
+                del self._sessions[session_id]
+                return None
+            state.touch()
+            return state
 
     def remove(self, session_id: str) -> bool:
         """Remove a session.  Returns True if it existed."""
-        return self._sessions.pop(session_id, None) is not None
+        with self._lock:
+            return self._sessions.pop(session_id, None) is not None
 
-    def _purge_expired(self) -> None:
-        """Remove all expired sessions."""
+    def _purge_expired_locked(self) -> None:
+        """Remove all expired sessions.  Caller must hold _lock."""
         expired = [sid for sid, s in self._sessions.items() if s.is_expired]
         for sid in expired:
             del self._sessions[sid]
diff --git a/codewiki/mcp/tools/analysis.py b/codewiki/mcp/tools/analysis.py
@@ -33,7 +33,8 @@ def _build_component_index(
     """
     all_ids = list(components.keys())
     total = len(all_ids)
-    limit = min(max(limit, 1), 200)  # clamp to [1, 200]
+    offset = max(0, int(offset))       # prevent negative-index wrapping
+    limit = min(max(int(limit), 1), 200)  # clamp to [1, 200]
     page_ids = all_ids[offset : offset + limit]
     index: list[dict] = []
     for comp_id in page_ids:
@@ -336,7 +337,7 @@ def handle_analyze_repo(
     if pagination["has_more"]:
         result["hint"] += (
             f" Component index has {pagination['total']} items; "
-            f"call analyze_repo again with offset={offset + limit} to see the next page."
+            f"call list_components(session_id='{session.session_id}', offset={offset + limit}) to see the next page."
         )
     if changes and not changes.get("no_changes"):
         result["hint"] = (
diff --git a/codewiki/mcp/tools/code_reader.py b/codewiki/mcp/tools/code_reader.py
@@ -8,10 +8,8 @@
 
 import json
 import logging
-import os
-import subprocess
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List
 
 from codewiki.mcp.session import SessionState, SessionStore
 
@@ -20,13 +18,25 @@
 # Truncation guard for very large responses
 _MAX_RESPONSE_LEN = 32000
 
+# Max components per read_code_components call
+_MAX_COMPONENTS_PER_CALL = 50
+
 
 def _maybe_truncate(text: str, limit: int = _MAX_RESPONSE_LEN) -> str:
     if len(text) <= limit:
         return text
     return text[:limit] + "\n\n<response clipped — use view_repo_file with view_range to read more>"
 
 
+def _is_within(path: Path, base: Path) -> bool:
+    """Return True if *path* resolves to somewhere inside *base*."""
+    try:
+        path.resolve().relative_to(base.resolve())
+        return True
+    except ValueError:
+        return False
+
+
 def handle_read_code_components(
     arguments: Dict[str, Any],
     store: SessionStore,
@@ -38,8 +48,11 @@ def handle_read_code_components(
         return json.dumps({"error": f"Session {session_id} not found or expired."})
 
     component_ids: List[str] = arguments["component_ids"]
-    components = session.components
+    # Cap the number of components to avoid oversized responses
+    if len(component_ids) > _MAX_COMPONENTS_PER_CALL:
+        component_ids = component_ids[:_MAX_COMPONENTS_PER_CALL]
 
+    components = session.components
     results = []
     for cid in component_ids:
         node = components.get(cid)
@@ -52,6 +65,8 @@ def handle_read_code_components(
             results.append(f"## {cid} ({getattr(node, 'component_type', '')})\n```{fence}\n{code}\n```\n")
 
     output = "\n".join(results)
+    if len(arguments["component_ids"]) > _MAX_COMPONENTS_PER_CALL:
+        output = f"<only first {_MAX_COMPONENTS_PER_CALL} components shown; call again with the remaining IDs>\n\n" + output
     return _maybe_truncate(output)
 
 
@@ -66,20 +81,43 @@ def handle_view_repo_file(
         return json.dumps({"error": f"Session {session_id} not found or expired."})
 
     rel_path = arguments["path"]
-    abs_path = Path(session.repo_path) / rel_path
+    repo_base = Path(session.repo_path).resolve()
+    abs_path = (repo_base / rel_path).resolve()
+
+    # Path traversal guard
+    if not _is_within(abs_path, repo_base):
+        return json.dumps({"error": "Path escapes repository directory."})
 
     if not abs_path.exists():
         return json.dumps({"error": f"Path not found: {rel_path}"})
 
-    # Directory listing
+    # Directory listing — use pathlib instead of shelling out
     if abs_path.is_dir():
-        out = subprocess.run(
-            rf"find {abs_path} -maxdepth 2 -not -path '*/\.*'",
-            shell=True,
-            capture_output=True,
-        )
-        listing = out.stdout.decode("utf-8", errors="replace")
-        listing = listing.replace(str(abs_path), rel_path)
+        entries: list[str] = []
+        for child in sorted(abs_path.iterdir()):
+            if child.name.startswith("."):
+                continue
+            rel_child = child.relative_to(repo_base)
+            suffix = "/" if child.is_dir() else ""
+            entries.append(f"{rel_child}{suffix}")
+        # Also list one level deeper if there aren't too many entries
+        if len(entries) <= 50:
+            expanded: list[str] = []
+            for child in sorted(abs_path.iterdir()):
+                if child.name.startswith("."):
+                    continue
+                rel_child = child.relative_to(repo_base)
+                suffix = "/" if child.is_dir() else ""
+                expanded.append(f"{rel_child}{suffix}")
+                if child.is_dir():
+                    for sub in sorted(child.iterdir()):
+                        if sub.name.startswith("."):
+                            continue
+                        rel_sub = sub.relative_to(repo_base)
+                        sub_suffix = "/" if sub.is_dir() else ""
+                        expanded.append(f"  {rel_sub}{sub_suffix}")
+            entries = expanded
+        listing = "\n".join(entries)
         return f"Directory listing for {rel_path}:\n{listing}"
 
     # File view
diff --git a/codewiki/mcp/tools/doc_writer.py b/codewiki/mcp/tools/doc_writer.py
@@ -6,17 +6,43 @@
 
 from __future__ import annotations
 
-import asyncio
 import json
 import logging
-import os
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Dict
 
 from codewiki.mcp.session import SessionState, SessionStore
 
 logger = logging.getLogger(__name__)
 
+# Max edit history entries per file (prevent unbounded memory growth)
+_MAX_HISTORY_PER_FILE = 20
+
+
+def _is_within(path: Path, base: Path) -> bool:
+    """Return True if *path* resolves to somewhere inside *base*."""
+    try:
+        path.resolve().relative_to(base.resolve())
+        return True
+    except ValueError:
+        return False
+
+
+def _safe_doc_path(session: SessionState, filename: str) -> Path | None:
+    """Resolve *filename* within session.output_dir, guarding against traversal."""
+    if not filename.endswith(".md"):
+        filename += ".md"
+    output_base = Path(session.output_dir).resolve()
+    doc_path = (output_base / filename).resolve()
+    if not _is_within(doc_path, output_base):
+        return None
+    return doc_path
+
+
+def _ensure_parent_dirs(path: Path) -> None:
+    """Create parent directories if they don't exist."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+
 
 async def _validate_mermaid(file_path: str, relative_path: str) -> str:
     """Run Mermaid validation and return the result string."""
@@ -27,9 +53,20 @@ async def _validate_mermaid(file_path: str, relative_path: str) -> str:
         return f"Mermaid validation skipped: {e}"
 
 
-def _ensure_parent_dirs(path: Path) -> None:
-    """Create parent directories if they don't exist."""
-    path.parent.mkdir(parents=True, exist_ok=True)
+def _save_history(session: SessionState, doc_path: Path, content: str) -> None:
+    """Append *content* to edit history for *doc_path*, capped at _MAX_HISTORY_PER_FILE."""
+    history = session.registry.get("file_history")
+    if history is None:
+        history = {}
+    elif isinstance(history, str):
+        history = json.loads(history)
+    key = str(doc_path)
+    entry = history.setdefault(key, [])
+    entry.append(content)
+    # Trim to last N entries
+    if len(entry) > _MAX_HISTORY_PER_FILE:
+        del entry[: len(entry) - _MAX_HISTORY_PER_FILE]
+    session.registry["file_history"] = history  # keep as native dict
 
 
 async def handle_write_doc_file(
@@ -43,11 +80,12 @@ async def handle_write_doc_file(
         return json.dumps({"error": f"Session {session_id} not found or expired."})
 
     filename = arguments["filename"]
-    if not filename.endswith(".md"):
-        filename += ".md"
+    doc_path = _safe_doc_path(session, filename)
+    if doc_path is None:
+        return json.dumps({"error": "Filename escapes output directory."})
+
     content = arguments["content"]
 
-    doc_path = Path(session.output_dir) / filename
     _ensure_parent_dirs(doc_path)
 
     if doc_path.exists():
@@ -81,36 +119,39 @@ async def handle_edit_doc_file(
         return json.dumps({"error": f"Session {session_id} not found or expired."})
 
     filename = arguments["filename"]
-    if not filename.endswith(".md"):
-        filename += ".md"
+    doc_path = _safe_doc_path(session, filename)
+    if doc_path is None:
+        return json.dumps({"error": "Filename escapes output directory."})
 
-    doc_path = Path(session.output_dir) / filename
     command = arguments["command"]
 
     if command == "undo":
         # Undo via registry history
-        history_key = str(doc_path)
-        history = session.registry.get("file_history", "{}")
-        file_history = json.loads(history) if isinstance(history, str) else history
-        path_history = file_history.get(history_key, [])
+        history = session.registry.get("file_history", {})
+        if isinstance(history, str):
+            history = json.loads(history)
+        path_history = history.get(str(doc_path), [])
         if not path_history:
             return json.dumps({"error": f"No edit history found for {filename}."})
         old_content = path_history.pop()
-        file_history[history_key] = path_history
-        session.registry["file_history"] = json.dumps(file_history)
+        history[str(doc_path)] = path_history
+        session.registry["file_history"] = history
         doc_path.write_text(old_content, encoding="utf-8")
-        return json.dumps({"status": "undone", "filename": filename})
+
+        # Validate Mermaid after undo
+        mermaid_result = await _validate_mermaid(str(doc_path), filename)
+        return json.dumps({
+            "status": "undone",
+            "filename": filename,
+            "mermaid_validation": mermaid_result,
+        }, ensure_ascii=False)
 
     if not doc_path.exists():
         return json.dumps({"error": f"File not found: {filename}. Use write_doc_file to create it."})
 
     # Save current content to history before editing
     current_content = doc_path.read_text(encoding="utf-8")
-    history_key = str(doc_path)
-    history = session.registry.get("file_history", "{}")
-    file_history = json.loads(history) if isinstance(history, str) else history
-    file_history.setdefault(history_key, []).append(current_content)
-    session.registry["file_history"] = json.dumps(file_history)
+    _save_history(session, doc_path, current_content)
 
     if command == "str_replace":
         old_str = arguments.get("old_str")
diff --git a/codewiki/src/be/utils.py b/codewiki/src/be/utils.py
@@ -239,16 +239,15 @@ async def validate_single_diagram(diagram_content: str, diagram_num: int, line_s
     """
     core_error = await _try_pythonmonkey_parse(diagram_content)
     if core_error is None:
-        # Both PythonMonkey (3.12+) and mermaid-py (env disabled) are unavailable
         if _MERMAID_PY_BROKEN:
-            return ""  # Skip validation gracefully
+            return f"  Diagram {diagram_num}: validation skipped (set MERMAID_VALIDATE=1 to enable)"
         try:
             core_error = await asyncio.wait_for(
                 asyncio.to_thread(_parse_via_mermaid_py, diagram_content),
                 timeout=15.0,
             )
         except asyncio.TimeoutError:
-            return ""  # Graceful skip on timeout
+            return f"  Diagram {diagram_num}: validation timed out (15s) — diagram may be invalid"
         except Exception as e:
             return f"  Diagram {diagram_num}: Exception during validation - {str(e)}"
 
diff --git a/tests/smoke_test_mcp.py b/tests/smoke_test_mcp.py
