From 157155c0d27c543d068c671e27e0387433b8ea1e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E8=88=AA=E5=AE=87?= <3364451258@qq.com>
Date: Sun, 5 Jul 2026 01:33:02 +0800
Subject: [PATCH] fix: add API key auth, rate limiting, and safe error messages

---
 src/server/app.py  | 25 ++++++++++++++++++------
 src/server/auth.py | 47 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 6 deletions(-)
 create mode 100644 src/server/auth.py

diff --git a/src/server/app.py b/src/server/app.py
index f77b78a..bfc5747 100644
--- a/src/server/app.py
+++ b/src/server/app.py
@@ -1,8 +1,10 @@
 """FastAPI 服务层."""
 import os
+import logging
 from pathlib import Path
 
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Depends, Request
+from fastapi.responses import RedirectResponse
 from pydantic import BaseModel
 
 from src.core.config import load_config, EmbedConfig, DEFAULT_CONFIG_PATH
@@ -10,6 +12,9 @@ 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
+from src.server.auth import verify_api_key, rate_limiter
+
+logger = logging.getLogger("md-vector-db")
 
 
 def _is_safe_path(path_str: str) -> bool:
@@ -94,6 +99,13 @@ app = FastAPI(
 )
 
 
+@app.middleware("http")
+async def rate_limit_middleware(request: Request, call_next):
+    await rate_limiter(request)
+    response = await call_next(request)
+    return response
+
+
 @app.get("/")
 def root():
     """根路径重定向到 API 文档."""
@@ -115,7 +127,7 @@ def list_collections():
 
 
 @app.post("/api/v1/ingest")
-def ingest_document(req: IngestRequest):
+def ingest_document(req: IngestRequest, _: bool = Depends(verify_api_key)):
     ingestor = _get_ingestor()
     try:
         if req.file_path:
@@ -136,19 +148,20 @@ def ingest_document(req: IngestRequest):
         return {"status": "ok", "chunks": count, "file": file_name}
     except HTTPException:
         raise
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+    except Exception:
+        logger.exception("入库失败")
+        raise HTTPException(status_code=500, detail="服务器内部错误")
 
 
 @app.post("/api/v1/search")
-def search_documents(req: SearchRequest):
+def search_documents(req: SearchRequest, _: bool = Depends(verify_api_key)):
     searcher = _get_searcher()
     results = 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):
+def delete_document(file_name: str, _: bool = Depends(verify_api_key)):
     searcher = _get_searcher()
     deleted = searcher.delete_by_source(file_name)
     if not deleted:
diff --git a/src/server/auth.py b/src/server/auth.py
new file mode 100644
index 0000000..691919b
--- /dev/null
+++ b/src/server/auth.py
@@ -0,0 +1,47 @@
+"""API 认证与安全中间件."""
+import os
+import time
+import threading
+from collections import defaultdict
+
+from fastapi import Header, HTTPException, Request
+
+# -- 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)