feat: 增加登录限流和跨进程退出机制

ai_knowledge_base_v2/apps/backend/app/services/admin_service.py

@@ -13,6 +13,7 @@ from app.models.chat import ChatMessage, ChatSession
 from app.models.knowledge import Knowledge
 from app.models.logs import AiRequestLog, OperationLog
 from app.models.user import User
+from app.services.security_state_service import SecurityStateService
 
 
 DEVELOPMENT_ENVS = {"local", "dev", "development", "docker", "test", "testing"}
@@ -55,13 +56,27 @@ def _get_bootstrap_admin_config() -> tuple[str, str, str]:
 
 class AdminAuthService:
     @classmethod
-    def login(cls, db: Session, username: str, password: str) -> dict:
+    def login(cls, db: Session, username: str, password: str, ip: str = "unknown") -> dict:
+        failure_key = f"auth:admin:fail:{username.lower()}:{ip}"
+        SecurityStateService.enforce_limit(
+            f"rate:admin_login:ip:{ip}", limit=30, window_seconds=600, message="登录请求过于频繁，请稍后再试"
+        )
+        SecurityStateService.ensure_not_locked(
+            failure_key, limit=5, message="账号或当前网络连续登录失败次数过多，请十五分钟后再试"
+        )
         cls.ensure_bootstrap_admin(db)
         admin = db.scalar(select(Admin).where(Admin.username == username))
         if admin is None or not verify_password(password, admin.password):
+            SecurityStateService.record_failure(
+                failure_key,
+                limit=5,
+                lock_seconds=900,
+                message="账号或当前网络连续登录失败次数过多，请十五分钟后再试",
+            )
             raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="管理员账号或密码错误")
         if admin.status != 1:
             raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="管理员已禁用")
+        SecurityStateService.clear(failure_key)
 
         admin.last_login_at = datetime.now(UTC).replace(tzinfo=None)
         db.add(admin)

ai_knowledge_base_v2/apps/backend/app/services/auth_service.py

@@ -10,21 +10,40 @@ from app.core.security import create_access_token
 from app.models.user import User
 from app.services.captcha_service import CaptchaService
 from app.services.sms_code_service import SmsCodeService
+from app.services.security_state_service import SecurityStateService
 
 
 class AuthService:
-    _revoked_token_jti: set[str] = set()
-
     @classmethod
-    def send_sms_code(cls, db: Session, phone: str, captcha_id: str, captcha_code: str) -> None:
+    def send_sms_code(cls, db: Session, phone: str, captcha_id: str, captcha_code: str, ip: str) -> None:
         CaptchaService.verify(captcha_id, captcha_code)
         user = cls._get_existing_user(db, phone)
         cls._ensure_user_can_login(user)
+        SecurityStateService.enforce_limit(
+            f"rate:sms:ip:{ip}", limit=20, window_seconds=3600, message="当前网络发送验证码过于频繁，请稍后再试"
+        )
+        SecurityStateService.enforce_limit(
+            f"rate:sms:phone:{phone}", limit=1, window_seconds=60, message="验证码发送过于频繁，请一分钟后再试"
+        )
         SmsCodeService.send_code(db, phone)
 
     @classmethod
     def login_with_sms(cls, db: Session, phone: str, code: str) -> dict:
-        SmsCodeService.verify_code(db, phone, code)
+        failure_key = f"auth:sms:fail:{phone}"
+        SecurityStateService.ensure_not_locked(
+            failure_key, limit=5, message="验证码连续错误次数过多，请十五分钟后重新获取"
+        )
+        try:
+            SmsCodeService.verify_code(db, phone, code)
+        except HTTPException:
+            SecurityStateService.record_failure(
+                failure_key,
+                limit=5,
+                lock_seconds=900,
+                message="验证码连续错误次数过多，请十五分钟后重新获取",
+            )
+            raise
+        SecurityStateService.clear(failure_key)
         user = cls._get_existing_user(db, phone)
         cls._ensure_user_can_login(user)
 
@@ -41,12 +60,12 @@ class AuthService:
     def logout(cls, token_payload: dict) -> None:
         jti = token_payload.get("jti")
         if jti:
-            cls._revoked_token_jti.add(str(jti))
+            SecurityStateService.revoke_token(str(jti), token_payload.get("exp"))
 
     @classmethod
     def is_token_revoked(cls, token_payload: dict) -> bool:
         jti = token_payload.get("jti")
-        return bool(jti and str(jti) in cls._revoked_token_jti)
+        return bool(jti and SecurityStateService.is_token_revoked(str(jti)))
 
     @classmethod
     def _get_existing_user(cls, db: Session, phone: str) -> User:

ai_knowledge_base_v2/apps/backend/app/services/security_state_service.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import UTC, datetime, timedelta
+from threading import Lock
+import time
+
+from fastapi import HTTPException, status
+
+from app.services.redis_client import get_sync_redis_client
+
+
+@dataclass
+class _Counter:
+    value: int
+    expires_at: float
+
+
+class SecurityStateService:
+    _counters: dict[str, _Counter] = {}
+    _lock = Lock()
+
+    @classmethod
+    def enforce_limit(cls, key: str, *, limit: int, window_seconds: int, message: str) -> None:
+        count = cls._increment(key, window_seconds)
+        if count > limit:
+            raise HTTPException(status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail=message)
+
+    @classmethod
+    def record_failure(cls, key: str, *, limit: int, lock_seconds: int, message: str) -> None:
+        count = cls._increment(key, lock_seconds)
+        if count >= limit:
+            raise HTTPException(status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail=message)
+
+    @classmethod
+    def ensure_not_locked(cls, key: str, *, limit: int, message: str) -> None:
+        if cls._get(key) >= limit:
+            raise HTTPException(status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail=message)
+
+    @classmethod
+    def clear(cls, key: str) -> None:
+        client = get_sync_redis_client()
+        if client is not None:
+            try:
+                client.delete(key)
+            except Exception:
+                pass
+        with cls._lock:
+            cls._counters.pop(key, None)
+
+    @classmethod
+    def revoke_token(cls, jti: str, expires_at: int | float | None) -> None:
+        now = datetime.now(UTC)
+        expiry = datetime.fromtimestamp(float(expires_at), UTC) if expires_at else now + timedelta(days=1)
+        ttl = max(1, int((expiry - now).total_seconds()))
+        key = f"auth:revoked:{jti}"
+        client = get_sync_redis_client()
+        if client is not None:
+            try:
+                client.set(key, "1", ex=ttl)
+                return
+            except Exception:
+                pass
+        with cls._lock:
+            cls._counters[key] = _Counter(1, time.monotonic() + ttl)
+
+    @classmethod
+    def is_token_revoked(cls, jti: str) -> bool:
+        return cls._get(f"auth:revoked:{jti}") > 0
+
+    @classmethod
+    def _increment(cls, key: str, ttl: int) -> int:
+        client = get_sync_redis_client()
+        if client is not None:
+            try:
+                pipe = client.pipeline()
+                pipe.incr(key)
+                pipe.expire(key, ttl, nx=True)
+                value, _ = pipe.execute()
+                return int(value)
+            except Exception:
+                pass
+
+        now = time.monotonic()
+        with cls._lock:
+            record = cls._counters.get(key)
+            if record is None or record.expires_at <= now:
+                record = _Counter(0, now + ttl)
+            record.value += 1
+            cls._counters[key] = record
+            return record.value
+
+    @classmethod
+    def _get(cls, key: str) -> int:
+        client = get_sync_redis_client()
+        if client is not None:
+            try:
+                value = client.get(key)
+                return int(value) if value is not None else 0
+            except Exception:
+                pass
+        now = time.monotonic()
+        with cls._lock:
+            record = cls._counters.get(key)
+            if record is None or record.expires_at <= now:
+                cls._counters.pop(key, None)
+                return 0
+            return record.value
+
+
+def client_ip(request) -> str:
+    forwarded = request.headers.get("x-forwarded-for", "").split(",", 1)[0].strip()
+    if forwarded:
+        return forwarded
+    return request.client.host if request.client else "unknown"