fix: 修复 44 个代码审查问题 (CRITICAL/HIGH/MEDIUM/LOW)
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:
+73
-32
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user