feat: 添加登录请求速率限制中间件，防止暴力破解密码

Author: xerrors

server/main.py

@@ -1,6 +1,11 @@
+import asyncio
+import time
+from collections import defaultdict, deque
+
 import uvicorn
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, status
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
 from starlette.middleware.base import BaseHTTPMiddleware
 
 from server.routers import router
@@ -11,6 +16,14 @@
 # 设置日志配置
 setup_logging()
 
+RATE_LIMIT_MAX_ATTEMPTS = 10
+RATE_LIMIT_WINDOW_SECONDS = 60
+RATE_LIMIT_ENDPOINTS = {("/api/auth/token", "POST")}
+
+# In-memory login attempt tracker to reduce brute-force exposure per worker
+_login_attempts: defaultdict[str, deque[float]] = defaultdict(deque)
+_attempt_lock = asyncio.Lock()
+
 app = FastAPI()
 app.include_router(router, prefix="/api")
 
@@ -24,6 +37,51 @@
 )
 
 
+def _extract_client_ip(request: Request) -> str:
+    forwarded_for = request.headers.get("x-forwarded-for")
+    if forwarded_for:
+        return forwarded_for.split(",")[0].strip()
+    if request.client:
+        return request.client.host
+    return "unknown"
+
+
+class LoginRateLimitMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        normalized_path = request.url.path.rstrip("/") or "/"
+        request_signature = (normalized_path, request.method.upper())
+
+        if request_signature in RATE_LIMIT_ENDPOINTS:
+            client_ip = _extract_client_ip(request)
+            now = time.monotonic()
+
+            async with _attempt_lock:
+                attempt_history = _login_attempts[client_ip]
+
+                while attempt_history and now - attempt_history[0] > RATE_LIMIT_WINDOW_SECONDS:
+                    attempt_history.popleft()
+
+                if len(attempt_history) >= RATE_LIMIT_MAX_ATTEMPTS:
+                    retry_after = int(max(1, RATE_LIMIT_WINDOW_SECONDS - (now - attempt_history[0])))
+                    return JSONResponse(
+                        status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+                        content={"detail": "登录尝试过于频繁，请稍后再试"},
+                        headers={"Retry-After": str(retry_after)},
+                    )
+
+                attempt_history.append(now)
+
+            response = await call_next(request)
+
+            if response.status_code < 400:
+                async with _attempt_lock:
+                    _login_attempts.pop(client_ip, None)
+
+            return response
+
+        return await call_next(request)
+
+
 # 鉴权中间件
 class AuthMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next):
@@ -58,6 +116,7 @@ async def dispatch(self, request: Request, call_next):
 
 
 # 添加鉴权中间件
+app.add_middleware(LoginRateLimitMiddleware)
 app.add_middleware(AuthMiddleware)
 
 

test/bruteforce_simulation.py

@@ -0,0 +1,123 @@
+"""
+Quick script to hammer the login endpoint with invalid credentials and observe rate limiting.
+
+Usage:
+    uv run python test/bruteforce_simulation.py --username demo --attempts 20
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import os
+import random
+import string
+import sys
+import time
+from collections import Counter
+
+import httpx
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Simulate brute-force login attempts.")
+    parser.add_argument("--base-url", default=os.getenv("TEST_BASE_URL", "http://localhost:5050"), help="API base URL")
+    parser.add_argument("--username", default=os.getenv("TEST_USERNAME", "admin"), help="Login identifier to attack")
+    parser.add_argument(
+        "--attempts", type=int, default=20, help="Total number of attempts to issue (default: 20)"
+    )
+    parser.add_argument(
+        "--concurrency",
+        type=int,
+        default=4,
+        help="Concurrent request limit (default: 4)",
+    )
+    parser.add_argument(
+        "--delay",
+        type=float,
+        default=0.05,
+        help="Delay in seconds between scheduling attempts (default: 0.05)",
+    )
+    parser.add_argument(
+        "--password",
+        default=None,
+        help="Explicit password to reuse for each request; random values are used when omitted",
+    )
+    return parser.parse_args()
+
+
+def build_payload(username: str, password: str) -> dict[str, str]:
+    return {"username": username, "password": password}
+
+
+def random_password() -> str:
+    base = string.ascii_letters + string.digits + "!@#$%^&*"
+    return "".join(random.choices(base, k=12))
+
+
+async def attempt_login(
+    client: httpx.AsyncClient,
+    semaphore: asyncio.Semaphore,
+    attempt_no: int,
+    username: str,
+    static_password: str | None,
+) -> tuple[int, float]:
+    async with semaphore:
+        password = static_password or random_password()
+        payload = build_payload(username, password)
+        started = time.perf_counter()
+        response = await client.post("/api/auth/token", data=payload)
+        elapsed = time.perf_counter() - started
+        detail = response.json().get("detail") if response.headers.get("content-type", "").startswith("application/json") else response.text
+        print(
+            f"[{attempt_no:02d}] {response.status_code} in {elapsed*1000:.1f} ms "
+            f"(pwd={password!r}) detail={detail!r}"
+        )
+        return response.status_code, elapsed
+
+
+async def run_simulation(args: argparse.Namespace) -> int:
+    timeout = httpx.Timeout(10.0, connect=3.0)
+    limits = httpx.Limits(max_connections=args.concurrency, max_keepalive_connections=args.concurrency)
+    semaphore = asyncio.Semaphore(args.concurrency)
+    status_counts: Counter[int] = Counter()
+
+    async with httpx.AsyncClient(base_url=args.base_url.rstrip("/"), timeout=timeout, limits=limits) as client:
+        tasks = []
+        for attempt_no in range(1, args.attempts + 1):
+            tasks.append(
+                asyncio.create_task(
+                    attempt_login(client, semaphore, attempt_no, args.username, args.password)
+                )
+            )
+            if args.delay:
+                await asyncio.sleep(args.delay)
+
+        for task in asyncio.as_completed(tasks):
+            status_code, _ = await task
+            status_counts[status_code] += 1
+
+    print("\nSummary:")
+    for code, total in sorted(status_counts.items()):
+        print(f"  HTTP {code}: {total} hits")
+
+    if 429 in status_counts:
+        print("Rate limiting engaged (HTTP 429 encountered).")
+        return 0
+
+    print("No rate limiting observed; consider increasing attempt count or reducing delay.")
+    return 1
+
+
+def main() -> None:
+    args = parse_args()
+    try:
+        exit_code = asyncio.run(run_simulation(args))
+    except KeyboardInterrupt:
+        print("\nInterrupted.")
+        exit_code = 130
+    sys.exit(exit_code)
+
+
+if __name__ == "__main__":
+    main()