fix: add API key auth, rate limiting, and safe error messages
This commit is contained in:
+19
-6
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user