fix: harden audit logger and add search input validation

src/comfyui_mcp/audit.py

@@ -2,13 +2,16 @@
 
 from __future__ import annotations
 
+import logging
 from datetime import UTC, datetime
 from pathlib import Path
 
 from pydantic import BaseModel, Field, model_serializer
 
 _SENSITIVE_KEYS = {"token", "password", "secret", "api_key", "authorization"}
 
+logger = logging.getLogger(__name__)
+
 
 def _redact_sensitive(data: dict[str, object]) -> dict[str, object]:
     """Remove sensitive keys from a dictionary."""
@@ -51,16 +54,31 @@ def serialize(self) -> dict[str, object]:
 class AuditLogger:
     def __init__(self, audit_file: Path) -> None:
         self._audit_file = Path(audit_file)
+        self._dir_created = False
+
+    def _ensure_directory(self) -> None:
+        """Create parent directories on first write, with symlink check."""
+        if self._dir_created:
+            return
+        parent = self._audit_file.parent
+        parent.mkdir(parents=True, exist_ok=True)
+        # Reject symlinked audit file — could redirect entries to attacker-controlled path
+        if self._audit_file.exists() and self._audit_file.is_symlink():
+            raise OSError(f"Audit log path is a symlink — refusing to write: {self._audit_file}")
+        self._dir_created = True
 
     def log(self, *, tool: str, action: str, **kwargs) -> AuditRecord:
-        """Write an audit record as a JSON line."""
+        """Write an audit record as a JSON line.
+
+        Raises OSError on write failure — audit log integrity is a security
+        requirement, so failures must not be silently swallowed.
+        """
         record = AuditRecord(tool=tool, action=action, **kwargs)
         try:
-            self._audit_file.parent.mkdir(parents=True, exist_ok=True)
+            self._ensure_directory()
             with open(self._audit_file, "a") as f:
                 f.write(record.model_dump_json() + "\n")
-        except OSError as e:
-            import logging
-
-            logging.getLogger(__name__).error("AUDIT LOG FAILURE: %s", e)
+        except OSError:
+            logger.exception("AUDIT LOG FAILURE — cannot write to %s", self._audit_file)
+            raise
         return record

src/comfyui_mcp/tools/models.py

@@ -23,6 +23,10 @@
 # Extensions we look for when finding the primary model file in HuggingFace repos
 _MODEL_EXTENSIONS = {".safetensors", ".ckpt", ".pt", ".pth", ".bin"}
 
+# Input validation limits for external API queries
+_MAX_QUERY_LENGTH = 200
+_MAX_MODEL_TYPE_LENGTH = 50
+
 
 async def _search_civitai(
     query: str,
@@ -173,6 +177,15 @@ async def search_models(
         if source not in ("civitai", "huggingface"):
             raise ValueError("source must be 'civitai' or 'huggingface'")
 
+        if not query or not query.strip():
+            raise ValueError("query must not be empty")
+        if len(query) > _MAX_QUERY_LENGTH:
+            raise ValueError(f"query too long ({len(query)} chars, max {_MAX_QUERY_LENGTH})")
+        if model_type and len(model_type) > _MAX_MODEL_TYPE_LENGTH:
+            raise ValueError(
+                f"model_type too long ({len(model_type)} chars, max {_MAX_MODEL_TYPE_LENGTH})"
+            )
+
         cap = max(1, min(limit, search_settings.max_search_results))
 
         audit.log(

tests/test_audit.py

@@ -1,6 +1,9 @@
 """Tests for structured audit logging."""
 
 import json
+import os
+
+import pytest
 
 from comfyui_mcp.audit import AuditLogger, AuditRecord
 
@@ -81,3 +84,22 @@ def test_log_strips_sensitive_keys_from_extra(self, tmp_path):
         content = log_file.read_text()
         assert "secret-value" not in content
         assert "a cat" in content
+
+    def test_log_raises_on_write_failure(self, tmp_path):
+        log_file = tmp_path / "readonly" / "audit.log"
+        (tmp_path / "readonly").mkdir()
+        (tmp_path / "readonly").chmod(0o444)
+        logger = AuditLogger(audit_file=log_file)
+        with pytest.raises(OSError):
+            logger.log(tool="test", action="called")
+        # Restore permissions for cleanup
+        (tmp_path / "readonly").chmod(0o755)
+
+    def test_log_rejects_symlinked_path(self, tmp_path):
+        real_file = tmp_path / "real.log"
+        real_file.touch()
+        link = tmp_path / "audit.log"
+        os.symlink(real_file, link)
+        logger = AuditLogger(audit_file=link)
+        with pytest.raises(OSError, match="symlink"):
+            logger.log(tool="test", action="called")

tests/test_tools_models.py

@@ -134,6 +134,24 @@ async def test_search_invalid_source(self, registered_tools):
         with pytest.raises(ValueError, match="source must be"):
             await registered_tools["search_models"](query="test", source="invalid")
 
+    async def test_search_empty_query_rejected(self, registered_tools):
+        with pytest.raises(ValueError, match="must not be empty"):
+            await registered_tools["search_models"](query="", source="civitai")
+
+    async def test_search_whitespace_query_rejected(self, registered_tools):
+        with pytest.raises(ValueError, match="must not be empty"):
+            await registered_tools["search_models"](query="   ", source="civitai")
+
+    async def test_search_query_too_long_rejected(self, registered_tools):
+        with pytest.raises(ValueError, match="query too long"):
+            await registered_tools["search_models"](query="x" * 201, source="civitai")
+
+    async def test_search_model_type_too_long_rejected(self, registered_tools):
+        with pytest.raises(ValueError, match="model_type too long"):
+            await registered_tools["search_models"](
+                query="test", source="civitai", model_type="x" * 51
+            )
+
     @respx.mock
     async def test_search_with_api_key(self, components):
         components["search_settings"] = ModelSearchSettings(civitai_api_key="test_key")