fix: add API key auth, rate limiting, and safe error messages
This commit is contained in:
+19
-6
@@ -1,8 +1,10 @@
|
|||||||
"""FastAPI 服务层."""
|
"""FastAPI 服务层."""
|
||||||
import os
|
import os
|
||||||
|
import logging
|
||||||
from pathlib import Path
|
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 pydantic import BaseModel
|
||||||
|
|
||||||
from src.core.config import load_config, EmbedConfig, DEFAULT_CONFIG_PATH
|
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.embedder import create_embedder
|
||||||
from src.core.ingest import DocumentIngestor
|
from src.core.ingest import DocumentIngestor
|
||||||
from src.core.search import Searcher
|
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:
|
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("/")
|
@app.get("/")
|
||||||
def root():
|
def root():
|
||||||
"""根路径重定向到 API 文档."""
|
"""根路径重定向到 API 文档."""
|
||||||
@@ -115,7 +127,7 @@ def list_collections():
|
|||||||
|
|
||||||
|
|
||||||
@app.post("/api/v1/ingest")
|
@app.post("/api/v1/ingest")
|
||||||
def ingest_document(req: IngestRequest):
|
def ingest_document(req: IngestRequest, _: bool = Depends(verify_api_key)):
|
||||||
ingestor = _get_ingestor()
|
ingestor = _get_ingestor()
|
||||||
try:
|
try:
|
||||||
if req.file_path:
|
if req.file_path:
|
||||||
@@ -136,19 +148,20 @@ def ingest_document(req: IngestRequest):
|
|||||||
return {"status": "ok", "chunks": count, "file": file_name}
|
return {"status": "ok", "chunks": count, "file": file_name}
|
||||||
except HTTPException:
|
except HTTPException:
|
||||||
raise
|
raise
|
||||||
except Exception as e:
|
except Exception:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
logger.exception("入库失败")
|
||||||
|
raise HTTPException(status_code=500, detail="服务器内部错误")
|
||||||
|
|
||||||
|
|
||||||
@app.post("/api/v1/search")
|
@app.post("/api/v1/search")
|
||||||
def search_documents(req: SearchRequest):
|
def search_documents(req: SearchRequest, _: bool = Depends(verify_api_key)):
|
||||||
searcher = _get_searcher()
|
searcher = _get_searcher()
|
||||||
results = searcher.search(req.query, top_k=req.top_k)
|
results = searcher.search(req.query, top_k=req.top_k)
|
||||||
return {"results": results}
|
return {"results": results}
|
||||||
|
|
||||||
|
|
||||||
@app.delete("/api/v1/documents/{file_name}")
|
@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()
|
searcher = _get_searcher()
|
||||||
deleted = searcher.delete_by_source(file_name)
|
deleted = searcher.delete_by_source(file_name)
|
||||||
if not deleted:
|
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