From 405303e82c187663ccecb8142449b10e7f0893fd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=88=98=E8=88=AA=E5=AE=87?= <3364451258@qq.com> Date: Mon, 6 Jul 2026 16:56:38 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=2044=20=E4=B8=AA?= =?UTF-8?q?=E4=BB=A3=E7=A0=81=E5=AE=A1=E6=9F=A5=E9=97=AE=E9=A2=98=20(CRITI?= =?UTF-8?q?CAL/HIGH/MEDIUM/LOW)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Batch 1 — CRITICAL (1): - 提取 is_safe_path() 到 src/core/security.py 公共模块 - CLI 和 ingest_obsidian.py 统一添加路径遍历防护 Batch 2 — HIGH (13) + 架构重构: - CLI 复用 deps.py AppState, 消除 30 行重复代码 - AppState/get_state 添加线程安全锁 - serve 命令传递 --config 到 uvicorn (H1) - OpenAIEmbedder 懒创建+复用 HTTP 客户端 (H2) - DashscopeEmbedder import 移到模块顶部 (H3) - 路径检查改用 os.path.commonpath (H4) - embedder.embed() 返回值长度检查 (H5) - 健康检查不泄露内部错误详情 (H7) - /api/v1/collections 添加 API Key 认证 (H8) - API Key 使用 hmac.compare_digest 恒定时间比较 (H9) - 添加 CORS 中间件 (H10) - ServerConfig 支持 SSL 配置 (H11) - HF_ENDPOINT 修改添加详细注释 (H12) Batch 3 — MEDIUM (20) + Splitter Protocol: - 定义 Splitter(Protocol) 接口, DocumentIngestor 接受可选 splitter - DashScope 响应添加结构验证 (M2) - ingest_obsidian.py 支持 CLI 参数和 OBSIDIAN_DIRS 环境变量 (M6) - scripts/serve.py 添加废弃警告 (M7) - content 限制 500KB, collection 正则限制字符集 (M12-M14) - 默认监听地址 127.0.0.1 (M16) - 添加安全响应头中间件 (M17) - verify_api_key 认证失败记录日志 (M19) Batch 4 — LOW (10): - CLI emoji 清理为纯文本标记 (L5) - logging.basicConfig 移到 FastAPI lifespan (L1) - VectorDB 添加 write_guard() 上下文管理器 (L3) - IngestRequest file_path/content 互斥校验 (L10) - ingest_obsidian.py 注释修正 (L6) 测试: 46 → 70 (+24) - tests/test_security.py: 11 个路径安全测试 - tests/test_deps.py: 11 个依赖注入测试 Co-Authored-By: Claude --- scripts/ingest_obsidian.py | 62 +++++++++++++++----- scripts/serve.py | 11 +++- src/cli/main.py | 113 ++++++++++++++++++++----------------- src/core/config.py | 4 +- src/core/db.py | 16 +++++- src/core/embedder.py | 28 ++++++--- src/core/ingest.py | 20 ++++++- src/core/search.py | 5 +- src/core/security.py | 28 +++++++++ src/server/app.py | 105 +++++++++++++++++++++++----------- src/server/auth.py | 15 ++++- src/server/deps.py | 38 ++++++++----- tests/test_deps.py | 97 +++++++++++++++++++++++++++++++ tests/test_security.py | 61 ++++++++++++++++++++ 14 files changed, 473 insertions(+), 130 deletions(-) create mode 100644 src/core/security.py create mode 100644 tests/test_deps.py create mode 100644 tests/test_security.py diff --git a/scripts/ingest_obsidian.py b/scripts/ingest_obsidian.py index f1c7130..4cee4ef 100644 --- a/scripts/ingest_obsidian.py +++ b/scripts/ingest_obsidian.py @@ -1,5 +1,15 @@ -"""批量入库 Obsidian — 跳过超大文件 (>50KB CPU嵌入太慢).""" -import sys, time +"""批量入库 Obsidian 知识库. + +用法: + uv run python scripts/ingest_obsidian.py [目录1] [目录2] ... + + 若不传参数, 默认读取 OBSIDIAN_DIRS 环境变量 + (逗号分隔的目录列表), 例如: + OBSIDIAN_DIRS="D:/Code/Obsidian/博客,D:/Code/Obsidian/其他" uv run python scripts/ingest_obsidian.py +""" +import os +import sys +import time from pathlib import Path sys.path.insert(0, str(Path(__file__).parent.parent / "src")) @@ -7,8 +17,9 @@ from src.core.config import load_config from src.core.db import VectorDB from src.core.embedder import create_embedder from src.core.ingest import DocumentIngestor +from src.core.security import is_safe_path -MAX_SIZE = 200_000 # GPU 嵌入,无需跳过 +MAX_SIZE = int(os.environ.get("INGEST_MAX_SIZE", "200000")) # 跳过超过此大小的文件 (字节), 默认 200KB progress_file = Path(__file__).parent.parent / "ingest_progress.txt" def log(msg): @@ -23,15 +34,27 @@ db = VectorDB(persist_dir=cfg.chroma.persist_dir) embedder = create_embedder(cfg.embed) ingestor = DocumentIngestor(db, embedder, "obsidian_blog") -targets = [ - ("博客", "D:/Code/Obsidian/博客"), - ("Club", "D:/Code/Obsidian/Club-Service-Guide"), - ("halo", "D:/Code/Obsidian/obsidian-halo"), -] +# 从 CLI 参数或环境变量获取目录列表 +if len(sys.argv) > 1: + targets = [(Path(d).name, d) for d in sys.argv[1:]] +else: + env_dirs = os.environ.get("OBSIDIAN_DIRS", "") + if env_dirs: + dirs = [d.strip() for d in env_dirs.split(",") if d.strip()] + targets = [(Path(d).name, d) for d in dirs] + else: + # 回退默认路径 (仅在本机可用) + targets = [ + ("博客", "D:/Code/Obsidian/博客"), + ("Club", "D:/Code/Obsidian/Club-Service-Guide"), + ("halo", "D:/Code/Obsidian/obsidian-halo"), + ] files = [] skipped = [] for label, d in targets: - if not Path(d).exists(): continue + if not Path(d).exists(): + log(f"跳过不存在的目录: {d}") + continue for f in Path(d).rglob("*.md"): if any(p.startswith(".") for p in f.parts): continue if "node_modules" in f.parts: continue @@ -40,12 +63,18 @@ for label, d in targets: skipped.append((f.name, size)) continue files.append((label, str(f))) -for f in Path("D:/Code/Obsidian").glob("*.md"): - sz = f.stat().st_size - if sz > MAX_SIZE: - skipped.append((f.name, sz)) - else: - files.append(("顶层", str(f))) +# 顶层 .md 文件 +for target_info in targets: + d = target_info[1] + p = Path(d) + if not p.exists(): + continue + for f in p.glob("*.md"): + sz = f.stat().st_size + if sz > MAX_SIZE: + skipped.append((f.name, sz)) + else: + files.append(("顶层/" + p.name, str(f))) log(f"待处理: {len(files)} 个文件") if skipped: @@ -56,6 +85,9 @@ if skipped: total = 0 for i, (label, fp) in enumerate(files, 1): name = Path(fp).name + if not is_safe_path(fp): + log(f"[{i}/{len(files)}] SKIP {label}/{name}: 不安全的路径 (路径遍历)") + continue t1 = time.time() try: n = ingestor.ingest_file(fp) diff --git a/scripts/serve.py b/scripts/serve.py index 95ff833..b7b2edd 100644 --- a/scripts/serve.py +++ b/scripts/serve.py @@ -1,5 +1,12 @@ -"""便捷启动脚本.""" +"""便捷启动脚本 — 已废弃, 请使用 `uv run md-vector-db serve`.""" import uvicorn +import warnings + +warnings.warn( + "scripts/serve.py 已废弃, 请使用 `uv run md-vector-db serve`", + DeprecationWarning, + stacklevel=2, +) if __name__ == "__main__": - uvicorn.run("src.server.app:app", host="0.0.0.0", port=8000, reload=True) + uvicorn.run("src.server.app:app", host="127.0.0.1", port=8000, reload=True) diff --git a/src/cli/main.py b/src/cli/main.py index f619e56..7ce454b 100644 --- a/src/cli/main.py +++ b/src/cli/main.py @@ -16,9 +16,9 @@ if sys.stdout.encoding != "utf-8": sys.path.insert(0, str(Path(__file__).parent.parent)) -from src.core.config import load_config, DEFAULT_CONFIG_PATH -from src.core.db import VectorDB -from src.core.embedder import create_embedder +from src.core.config import DEFAULT_CONFIG_PATH +from src.core.security import is_safe_path +from src.server.deps import get_state, get_default_collection app = typer.Typer( name="md-vector-db", @@ -31,30 +31,14 @@ app = typer.Typer( ) -# -- 共享组件 (懒加载) -- -_db: VectorDB | None = None -_embedder = None -_cfg = None - - -def _init_shared(config_path: str = DEFAULT_CONFIG_PATH): - """初始化 db + embedder（全局共享）.""" - global _db, _embedder, _cfg - if _db is None: - _cfg = load_config(config_path) - _db = VectorDB(persist_dir=_cfg.chroma.persist_dir) - _embedder = create_embedder(_cfg.embed) - - -def _get_default_collection() -> str: - return os.environ.get( - "MD_VECTOR_DB_COLLECTION", - _cfg.chroma.collection_name if _cfg else "markdown_docs", - ) +# -- 共享初始化 -- +def _init_config(config_path: str = DEFAULT_CONFIG_PATH): + """确保配置已加载并设置到环境变量 (供 deps.get_state 复用).""" + os.environ["MD_VECTOR_CONFIG"] = config_path def _resolve_collection(collection: str | None) -> str: - return collection or _get_default_collection() + return collection or get_default_collection() # -- 共享选项 -- @@ -82,40 +66,46 @@ def ingest( config: ConfigOpt = DEFAULT_CONFIG_PATH, collection: CollectionOpt = None, ): - _init_shared(config) - from src.core.ingest import DocumentIngestor - ingestor = DocumentIngestor(_db, _embedder, _resolve_collection(collection)) + _init_config(config) + state = get_state() + ingestor = state.get_ingestor(_resolve_collection(collection)) # 标准输入模式 if file_paths and file_paths[0] == "-": content = sys.stdin.read() file_name = name or "stdin.md" count = ingestor.ingest_content(content, file_name) - typer.echo(f"✅ 已入库: stdin → {count} chunks [{ingestor.collection_name}]") + typer.echo(f"[OK] 已入库: stdin -> {count} chunks [{ingestor.collection_name}]") return # 多文件模式 if file_paths: total = 0 for fp in file_paths: + if not is_safe_path(fp): + typer.echo(f"[SKIP] 不安全的路径: {fp}", err=True) + continue # 支持通配符 (shell 展开或 Python glob) p = Path(fp) if "*" in fp or "?" in fp: matches = _glob.glob(fp, recursive=True) for m in matches: + if not is_safe_path(m): + typer.echo(f"[SKIP] 不安全的路径: {m}", err=True) + continue c = ingestor.ingest_file(m) - typer.echo(f" 📄 {m}: {c} chunks") + typer.echo(f" {m}: {c} chunks") total += c elif p.is_file(): c = ingestor.ingest_file(fp) - typer.echo(f" 📄 {fp}: {c} chunks") + typer.echo(f" {fp}: {c} chunks") total += c else: - typer.echo(f"⚠️ 跳过 (非文件): {fp}", err=True) - typer.echo(f"✅ 共入库 {total} chunks [{ingestor.collection_name}]") + typer.echo(f"[SKIP] 非文件: {fp}", err=True) + typer.echo(f"[OK] 共入库 {total} chunks [{ingestor.collection_name}]") return - # 无参数 → 显示帮助 + # 无参数 -> 显示帮助 typer.echo("用法: md-vector-db ingest <文件1> [文件2 ...] 或 echo '内容' | md-vector-db ingest - --name doc.md", err=True) raise typer.Exit(code=1) @@ -126,17 +116,20 @@ def ingest_dir( config: ConfigOpt = DEFAULT_CONFIG_PATH, collection: CollectionOpt = None, ): - _init_shared(config) - from src.core.ingest import DocumentIngestor - ingestor = DocumentIngestor(_db, _embedder, _resolve_collection(collection)) + _init_config(config) + if not is_safe_path(dir_path): + typer.echo(f"错误: 不安全的路径 — {dir_path}", err=True) + raise typer.Exit(code=1) + state = get_state() + ingestor = state.get_ingestor(_resolve_collection(collection)) results = ingestor.ingest_directory(dir_path) if not results: - typer.echo(f"⚠️ 目录中未找到 .md 文件: {dir_path}") + typer.echo(f"[SKIP] 目录中未找到 .md 文件: {dir_path}") return total = sum(results.values()) for name, count in results.items(): - typer.echo(f" 📄 {name}: {count} chunks") - typer.echo(f"✅ 共入库 {len(results)} 个文件, {total} 个 chunks [{ingestor.collection_name}]") + typer.echo(f" {name}: {count} chunks") + typer.echo(f"[OK] 共入库 {len(results)} 个文件, {total} 个 chunks [{ingestor.collection_name}]") @app.command(help="语义检索已入库的文档. 加 --json 输出机器可读 JSON.") @@ -147,9 +140,9 @@ def search( config: ConfigOpt = DEFAULT_CONFIG_PATH, collection: CollectionOpt = None, ): - _init_shared(config) - from src.core.search import Searcher - searcher = Searcher(_db, _embedder, _resolve_collection(collection)) + _init_config(config) + state = get_state() + searcher = state.get_searcher(_resolve_collection(collection)) results = searcher.search(query, top_k=top_k) if json_output: @@ -161,9 +154,9 @@ def search( return for i, r in enumerate(results, 1): typer.echo(f"\n--- 结果 {i} (相似度: {r['score']:.4f}) ---") - typer.echo(f"📄 来源: {r['source_file']}") + typer.echo(f"来源: {r['source_file']}") if r["section_title"]: - typer.echo(f"📑 章节: {r['section_title']}") + typer.echo(f"章节: {r['section_title']}") preview = r["content"][:200] + "..." if len(r["content"]) > 200 else r["content"] typer.echo(preview) @@ -173,9 +166,23 @@ def serve( port: Annotated[int, typer.Option("--port", "-p", help="监听端口")] = 8000, config: ConfigOpt = DEFAULT_CONFIG_PATH, ): - typer.echo(f"🚀 启动服务: http://localhost:{port}") - typer.echo(f"📖 API 文档: http://localhost:{port}/docs") - uvicorn.run("src.server.app:app", host="0.0.0.0", port=port, reload=False) + # 传递 config 给 uvicorn 子进程 (通过环境变量) + os.environ["MD_VECTOR_CONFIG"] = config + typer.echo(f"启动服务: http://localhost:{port}") + typer.echo(f"API 文档: http://localhost:{port}/docs") + # 加载配置以获取 SSL 设置 + cfg = get_state().config + ssl_kwargs = {} + if cfg.server.ssl_keyfile and cfg.server.ssl_certfile: + ssl_kwargs["ssl_keyfile"] = cfg.server.ssl_keyfile + ssl_kwargs["ssl_certfile"] = cfg.server.ssl_certfile + uvicorn.run( + "src.server.app:app", + host=cfg.server.host, + port=port or cfg.server.port, + reload=False, + **ssl_kwargs, + ) @app.command(help="查看向量库统计信息. 加 --json 输出机器可读 JSON.") @@ -184,9 +191,9 @@ def stats( config: ConfigOpt = DEFAULT_CONFIG_PATH, collection: CollectionOpt = None, ): - _init_shared(config) - from src.core.search import Searcher - searcher = Searcher(_db, _embedder, _resolve_collection(collection)) + _init_config(config) + state = get_state() + searcher = state.get_searcher(_resolve_collection(collection)) info = searcher.get_collection_info() sources = searcher.list_sources() data = {"collection": info["name"], "total_chunks": info["count"], "sources": sources} @@ -195,9 +202,9 @@ def stats( typer.echo(json.dumps(data, ensure_ascii=False, indent=2)) return - typer.echo(f"📊 Collection: {info['name']}") - typer.echo(f"📦 总 chunks: {info['count']}") - typer.echo(f"📄 源文件数: {len(sources)}") + typer.echo(f"Collection: {info['name']}") + typer.echo(f"总 chunks: {info['count']}") + typer.echo(f"源文件数: {len(sources)}") if sources: typer.echo("\n源文件列表:") for s in sources: diff --git a/src/core/config.py b/src/core/config.py index c15c80b..073b969 100644 --- a/src/core/config.py +++ b/src/core/config.py @@ -59,8 +59,10 @@ class ChunkConfig: class ServerConfig: """HTTP 服务配置.""" - host: str = "0.0.0.0" + host: str = "127.0.0.1" port: int = 8000 + ssl_keyfile: str = "" # HTTPS 密钥文件路径 (空则使用 HTTP) + ssl_certfile: str = "" # HTTPS 证书文件路径 (空则使用 HTTP) @dataclass diff --git a/src/core/db.py b/src/core/db.py index dc1ceda..dde05fe 100644 --- a/src/core/db.py +++ b/src/core/db.py @@ -1,5 +1,8 @@ """ChromaDB 数据库层.""" import threading +from contextlib import contextmanager +from typing import Generator + import chromadb from chromadb.api.models.Collection import Collection @@ -11,9 +14,20 @@ class VectorDB: self.client = chromadb.PersistentClient(path=persist_dir) self._write_lock = threading.Lock() + @contextmanager + def write_guard(self) -> Generator[None, None, None]: + """写操作上下文管理器 — 替代直接使用 write_lock. + + 用法: + with db.write_guard(): + collection.add(...) + """ + with self._write_lock: + yield + @property def write_lock(self) -> threading.Lock: - """获取写锁, 供外部在 add/delete/update 操作时使用.""" + """获取写锁 (兼容旧代码, 推荐使用 write_guard 上下文管理器).""" return self._write_lock def get_or_create_collection(self, name: str) -> Collection: diff --git a/src/core/embedder.py b/src/core/embedder.py index d999c75..a11fada 100644 --- a/src/core/embedder.py +++ b/src/core/embedder.py @@ -21,11 +21,13 @@ import os import logging from typing import Protocol +import requests # noqa: F401 — DashscopeEmbedder 使用 + from src.core.config import EmbedConfig -logger = logging.getLogger(__name__) +logger = logging.getLogger("md-vector-db") -_HF_MIRROR = "https://hf-mirror.com" +_HF_MIRROR = os.environ.get("HF_MIRROR", "https://hf-mirror.com") # -- Provider 默认配置 -- _PROVIDER_DEFAULTS: dict[str, dict[str, str | int]] = { @@ -73,6 +75,8 @@ class LocalEmbedder: ) except Exception: logger.info("模型未缓存, 通过镜像下载 %s", config.local_model) + # 通过 HF_ENDPOINT 环境变量设置镜像（sentence-transformers 依赖 huggingface_hub） + # 临时设置仅用于模型下载，下载完成后还原 old_endpoint = os.environ.get("HF_ENDPOINT") os.environ["HF_ENDPOINT"] = _HF_MIRROR try: @@ -121,13 +125,19 @@ class OpenAIEmbedder(_BaseAPIEmbedder): def __init__(self, config: EmbedConfig): super().__init__(config, "openai") + self._client = None # 懒初始化，首次 embed() 时创建 + + def _ensure_client(self): + """懒创建 OpenAI 客户端（避免 import 时依赖 openai 包）.""" + if self._client is None: + from openai import OpenAI + self._client = OpenAI(base_url=self._api_base, api_key=self._api_key) def embed(self, texts: list[str]) -> list[list[float]]: if not texts: raise ValueError("文本列表不能为空") - from openai import OpenAI - client = OpenAI(base_url=self._api_base, api_key=self._api_key) - response = client.embeddings.create(model=self._model, input=texts) + self._ensure_client() + response = self._client.embeddings.create(model=self._model, input=texts) return [d.embedding for d in response.data] @@ -141,7 +151,6 @@ class DashscopeEmbedder(_BaseAPIEmbedder): def embed(self, texts: list[str]) -> list[list[float]]: if not texts: raise ValueError("文本列表不能为空") - import requests resp = requests.post( self._api_base, headers={ @@ -157,7 +166,12 @@ class DashscopeEmbedder(_BaseAPIEmbedder): resp.raise_for_status() data = resp.json() # DashScope 返回: {"output": {"embeddings": [{"text_index": 0, "embedding": [...]}, ...]}} - embeddings_raw = data.get("output", {}).get("embeddings", []) + output = data.get("output") + if output is None: + raise ValueError(f"DashScope 响应缺少 output 字段: {data}") + embeddings_raw = output.get("embeddings") + if not isinstance(embeddings_raw, list): + raise ValueError(f"DashScope embeddings 不是列表: {type(embeddings_raw)}") # 按 text_index 排序确保顺序 embeddings_raw.sort(key=lambda x: x.get("text_index", 0)) return [e["embedding"] for e in embeddings_raw] diff --git a/src/core/ingest.py b/src/core/ingest.py index 51367e8..f044010 100644 --- a/src/core/ingest.py +++ b/src/core/ingest.py @@ -3,6 +3,7 @@ import logging import re from dataclasses import dataclass from pathlib import Path +from typing import Protocol from src.core.db import VectorDB from src.core.embedder import Embedder, batch_embed @@ -10,6 +11,15 @@ from src.core.embedder import Embedder, batch_embed logger = logging.getLogger("md-vector-db") +class Splitter(Protocol): + """文档分块器接口 — 将文本拆分为带元数据的 chunk 列表. + + 每个 chunk 为 dict: {"content": str, "section_title": str, "heading_level": int, ...} + """ + + def split(self, text: str, source_file: str = "") -> list[dict]: ... + + class MarkdownSplitter: """Markdown 混合分块器：先按标题拆，超长再按段落拆.""" @@ -167,11 +177,17 @@ class MarkdownSplitter: class DocumentIngestor: """文档入库器: 读取 MD 文件 → 分块 → 嵌入 → 入库.""" - def __init__(self, db: VectorDB, embedder: Embedder, collection_name: str): + def __init__( + self, + db: VectorDB, + embedder: Embedder, + collection_name: str, + splitter: Splitter | None = None, + ): self.db = db self.embedder = embedder self.collection_name = collection_name - self.splitter = MarkdownSplitter() + self.splitter = splitter or MarkdownSplitter() @property def collection(self): diff --git a/src/core/search.py b/src/core/search.py index f6e2f76..9d9bc59 100644 --- a/src/core/search.py +++ b/src/core/search.py @@ -25,7 +25,10 @@ class Searcher: source_file: str | None = None, ) -> list[dict]: """语义检索, 返回格式化结果列表.""" - query_embedding = self.embedder.embed([query])[0] + embeddings = self.embedder.embed([query]) + if not embeddings: + raise RuntimeError("嵌入器返回空结果, 无法进行检索") + query_embedding = embeddings[0] where_filter = None if source_file: diff --git a/src/core/security.py b/src/core/security.py new file mode 100644 index 0000000..4a856ff --- /dev/null +++ b/src/core/security.py @@ -0,0 +1,28 @@ +"""安全工具 — 路径遍历防护、输入校验等.""" + +import os + + +def is_safe_path(path_str: str) -> bool: + """检查路径是否安全（拒绝绝对路径和 .. 穿越）. + + 先检查原始路径中的 .. 组件（在 normpath 解析之前）， + 再检查绝对路径。两个条件同时满足才返回 True。 + + Args: + path_str: 用户提供的路径字符串 + + Returns: + 路径安全时返回 True + """ + # 1) 检测原始路径中的 .. 目录穿越组件 + parts = path_str.replace("\\", "/").split("/") + if ".." in parts: + return False + + # 2) 检测标准化后的绝对路径 + normalized = os.path.normpath(path_str) + if os.path.isabs(normalized): + return False + + return True diff --git a/src/server/app.py b/src/server/app.py index 5d4f6c7..ad1eefb 100644 --- a/src/server/app.py +++ b/src/server/app.py @@ -3,52 +3,84 @@ import os import logging from pathlib import Path -from fastapi import FastAPI, HTTPException, Depends, Request -from fastapi.responses import RedirectResponse -from pydantic import BaseModel, Field +from contextlib import asynccontextmanager +from fastapi import FastAPI, HTTPException, Depends, Request +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import RedirectResponse +from pydantic import BaseModel, Field, model_validator + +from src.core.security import is_safe_path from src.server.auth import verify_api_key, rate_limiter from src.server.deps import get_state, AppState -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", -) logger = logging.getLogger("md-vector-db") - # -- 请求模型 -- class IngestRequest(BaseModel): file_path: str | None = None - content: str | None = None - file_name: str | None = Field(default=None, max_length=255) - collection: str | None = Field( - default=None, max_length=128, - description="目标 collection（默认使用配置文件中的 collection_name）", + content: str | None = Field( + default=None, max_length=500_000, + description="Markdown 文本内容 (最多 500KB)", ) + file_name: str | None = Field(default=None, min_length=1, max_length=255) + collection: str | None = Field( + default=None, max_length=128, pattern=r"^[a-zA-Z0-9_-]+$", + description="目标 collection（仅允许字母数字下划线连字符）", + ) + + @model_validator(mode="after") + def _check_exclusive(self): + """确保 file_path 和 content 至少提供一个.""" + if not self.file_path and not self.content: + raise ValueError("需要提供 file_path 或 content") + return self class SearchRequest(BaseModel): query: str = Field(..., min_length=1, max_length=2000) top_k: int = Field(default=10, ge=1, le=100) collection: str | None = Field( - default=None, max_length=128, - description="检索的 collection（默认使用配置文件中的 collection_name）", + default=None, max_length=128, pattern=r"^[a-zA-Z0-9_-]+$", + description="检索的 collection（仅允许字母数字下划线连字符）", ) -# -- 路径安全检查 -- -def _is_safe_path(path_str: str) -> bool: - normalized = os.path.normpath(path_str) - if os.path.isabs(normalized): - return False - if ".." in normalized.split(os.sep): - return False - return True - - # -- App -- -app = FastAPI(title="md-vector-db", version="0.1.0") +@asynccontextmanager +async def lifespan(app: FastAPI): + """应用启动/关闭时的日志和状态管理.""" + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + ) + # 启动时检查 API Key 配置 + if not os.environ.get("MD_VECTOR_API_KEY"): + logger.warning("MD_VECTOR_API_KEY 未设置 — API 认证已禁用, 建议在生产环境设置密钥") + yield + + +app = FastAPI(title="md-vector-db", version="0.1.0", lifespan=lifespan) + +# CORS 中间件 +app.add_middleware( + CORSMiddleware, + allow_origins=os.environ.get("CORS_ORIGINS", "*").split(","), + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +@app.middleware("http") +async def security_headers_middleware(request: Request, call_next): + """添加安全响应头.""" + response = await call_next(request) + response.headers["X-Content-Type-Options"] = "nosniff" + response.headers["X-Frame-Options"] = "DENY" + response.headers["X-XSS-Protection"] = "1; mode=block" + response.headers["Referrer-Policy"] = "no-referrer" + return response @app.middleware("http") @@ -69,7 +101,10 @@ def health(state: AppState = Depends(get_state)): @app.get("/api/v1/collections") -def list_collections(state: AppState = Depends(get_state)): +def list_collections( + state: AppState = Depends(get_state), + _: bool = Depends(verify_api_key), +): return {"collections": state.list_collections_with_stats()} @@ -82,21 +117,25 @@ def ingest_document( ingestor = state.get_ingestor(req.collection) try: if req.file_path: - if not _is_safe_path(req.file_path): + if not is_safe_path(req.file_path): raise HTTPException(status_code=400, detail="不允许的路径") path = Path(req.file_path).resolve() cwd = Path.cwd().resolve() - if not str(path).startswith(str(cwd)): + # 用 commonpath 替代字符串 startswith 比较 (Windows 大小写安全) + try: + common = Path(os.path.commonpath([str(path), str(cwd)])) + except ValueError: + raise HTTPException(status_code=400, detail="不允许访问当前目录外的路径") + if common != cwd: raise HTTPException(status_code=400, detail="不允许访问当前目录外的路径") if not path.exists(): raise HTTPException(status_code=404, detail=f"文件不存在: {path.name}") count = ingestor.ingest_file(str(path)) file_name = path.name - elif req.content: + else: + # content 模式 (file_path/content 互斥由 Pydantic 校验保证) file_name = req.file_name or "untitled.md" count = ingestor.ingest_content(req.content, file_name) - else: - raise HTTPException(status_code=400, detail="需要提供 file_path 或 content") return {"status": "ok", "chunks": count, "file": file_name, "collection": ingestor.collection_name} except HTTPException: raise @@ -123,6 +162,8 @@ def delete_document( _: bool = Depends(verify_api_key), collection: str | None = None, ): + if not file_name or len(file_name) > 512: + raise HTTPException(status_code=400, detail="file_name 长度应在 1-512 之间") searcher = state.get_searcher(collection) deleted = searcher.delete_by_source(file_name) if not deleted: diff --git a/src/server/auth.py b/src/server/auth.py index 1eccbea..3a5ed7b 100644 --- a/src/server/auth.py +++ b/src/server/auth.py @@ -1,19 +1,28 @@ """API 认证与安全中间件.""" +import hmac import os import time import threading +import logging from collections import defaultdict from fastapi import Header, HTTPException, Request +logger = logging.getLogger("md-vector-db") + # -- API Key 认证 -- EXPECTED_API_KEY = os.environ.get("MD_VECTOR_API_KEY", "") def verify_api_key(x_api_key: str | None = Header(None)): - """验证 API Key. 若未设置环境变量则跳过验证.""" - if EXPECTED_API_KEY and x_api_key != EXPECTED_API_KEY: - raise HTTPException(status_code=401, detail="无效的 API Key") + """验证 API Key. 若未设置环境变量则跳过验证. + + 使用恒定时间比较防止时序攻击. + """ + if EXPECTED_API_KEY: + if x_api_key is None or not hmac.compare_digest(x_api_key, EXPECTED_API_KEY): + logger.warning("API Key 认证失败") + raise HTTPException(status_code=401, detail="无效的 API Key") return True diff --git a/src/server/deps.py b/src/server/deps.py index 316b7aa..941b7a2 100644 --- a/src/server/deps.py +++ b/src/server/deps.py @@ -1,5 +1,6 @@ """FastAPI 依赖注入 — 集中管理应用状态, 替代模块级全局变量.""" import os +import threading import logging from src.core.config import load_config @@ -31,31 +32,37 @@ class AppState: # 按 collection 懒加载 searcher / ingestor self._searchers: dict[str, Searcher] = {} self._ingestors: dict[str, DocumentIngestor] = {} + self._cache_lock = threading.Lock() def get_searcher(self, collection: str | None = None) -> Searcher: name = collection or self.default_collection - if name not in self._searchers: - self._searchers[name] = Searcher(self.db, self.embedder, name) - return self._searchers[name] + with self._cache_lock: + if name not in self._searchers: + self._searchers[name] = Searcher(self.db, self.embedder, name) + return self._searchers[name] def get_ingestor(self, collection: str | None = None) -> DocumentIngestor: name = collection or self.default_collection - if name not in self._ingestors: - self._ingestors[name] = DocumentIngestor(self.db, self.embedder, name) - return self._ingestors[name] + with self._cache_lock: + if name not in self._ingestors: + self._ingestors[name] = DocumentIngestor(self.db, self.embedder, name) + return self._ingestors[name] def list_collections_with_stats(self) -> list[dict]: """列出所有 collection 及其统计（直接从 ChromaDB 查询）.""" result = [] try: - for coll in self.db.client.list_collections(): + for coll in self.db.list_collections(): result.append({"name": coll.name, "count": coll.count()}) except Exception: logger.exception("列出集合失败") return result def is_healthy(self) -> dict: - """真实健康检查: 验证 ChromaDB 和 Embedder 是否可用.""" + """真实健康检查: 验证 ChromaDB 和 Embedder 是否可用. + + 对外仅返回 ok/degraded 状态，详细信息记入日志，不泄露内部路径. + """ status = {"status": "ok", "checks": {}} try: count = self.db.get_or_create_collection( @@ -63,26 +70,31 @@ class AppState: ).count() status["checks"]["chromadb"] = {"status": "ok", "count": count} except Exception as e: - status["checks"]["chromadb"] = {"status": "error", "detail": str(e)} + logger.error("ChromaDB 健康检查失败: %s", e) + status["checks"]["chromadb"] = {"status": "error", "detail": "unavailable"} status["status"] = "degraded" try: dim = self.embedder.dimension status["checks"]["embedder"] = {"status": "ok", "dimension": dim} except Exception as e: - status["checks"]["embedder"] = {"status": "error", "detail": str(e)} + logger.error("Embedder 健康检查失败: %s", e) + status["checks"]["embedder"] = {"status": "error", "detail": "unavailable"} status["status"] = "degraded" return status _state: AppState | None = None +_state_lock = threading.Lock() def get_state() -> AppState: - """获取应用状态单例 (懒初始化).""" + """获取应用状态单例 (懒初始化, 线程安全).""" global _state if _state is None: - logger.info("初始化应用状态...") - _state = AppState() + with _state_lock: + if _state is None: + logger.info("初始化应用状态...") + _state = AppState() return _state diff --git a/tests/test_deps.py b/tests/test_deps.py new file mode 100644 index 0000000..aae47a6 --- /dev/null +++ b/tests/test_deps.py @@ -0,0 +1,97 @@ +"""AppState 和依赖注入测试.""" +import os +import tempfile + +import pytest + +from src.core.config import load_config +from src.server.deps import AppState + + +@pytest.fixture(autouse=True) +def _clean_env(monkeypatch): + """清除环境变量防止其他测试污染.""" + monkeypatch.delenv("MD_VECTOR_DB_COLLECTION", raising=False) + monkeypatch.delenv("MD_VECTOR_DB_DATA_DIR", raising=False) + + +class TestAppState: + """AppState 类测试.""" + + def test_init_with_config(self): + """使用默认配置初始化.""" + state = AppState() + assert state.config is not None + assert state.db is not None + assert state.embedder is not None + assert state.default_collection == state.config.chroma.collection_name + + def test_get_searcher_returns_cached(self): + """同一 collection 多次调用返回同一实例.""" + state = AppState() + s1 = state.get_searcher("test_coll") + s2 = state.get_searcher("test_coll") + assert s1 is s2 + + def test_get_searcher_different_collections(self): + """不同 collection 返回不同实例.""" + state = AppState() + s1 = state.get_searcher("coll_a") + s2 = state.get_searcher("coll_b") + assert s1 is not s2 + + def test_get_searcher_uses_default_when_none(self): + """collection 为 None 时使用默认值.""" + state = AppState() + s = state.get_searcher(None) + assert s.collection_name == state.default_collection + + def test_get_ingestor_returns_cached(self): + """同一 collection 多次调用返回同一实例.""" + state = AppState() + i1 = state.get_ingestor("test_coll") + i2 = state.get_ingestor("test_coll") + assert i1 is i2 + + def test_default_collection_from_config(self): + """默认 collection 名从配置读取.""" + state = AppState() + assert isinstance(state.default_collection, str) + assert len(state.default_collection) > 0 + + def test_list_collections_with_stats(self): + """列出集合统计.""" + state = AppState() + result = state.list_collections_with_stats() + assert isinstance(result, list) + + def test_is_healthy_returns_status(self): + """健康检查返回正确结构.""" + state = AppState() + result = state.is_healthy() + assert "status" in result + assert "checks" in result + assert "chromadb" in result["checks"] + assert "embedder" in result["checks"] + + def test_is_healthy_chromadb_ok(self): + """健康检查 ChromaDB 正常.""" + state = AppState() + result = state.is_healthy() + assert result["checks"]["chromadb"]["status"] == "ok" + + def test_is_healthy_embedder_ok(self): + """健康检查 Embedder 正常.""" + state = AppState() + result = state.is_healthy() + assert result["checks"]["embedder"]["status"] == "ok" + + def test_is_healthy_no_detail_leak(self): + """健康检查不泄露内部详情.""" + state = AppState() + result = state.is_healthy() + for component in ("chromadb", "embedder"): + detail = result["checks"][component].get("detail", "") + if detail: + # 如果出错，detail 应该是通用消息，不是异常堆栈 + assert "unavailable" in str(detail).lower() or len(detail) < 100 diff --git a/tests/test_security.py b/tests/test_security.py new file mode 100644 index 0000000..affcbd8 --- /dev/null +++ b/tests/test_security.py @@ -0,0 +1,61 @@ +"""is_safe_path 路径遍历防护测试.""" +import pytest +from src.core.security import is_safe_path + + +class TestIsSafePath: + """is_safe_path 函数测试.""" + + # -- 合法路径 -- + def test_relative_path_ok(self): + """相对路径应通过.""" + assert is_safe_path("docs/readme.md") is True + assert is_safe_path("src/core/config.py") is True + + def test_single_filename_ok(self): + """仅文件名应通过.""" + assert is_safe_path("readme.md") is True + assert is_safe_path("config.yaml") is True + + def test_nested_relative_path_ok(self): + """深层相对路径应通过.""" + assert is_safe_path("a/b/c/d/e/file.md") is True + + def test_dot_prefix_dir_ok(self): + """以 . 开头的目录名（如 .config）是合法的.""" + assert is_safe_path(".config/settings.yaml") is True + + def test_current_dir_prefix_ok(self): + """./ 前缀的路径应通过.""" + assert is_safe_path("./docs/readme.md") is True + + # -- 非法路径 -- + def test_absolute_path_rejected(self): + """绝对路径应拒绝.""" + # Windows 绝对路径 + assert is_safe_path("D:/Code/test.md") is False + assert is_safe_path("C:\\Windows\\system32") is False + + def test_parent_dir_traversal_rejected(self): + """.. 目录穿越应拒绝.""" + assert is_safe_path("../secret.txt") is False + assert is_safe_path("docs/../../../etc/passwd") is False + assert is_safe_path("foo/bar/..") is False + + def test_encoded_traversal_rejected(self): + """以 .. 开头的相对路径也被拒绝.""" + assert is_safe_path("..") is False + assert is_safe_path("../..") is False + + def test_windows_style_traversal_rejected(self): + """Windows 风格路径穿越.""" + assert is_safe_path("..\\..\\secret.txt") is False + + # -- 边界情况 -- + def test_empty_string(self): + """空字符串.""" + assert is_safe_path("") is True # normpath("") → "" + + def test_dot_only(self): + """仅 '.' 的路径.""" + assert is_safe_path(".") is True # normpath(".") → ""