md-vector-db/docs/superpowers/plans/2026-07-05-audit-fixes.md

审计问题修复计划

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) to implement this plan task-by-task.

Goal: 修复架构审计报告中 Phase 1+2 的全部严重和高风险问题，以及部分中风险问题。

Architecture: 按影响面从底层到上层依次修复：config → embedder → db → ingest → search → server → cli。每层修复后运行测试验证不引入回归。

Tech Stack: Python 3.13, FastAPI, ChromaDB, sentence-transformers

---

修复范围

审计编号	问题	本次修复
S-01	API 密钥明文	✅
S-02	路径遍历	✅
S-03	无认证	✅
S-05	错误信息泄露	✅
C-01	并发不安全	✅
C-02	全局环境变量污染	✅
C-04	top_k 无上界	✅
C-05	大文件 OOM	✅
C-06	配置路径硬编码	✅
C-07	假健康检查	✅
D-02	全局单例	✅
D-03	同步嵌入阻塞	✅
M-02	无日志	✅
T-02	DELETE 端点无测试	✅

---

Task 1: 配置安全修复 (S-01, C-06)

Files:

• Modify: config.yaml — 移除 api_key 字段

• Modify: src/core/config.py — api_key 默认从环境变量读取

• Modify: src/cli/main.py:14 — 移除重复 sys.path 操作（顺带）

• Add Test: tests/test_config.py — 补充 API key 环境变量测试

• Step 1: 补充测试 — tests/test_config.py 追加

class TestEmbedConfigEnvVar:
    """api_key 从环境变量读取."""

    def test_api_key_from_env(self, monkeypatch):
        monkeypatch.setenv("EMBED_API_KEY", "sk-env-test")
        from src.core.config import EmbedConfig
        cfg = EmbedConfig(mode="api")
        assert cfg.api_key == "sk-env-test"

    def test_api_key_empty_when_not_set(self, monkeypatch):
        monkeypatch.delenv("EMBED_API_KEY", raising=False)
        from src.core.config import EmbedConfig
        cfg = EmbedConfig(mode="api")
        assert cfg.api_key == ""

• Step 2: 运行测试确认失败

cd D:/Code/doing_exercises/programs/md-vector-db && uv run pytest tests/test_config.py -v

Expected: 新测试 FAIL

• Step 3: 修改 config.yaml — 删除 api_key 行

embed:
  mode: local           # local | api
  local_model: BAAI/bge-small-zh-v1.5
  api_base: ""          # api 模式下填写
  # api_key 请通过环境变量 EMBED_API_KEY 设置

• Step 4: 修改 config.py — api_key 默认从环境变量读取

将 EmbedConfig 的 api_key 字段改为:

@dataclass
class EmbedConfig:
    """嵌入模型配置."""
    mode: str = "local"
    local_model: str = "BAAI/bge-small-zh-v1.5"
    api_base: str = ""
    api_key: str = field(default_factory=lambda: os.environ.get("EMBED_API_KEY", ""))

并在文件顶部加 import os。

• Step 5: 统一配置路径常量

在 config.py 末尾添加:

DEFAULT_CONFIG_PATH = "config.yaml"

其他文件中的 "config.yaml" 字符串引用改为 from src.core.config import DEFAULT_CONFIG_PATH。

修改 src/server/app.py:47,59 和 src/cli/main.py:26 中的硬编码 "config.yaml"。

• Step 6: 运行全部测试

cd D:/Code/doing_exercises/programs/md-vector-db && uv run pytest tests/ -v --tb=short

Expected: 33 passed (31 + 2 new)

• Step 7: Commit

git add -A && git commit -m "fix: remove api_key from config.yaml, read from env var; unify config path"

---

Task 2: 路径遍历防护 (S-02)

Files:

• Modify: src/server/app.py — ingest_document() 增加路径白名单

• Add Test: tests/test_api.py — 路径遍历攻击测试

• Step 1: 补充测试 — tests/test_api.py 追加

class TestSecurity:
    """安全测试."""

    def test_ingest_rejects_path_traversal(self, client):
        """拒绝路径遍历攻击."""
        response = client.post(
            "/api/v1/ingest",
            json={"file_path": "../../../etc/passwd", "file_name": "hack.md"},
        )
        assert response.status_code in (400, 403)

    def test_ingest_rejects_absolute_path(self, client):
        """拒绝绝对路径."""
        response = client.post(
            "/api/v1/ingest",
            json={"file_path": "C:\\Windows\\System32\\config\\SAM"},
        )
        assert response.status_code in (400, 403)

• Step 2: 运行测试确认失败

uv run pytest tests/test_api.py::TestSecurity -v

Expected: FAIL (当前无防护)

• Step 3: 修改 app.py — 路径校验函数

在 ingest_document() 最前面加:

import os as _os

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.post("/api/v1/ingest")
def ingest_document(req: IngestRequest):
    if req.file_path:
        if not _is_safe_path(req.file_path):
            raise HTTPException(status_code=400, detail="不允许的路径")
        path = Path(req.file_path)
        ...

• Step 4: 运行测试确认通过

uv run pytest tests/test_api.py -v

Expected: 8 passed (6 original + 2 new security tests)

• Step 5: Commit

git add -A && git commit -m "fix: add path traversal protection to ingest endpoint"

---

Task 3: API 认证 + 速率限制 + 安全错误消息 (S-03, S-04, S-05)

Files:

• Modify: src/server/app.py — 添加 API Key 中间件、速率限制、安全错误消息

• Create: src/server/auth.py — 认证依赖

• Add Test: tests/test_api.py — 认证/速率测试

• Step 1: 创建 auth.py

"""API 认证与安全中间件."""
import os
from fastapi import Header, HTTPException, Request
from fastapi.responses import JSONResponse
import time
from collections import defaultdict
import threading


# -- 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")
    return True


# -- 简易速率限制 --
class RateLimiter:
    """基于内存的简易速率限制器."""

    def __init__(self, max_requests: int = 30, window_seconds: int = 60):
        self.max_requests = max_requests
        self.window = window_seconds
        self._store: dict[str, list[float]] = defaultdict(list)
        self._lock = threading.Lock()

    def is_allowed(self, client_id: str) -> bool:
        now = time.time()
        with self._lock:
            records = self._store[client_id]
            # 清理过期记录
            records[:] = [t for t in records if now - t < self.window]
            if len(records) >= self.max_requests:
                return False
            records.append(now)
            return True

    async def __call__(self, request: Request):
        client_id = request.client.host if request.client else "unknown"
        if not self.is_allowed(client_id):
            raise HTTPException(status_code=429, detail="请求过于频繁，请稍后再试")
        return True


rate_limiter = RateLimiter(max_requests=30, window_seconds=60)

• Step 2: 修改 app.py — 集成认证和速率限制

在 app.py 中:

from src.server.auth import verify_api_key, rate_limiter
import logging

logger = logging.getLogger("md-vector-db")

# 速率限制中间件
@app.middleware("http")
async def rate_limit_middleware(request: Request, call_next):
    await rate_limiter(request)
    response = await call_next(request)
    return response


@app.post("/api/v1/ingest")
def ingest_document(req: IngestRequest, _: bool = Depends(verify_api_key)):
    ...
    except HTTPException:
        raise
    except Exception as e:
        logger.exception("入库失败: %s", e)
        raise HTTPException(status_code=500, detail="服务器内部错误")


@app.post("/api/v1/search")
def search_documents(req: SearchRequest, _: bool = Depends(verify_api_key)):
    ...


@app.delete("/api/v1/documents/{file_name}")
def delete_document(file_name: str, _: bool = Depends(verify_api_key)):
    ...

• Step 3: 运行测试确认通过

uv run pytest tests/test_api.py -v

Expected: 全部通过 (未设 API Key 环境变量时认证跳过)

• Step 4: Commit

git add -A && git commit -m "fix: add API key auth, rate limiting, and safe error messages"

---

Task 4: 请求校验 (C-04, M-03)

Files:

• Modify: src/server/app.py — SearchRequest.top_k 加 Pydantic Field 校验；IngestRequest.file_name 校验

• Step 1: 修改请求模型

from pydantic import BaseModel, Field

class IngestRequest(BaseModel):
    file_path: str | None = None
    content: str | None = None
    file_name: str | None = Field(default=None, max_length=255, pattern=r"^[^\\/:*?\"<>|]+\.md$")


class SearchRequest(BaseModel):
    query: str = Field(..., min_length=1, max_length=2000)
    top_k: int = Field(default=10, ge=1, le=100)

• Step 2: 运行测试

uv run pytest tests/test_api.py -v

Expected: 全部通过

• Step 3: Commit

git add -A && git commit -m "fix: add field validation to API request models"

---

Task 5: ChromaDB 线程安全 (C-01)

Files:

• Modify: src/core/db.py — VectorDB 写操作加锁

• Step 1: 修改 db.py

"""ChromaDB 数据库层."""
import threading
import chromadb
from chromadb.api.models.Collection import Collection


class VectorDB:
    """线程安全的向量数据库封装."""

    def __init__(self, persist_dir: str = "./data"):
        self.client = chromadb.PersistentClient(path=persist_dir)
        self._write_lock = threading.Lock()

    def get_or_create_collection(self, name: str) -> Collection:
        """获取或创建 collection."""
        return self.client.get_or_create_collection(name=name)

    def delete_collection(self, name: str) -> None:
        """删除 collection (线程安全)."""
        with self._write_lock:
            try:
                self.client.delete_collection(name=name)
            except ValueError:
                pass

    def close(self) -> None:
        """释放数据库连接."""
        self.client.close()

    @property
    def write_lock(self) -> threading.Lock:
        """获取写锁，供外部在 add/delete/update 操作时使用."""
        return self._write_lock

• Step 2: 修改 ingest.py — 使用写锁

class DocumentIngestor:
    ...
    def ingest_content(self, content: str, file_name: str) -> int:
        ...
        with self.db.write_lock:
            self.collection.add(
                ids=ids, embeddings=embeddings,
                documents=texts, metadatas=metadatas,
            )
        return len(chunks)

• Step 3: 修改 search.py — delete 使用写锁

def delete_by_source(self, file_name: str) -> bool:
    ...
    with self.db.write_lock:
        existing = self.collection.get(where={"source_file": file_name})
        if existing and existing["ids"]:
            self.collection.delete(ids=existing["ids"])
            return True
    return False

• Step 4: 运行测试

uv run pytest tests/ -v --tb=short

Expected: 31 passed

• Step 5: Commit

git add -A && git commit -m "fix: add thread-safe write lock to ChromaDB operations"

---

Task 6: Embedder 重构 (C-02, C-05, D-01)

Files:

• Modify: src/core/embedder.py — 拆分为策略模式；移除全局 env var 副作用；加分批嵌入

• Modify: tests/test_embedder.py — 补充 API 模式测试

• Step 1: 重写 embedder.py

"""嵌入模型抽象层 — 策略模式."""
import os
import logging
from dataclasses import dataclass
from typing import Protocol

from src.core.config import EmbedConfig

logger = logging.getLogger(__name__)

# 模型下载镜像
HF_MIRROR = "https://hf-mirror.com"


class Embedder(Protocol):
    """嵌入器接口."""
    @property
    def dimension(self) -> int: ...
    def embed(self, texts: list[str]) -> list[list[float]]: ...


class LocalEmbedder:
    """本地 sentence-transformers 模型嵌入器."""

    def __init__(self, config: EmbedConfig):
        from sentence_transformers import SentenceTransformer
        self._config = config
        try:
            self._model = SentenceTransformer(
                config.local_model, local_files_only=True
            )
        except Exception:
            logger.info("模型未缓存, 通过镜像下载 %s", config.local_model)
            os.environ["HF_ENDPOINT"] = HF_MIRROR
            self._model = SentenceTransformer(config.local_model)
            os.environ.pop("HF_ENDPOINT", None)

    @property
    def dimension(self) -> int:
        return self._model.get_embedding_dimension()

    def embed(self, texts: list[str]) -> list[list[float]]:
        if not texts:
            raise ValueError("文本列表不能为空")
        embeddings = self._model.encode(texts, normalize_embeddings=True)
        return embeddings.tolist()


class APIEmbedder:
    """OpenAI 兼容 API 嵌入器."""

    def __init__(self, config: EmbedConfig):
        self._api_base = config.api_base
        self._api_key = config.api_key

    @property
    def dimension(self) -> int:
        return 1536

    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="text-embedding-3-small", input=texts,
        )
        return [d.embedding for d in response.data]


def create_embedder(config: EmbedConfig) -> Embedder:
    """工厂函数."""
    if config.mode == "local":
        return LocalEmbedder(config)
    elif config.mode == "api":
        return APIEmbedder(config)
    else:
        raise ValueError(f"不支持的嵌入模式: {config.mode}")


def batch_embed(embedder: Embedder, texts: list[str], batch_size: int = 32) -> list[list[float]]:
    """分批嵌入，避免一次性传入过多文本导致 OOM."""
    all_embeddings = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        all_embeddings.extend(embedder.embed(batch))
    return all_embeddings

• Step 2: 修改 ingest.py — 使用 batch_embed

将 ingest_content() 中的:

embeddings = self.embedder.embed(texts)

改为:

from src.core.embedder import batch_embed
embeddings = batch_embed(self.embedder, texts)

• Step 3: 补充 API 模式测试

在 tests/test_embedder.py 追加:

class TestAPIEmbedder:
    """API 嵌入器测试."""

    def test_api_embedder_init(self):
        from src.core.embedder import APIEmbedder
        cfg = EmbedConfig(mode="api", api_base="https://api.test.com", api_key="sk-test")
        emb = APIEmbedder(cfg)
        assert emb.dimension == 1536

    def test_api_embedder_empty_raises(self):
        from src.core.embedder import APIEmbedder
        cfg = EmbedConfig(mode="api")
        emb = APIEmbedder(cfg)
        with pytest.raises(ValueError):
            emb.embed([])

class TestBatchEmbed:
    """分批嵌入测试."""

    def test_batch_embed(self):
        from src.core.embedder import batch_embed, create_embedder
        embedder = create_embedder(EmbedConfig(mode="local"))
        texts = ["测试文本"] * 70  # 超过 batch_size 32
        results = batch_embed(embedder, texts, batch_size=32)
        assert len(results) == 70
        assert len(results[0]) == embedder.dimension

同时把原来的 TestEmbedder 中的 create_embedder(EmbedConfig(mode="local")) 改为检查返回类型是 LocalEmbedder。

• Step 4: 运行测试

uv run pytest tests/test_embedder.py tests/test_ingest.py tests/test_search.py tests/test_api.py -v

Expected: 全部通过

• Step 5: Commit

git add -A && git commit -m "refactor: strategy-pattern embedder, remove global env vars, add batch embedding"

---

Task 7: 服务层重构 (D-02, D-03, C-07)

Files:

• Modify: src/server/app.py — FastAPI Depends 注入替代全局单例；异步嵌入；真实健康检查

• Create: src/server/deps.py — 依赖注入工厂

• Step 1: 创建 deps.py

"""FastAPI 依赖注入."""
import os
import functools
import logging

from src.core.config import load_config, EmbedConfig
from src.core.db import VectorDB
from src.core.embedder import create_embedder
from src.core.ingest import DocumentIngestor
from src.core.search import Searcher

logger = logging.getLogger(__name__)


class AppState:
    """应用级共享状态 (替代模块级全局变量)."""

    def __init__(self):
        config_path = os.environ.get("MD_VECTOR_CONFIG", "config.yaml")
        self.config = load_config(config_path)

        data_dir = os.environ.get("MD_VECTOR_DB_DATA_DIR", self.config.chroma.persist_dir)
        self.db = VectorDB(persist_dir=data_dir)

        self.embedder = create_embedder(self.config.embed)

        collection = os.environ.get("MD_VECTOR_DB_COLLECTION", self.config.chroma.collection_name)
        self.searcher = Searcher(self.db, self.embedder, collection)
        self.ingestor = DocumentIngestor(self.db, self.embedder, collection)

    def is_healthy(self) -> dict:
        """真实健康检查."""
        status = {"status": "ok", "checks": {}}
        try:
            count = self.searcher.collection.count()
            status["checks"]["chromadb"] = {"status": "ok", "count": count}
        except Exception as e:
            status["checks"]["chromadb"] = {"status": "error", "detail": str(e)}
            status["status"] = "degraded"
        try:
            _ = self.embedder.dimension
            status["checks"]["embedder"] = {"status": "ok"}
        except Exception as e:
            status["checks"]["embedder"] = {"status": "error", "detail": str(e)}
            status["status"] = "degraded"
        return status


_state: AppState | None = None


def get_state() -> AppState:
    """获取应用状态单例."""
    global _state
    if _state is None:
        logger.info("初始化应用状态...")
        _state = AppState()
    return _state

• Step 2: 重写 app.py — 使用 Depends 注入

"""FastAPI 服务层."""
import logging
from pathlib import Path

from fastapi import FastAPI, HTTPException, Depends, Request
from fastapi.responses import RedirectResponse, JSONResponse
from pydantic import BaseModel, Field

from src.server.auth import verify_api_key, rate_limiter
from src.server.deps import get_state, AppState

logger = logging.getLogger(__name__)

# -- 请求模型 --
class IngestRequest(BaseModel):
    file_path: str | None = None
    content: str | None = None
    file_name: str | None = Field(default=None, max_length=255, pattern=r"^[^\\/:*?\"<>|]+\.md$")

class SearchRequest(BaseModel):
    query: str = Field(..., min_length=1, max_length=2000)
    top_k: int = Field(default=10, ge=1, le=100)

# -- App --
app = FastAPI(title="md-vector-db", version="0.1.0")


@app.middleware("http")
async def rate_limit_middleware(request: Request, call_next):
    await rate_limiter(request)
    return await call_next(request)


@app.get("/")
def root():
    return RedirectResponse(url="/docs")


@app.get("/api/v1/health")
def health(state: AppState = Depends(get_state)):
    return state.is_healthy()


@app.get("/api/v1/collections")
def list_collections(state: AppState = Depends(get_state)):
    info = state.searcher.get_collection_info()
    sources = state.searcher.list_sources()
    return {"collections": [info], "sources": sources}


@app.post("/api/v1/ingest")
def ingest_document(req: IngestRequest, state: AppState = Depends(get_state),
                    _: bool = Depends(verify_api_key)):
    try:
        if req.file_path:
            path = Path(req.file_path).resolve()
            if not str(path).startswith(str(Path.cwd().resolve())):
                raise HTTPException(status_code=400, detail="不允许访问该路径")
            if not path.exists():
                raise HTTPException(status_code=404, detail=f"文件不存在: {path.name}")
            count = state.ingestor.ingest_file(str(path))
            file_name = path.name
        elif req.content:
            file_name = req.file_name or "untitled.md"
            count = state.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}
    except HTTPException:
        raise
    except Exception:
        logger.exception("入库失败")
        raise HTTPException(status_code=500, detail="服务器内部错误")


@app.post("/api/v1/search")
def search_documents(req: SearchRequest, state: AppState = Depends(get_state),
                     _: bool = Depends(verify_api_key)):
    results = state.searcher.search(req.query, top_k=req.top_k)
    return {"results": results}


@app.delete("/api/v1/documents/{file_name}")
def delete_document(file_name: str, state: AppState = Depends(get_state),
                    _: bool = Depends(verify_api_key)):
    deleted = state.searcher.delete_by_source(file_name)
    if not deleted:
        raise HTTPException(status_code=404, detail=f"文档不存在: {file_name}")
    return {"status": "ok", "file": file_name}

• Step 3: 运行全部测试

uv run pytest tests/ -v --tb=short

Expected: 全部通过

• Step 4: Commit

git add -A && git commit -m "refactor: FastAPI Depends injection, async-safe, real health check"

---

Task 8: 测试补充与回归 (T-02, D-05, T-03)

Files:

• Modify: tests/test_api.py — 补充 DELETE 端点测试、ingest_file 测试

• Modify: tests/test_ingest.py — 补充 DocumentIngestor 集成测试

• Step 1: 补充 API 测试

在 tests/test_api.py 追加:

class TestDeleteEndpoint:
    """删除端点."""

    def test_delete_nonexistent(self, client):
        response = client.delete("/api/v1/documents/nonexistent.md")
        assert response.status_code == 404

    def test_delete_ingested(self, client):
        client.post("/api/v1/ingest", json={"content": "# Test", "file_name": "del.md"})
        response = client.delete("/api/v1/documents/del.md")
        assert response.status_code == 200
        # 删除后搜索不应返回结果
        search_resp = client.post("/api/v1/search", json={"query": "Test", "top_k": 3})
        results = search_resp.json()["results"]
        sources = [r["source_file"] for r in results]
        assert "del.md" not in sources

• Step 2: 补充 ingest 集成测试

在 tests/test_ingest.py 追加:

class TestIngestorIntegration:
    """入库器集成测试 (使用真实 embedder)."""

    def test_ingest_content_real(self, tmp_path):
        import tempfile
        from src.core.config import EmbedConfig
        from src.core.db import VectorDB
        from src.core.embedder import create_embedder
        from src.core.ingest import DocumentIngestor

        db = VectorDB(persist_dir=str(tmp_path))
        embedder = create_embedder(EmbedConfig(mode="local"))
        ingestor = DocumentIngestor(db, embedder, "test_integration")

        count = ingestor.ingest_content("# Hello\nWorld.", "hello.md")
        assert count > 0
        assert ingestor.collection.count() == count

    def test_ingest_deduplicates(self, tmp_path):
        """重复入库同一文件会去重."""
        from src.core.config import EmbedConfig
        from src.core.db import VectorDB
        from src.core.embedder import create_embedder
        from src.core.ingest import DocumentIngestor

        db = VectorDB(persist_dir=str(tmp_path))
        embedder = create_embedder(EmbedConfig(mode="local"))
        ingestor = DocumentIngestor(db, embedder, "test_dedup")

        c1 = ingestor.ingest_content("# A", "dup.md")
        c2 = ingestor.ingest_content("# B", "dup.md")
        # 第二次入库应覆盖第一次，总数不翻倍
        assert ingestor.collection.count() == c2

• Step 3: 运行全部测试

uv run pytest tests/ -v --tb=short

Expected: 35+ passed

• Step 4: 最终 Commit

git add -A && git commit -m "test: add DELETE endpoint tests and ingest integration tests"

---

Task 9: 日志系统 + README 更新

Files:

• Modify: src/server/app.py — 启动时配置 logging

• Modify: src/cli/main.py — 添加 logging

• Modify: README.md — 更新 API 文档

• Step 1: 在 app.py 顶部配置日志

import logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    handlers=[logging.StreamHandler()]
)

• Step 2: 运行测试确认无回归

uv run pytest tests/ -v --tb=short

• Step 3: Commit

git add -A && git commit -m "feat: add structured logging and update README"

---

最终验证

cd D:/Code/doing_exercises/programs/md-vector-db
uv run pytest tests/ -v
uv run python -m src.cli.main search "测试"