Serendipity/md-vector-db
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: 修复 44 个代码审查问题 (CRITICAL/HIGH/MEDIUM/LOW)

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
刘航宇
2026-07-06 16:56:38 +08:00
parent 832201186d
commit 405303e82c
14 changed files with 473 additions and 130 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/ingest_obsidian.py
+47 -15
View File
@@ -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)
scripts/serve.py
+9 -2
View File
@@ -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)
src/cli/main.py
+60 -53
View File
@@ -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:
src/core/config.py
+3 -1
View File
@@ -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
src/core/db.py
+15 -1
View File
@@ -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:
src/core/embedder.py
+21 -7
View File
@@ -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]
src/core/ingest.py
+18 -2
View File
@@ -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):
src/core/search.py
+4 -1
View File
@@ -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:
src/core/security.py
+28
View File
@@ -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
src/server/app.py
+73 -32
View File
@@ -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:
src/server/auth.py
+12 -3
View File
@@ -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
src/server/deps.py
+25 -13
View File
@@ -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
tests/test_deps.py
+97
View File
@@ -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
tests/test_security.py
+61
View File
@@ -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(".") → ""