fix: add API key auth, rate limiting, and safe error messages

This commit is contained in:
2026-07-05 01:33:02 +08:00
parent 2ef0cf4a4e
commit 157155c0d2
2 changed files with 66 additions and 6 deletions
+19 -6
View File
@@ -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:
+47
View File
@@ -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)