Store login verification codes in Redis

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

@@ -12,6 +12,8 @@ from uuid import uuid4
 from fastapi import HTTPException, status
 from PIL import Image, ImageDraw, ImageFilter, ImageFont
 
+from app.services.redis_client import get_sync_redis_client
+
 logger = logging.getLogger(__name__)
 
 CAPTCHA_EXPIRES_SECONDS = 300
@@ -34,10 +36,13 @@ class CaptchaService:
         cls._cleanup_expired()
         captcha_id = str(uuid4())
         code = _random_code()
-        cls._records[captcha_id] = CaptchaRecord(
-            code=code,
-            expires_at=datetime.now(UTC) + timedelta(seconds=CAPTCHA_EXPIRES_SECONDS),
-        )
+        if _set_redis_captcha(captcha_id, code):
+            cls._records.pop(captcha_id, None)
+        else:
+            cls._records[captcha_id] = CaptchaRecord(
+                code=code,
+                expires_at=datetime.now(UTC) + timedelta(seconds=CAPTCHA_EXPIRES_SECONDS),
+            )
         return {
             "captchaId": captcha_id,
             "imageBase64": _render_captcha_image(code),
@@ -46,6 +51,14 @@ class CaptchaService:
 
     @classmethod
     def verify(cls, captcha_id: str, captcha_code: str) -> None:
+        redis_available, redis_code = _get_redis_captcha(captcha_id)
+        if redis_available and redis_code is not None:
+            if redis_code.lower() != captcha_code.strip().lower():
+                raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="图形验证码错误")
+            _delete_redis_captcha(captcha_id)
+            cls._records.pop(captcha_id, None)
+            return
+
         record = cls._records.get(captcha_id)
         if record is None:
             raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="图形验证码不存在或已过期")
@@ -69,6 +82,43 @@ def _random_code() -> str:
     return "".join(random.SystemRandom().choice(alphabet) for _ in range(CAPTCHA_LENGTH))
 
 
+def _captcha_key(captcha_id: str) -> str:
+    return f"auth:captcha:{captcha_id}"
+
+
+def _set_redis_captcha(captcha_id: str, code: str) -> bool:
+    client = get_sync_redis_client()
+    if client is None:
+        return False
+    try:
+        return bool(client.set(_captcha_key(captcha_id), code, ex=CAPTCHA_EXPIRES_SECONDS))
+    except Exception:
+        logger.warning("Redis captcha write failed for %s", captcha_id, exc_info=True)
+        return False
+
+
+def _get_redis_captcha(captcha_id: str) -> tuple[bool, str | None]:
+    client = get_sync_redis_client()
+    if client is None:
+        return False, None
+    try:
+        value = client.get(_captcha_key(captcha_id))
+    except Exception:
+        logger.warning("Redis captcha read failed for %s", captcha_id, exc_info=True)
+        return False, None
+    return True, value if isinstance(value, str) else None
+
+
+def _delete_redis_captcha(captcha_id: str) -> None:
+    client = get_sync_redis_client()
+    if client is None:
+        return
+    try:
+        client.delete(_captcha_key(captcha_id))
+    except Exception:
+        logger.warning("Redis captcha delete failed for %s", captcha_id, exc_info=True)
+
+
 def _render_captcha_image(code: str) -> str:
     image = Image.new("RGB", (CAPTCHA_WIDTH, CAPTCHA_HEIGHT), "#f7fbfa")
     draw = ImageDraw.Draw(image)

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

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from functools import lru_cache
 
+from redis import Redis as SyncRedis
 from redis.asyncio import Redis
 
 from app.core.config import get_settings
@@ -15,6 +16,14 @@ def get_redis_client() -> Redis | None:
     return Redis.from_url(redis_url, decode_responses=True)
 
 
+@lru_cache(maxsize=1)
+def get_sync_redis_client() -> SyncRedis | None:
+    redis_url = get_settings().redis_url.strip()
+    if not redis_url:
+        return None
+    return SyncRedis.from_url(redis_url, decode_responses=True)
+
+
 async def ping_redis() -> bool:
     client = get_redis_client()
     if client is None:

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

@@ -13,6 +13,7 @@ from sqlalchemy.orm import Session
 
 from app.core.config import get_settings
 from app.models.ai_config import SystemConfig
+from app.services.redis_client import get_sync_redis_client
 
 logger = logging.getLogger(__name__)
 
@@ -46,16 +47,25 @@ class SmsCodeService:
         code = config.mock_sms_code if config.mock_sms_enabled else _generate_sms_code()
         if not config.mock_sms_enabled:
             _send_aliyun_sms(phone, code, config)
-        cls._codes[phone] = SmsCodeRecord(
-            code=code,
-            expires_at=datetime.now(UTC) + timedelta(minutes=config.sms_code_expire_minutes),
-        )
+        ttl_seconds = config.sms_code_expire_minutes * 60
+        if _set_redis_code(phone, code, ttl_seconds):
+            cls._codes.pop(phone, None)
+            return
+        cls._codes[phone] = SmsCodeRecord(code=code, expires_at=datetime.now(UTC) + timedelta(seconds=ttl_seconds))
 
     @classmethod
     def verify_code(cls, db: Session, phone: str, code: str) -> None:
         cls._validate_phone(phone)
-        record = cls._codes.get(phone)
         config = _load_sms_config(db)
+        redis_available, redis_code = _get_redis_code(phone)
+        if redis_available and redis_code is not None:
+            if redis_code != code:
+                raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="验证码错误")
+            _delete_redis_code(phone)
+            cls._codes.pop(phone, None)
+            return
+
+        record = cls._codes.get(phone)
         if config.mock_sms_enabled and record is None and code == config.mock_sms_code:
             return
         if record is None:
@@ -177,6 +187,43 @@ def _generate_sms_code() -> str:
     return f"{random.SystemRandom().randint(0, 999999):06d}"
 
 
+def _sms_code_key(phone: str) -> str:
+    return f"auth:sms_code:{phone}"
+
+
+def _set_redis_code(phone: str, code: str, ttl_seconds: int) -> bool:
+    client = get_sync_redis_client()
+    if client is None:
+        return False
+    try:
+        return bool(client.set(_sms_code_key(phone), code, ex=ttl_seconds))
+    except Exception:
+        logger.warning("Redis SMS code write failed for %s", _mask_phone(phone), exc_info=True)
+        return False
+
+
+def _get_redis_code(phone: str) -> tuple[bool, str | None]:
+    client = get_sync_redis_client()
+    if client is None:
+        return False, None
+    try:
+        value = client.get(_sms_code_key(phone))
+    except Exception:
+        logger.warning("Redis SMS code read failed for %s", _mask_phone(phone), exc_info=True)
+        return False, None
+    return True, value if isinstance(value, str) else None
+
+
+def _delete_redis_code(phone: str) -> None:
+    client = get_sync_redis_client()
+    if client is None:
+        return
+    try:
+        client.delete(_sms_code_key(phone))
+    except Exception:
+        logger.warning("Redis SMS code delete failed for %s", _mask_phone(phone), exc_info=True)
+
+
 def _config_text(values: dict[str, str], key: str, default: str) -> str:
     value = values.get(key)
     return value.strip() if value is not None else default