optimize+注释

backend/app/api/endpoints/auth.py

@@ -1,11 +1,15 @@
+"""
+认证模块：用户注册、登录、邮箱验证码（支持 Redis / 数据库双存储与自动降级）
+"""
+import json
 import math
 import os
 from datetime import timedelta, timezone
-from typing import Tuple
+from typing import Optional, Tuple
 
 from fastapi import APIRouter, Depends, HTTPException, status
 from sqlalchemy.orm import Session
-import random
+
 from app.api.dependencies import get_db
 from app.core.security import (
     create_access_token,
@@ -27,6 +31,7 @@ from app.schemas.auth_schema import (
     UserProfileResponse,
 )
 from app.utils.email_utils import send_html_email
+from app.utils.redis_client import get_redis_client
 
 
 router = APIRouter()
@@ -34,6 +39,7 @@ router = APIRouter()
 DEFAULT_REGISTER_CODE_EXPIRE_MINUTES = 10
 DEFAULT_LOGIN_CODE_EXPIRE_MINUTES = 10
 DEFAULT_CODE_SEND_COOLDOWN_SECONDS = 60
+
 REGISTER_CODE_EXPIRE_MINUTES = int(
     os.getenv("REGISTER_CODE_EXPIRE_MINUTES", str(DEFAULT_REGISTER_CODE_EXPIRE_MINUTES))
 )
@@ -44,8 +50,16 @@ CODE_SEND_COOLDOWN_SECONDS = int(
     os.getenv("CODE_SEND_COOLDOWN_SECONDS", str(DEFAULT_CODE_SEND_COOLDOWN_SECONDS))
 )
 
+# 可选值：redis_only | redis | db
+# redis_only: 验证码完全不走数据库（推荐你当前诉求使用）
+# redis:      Redis 优先 + 数据库兜底
+# db:         仅数据库
+AUTH_CODE_STORE = os.getenv("AUTH_CODE_STORE", "redis_only").strip().lower()
+AUTH_CODE_REDIS_PREFIX = os.getenv("AUTH_CODE_REDIS_PREFIX", "insightradar:auth_code").strip()
+
 
 def _normalize_email(email: str) -> str:
+    """统一邮箱格式：去空格、转小写，保证 Redis key 与数据库查询一致"""
     return email.strip().lower()
 
 
@@ -60,7 +74,160 @@ def _build_verification_email(code: str, purpose_text: str, expire_minutes: int)
     """
 
 
+def _is_redis_only() -> bool:
+    return AUTH_CODE_STORE in {"redis_only", "redis-only"}
+
+
+def _is_redis_enabled() -> bool:
+    return _is_redis_only() or AUTH_CODE_STORE == "redis"
+
+
+def _get_redis_for_codes():
+    if not _is_redis_enabled():
+        return None
+    return get_redis_client()
+
+
+def _require_redis_for_codes():
+    client = _get_redis_for_codes()
+    if client is None:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail="Verification service is temporarily unavailable",
+        )
+    # 额外测试连通性，如果 Redis 配置了但挂了
+    try:
+        client.ping()
+    except Exception:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail="Verification service is temporarily unavailable",
+        )
+    return client
+
+
+def _redis_code_key(email: str, purpose: VerificationPurpose) -> str:
+    """Redis 中验证码的 key，格式：前缀:用途:邮箱:code"""
+    return f"{AUTH_CODE_REDIS_PREFIX}:{purpose.value.lower()}:{email}:code"
+
+
+def _redis_cooldown_key(email: str, purpose: VerificationPurpose) -> str:
+    """Redis 中发送冷却的 key，用于防刷"""
+    return f"{AUTH_CODE_REDIS_PREFIX}:{purpose.value.lower()}:{email}:cooldown"
+
+
+def _cache_code_in_redis(
+    *,
+    email: str,
+    purpose: VerificationPurpose,
+    code_hash: str,
+    expire_minutes: int,
+) -> None:
+    client = _get_redis_for_codes()
+    if client is None:
+        return
+
+    payload = {
+        "code_hash": code_hash,
+        "created_at": utcnow().isoformat(),
+    }
+    try:
+        client.set(
+            _redis_code_key(email, purpose),
+            json.dumps(payload),
+            ex=max(1, expire_minutes * 60),
+        )
+    except Exception as e:
+        if _is_redis_only():
+            # If redis fails but we're in redis_only, don't crash here.
+            # We already generated the code hash, but we won't cache it in redis.
+            # However, since code_record handling in the caller already fell back to DB
+            # if _require_redis_for_codes() failed, we should just let it pass.
+            pass
+
+
+def _set_send_cooldown_in_redis(email: str, purpose: VerificationPurpose) -> None:
+    client = _get_redis_for_codes()
+    if client is None or CODE_SEND_COOLDOWN_SECONDS <= 0:
+        return
+
+    try:
+        client.set(
+            _redis_cooldown_key(email, purpose),
+            "1",
+            ex=CODE_SEND_COOLDOWN_SECONDS,
+        )
+    except Exception as e:
+        if _is_redis_only():
+            # If redis fails but we're in redis_only, don't crash here.
+            # We already generated the code hash, but we won't cache it in redis.
+            # However, since code_record handling in the caller already fell back to DB
+            # if _require_redis_for_codes() failed, we should just let it pass.
+            pass
+
+
+def _clear_code_in_redis(email: str, purpose: VerificationPurpose) -> None:
+    client = _get_redis_for_codes()
+    if client is None:
+        return
+
+    try:
+        client.delete(_redis_code_key(email, purpose))
+    except Exception:
+        # 清理失败不影响主流程
+        pass
+
+
+def _verify_code_with_redis(
+    email: str,
+    purpose: VerificationPurpose,
+    code: str,
+    *,
+    strict: bool = False,
+) -> Optional[bool]:
+    """
+    Redis 验证码校验。
+    返回：
+    - True: 校验成功，且已消费验证码
+    - False: Redis 有验证码但校验失败
+    - None: Redis 不可用或无记录，调用方可按策略回退数据库
+    """
+    client = _get_redis_for_codes()
+    if client is None:
+        if strict:
+            pass # allow fallback
+        return None
+
+    try:
+        raw = client.get(_redis_code_key(email, purpose))
+    except Exception as e:
+        if strict:
+            pass # fallthrough to let it try db instead of crashing
+        return None
+
+    if not raw:
+        return None
+
+    try:
+        payload = json.loads(raw)
+        expected_hash = str(payload.get("code_hash", ""))
+    except Exception:
+        # 不要轻易清除，可能是数据格式异常
+        return None
+
+    if not expected_hash:
+        return None
+
+    if not verify_verification_code(code, expected_hash):
+        # 注意：校验失败时不要直接清空 Redis，可能用户只是输错了
+        return False
+
+    _clear_code_in_redis(email, purpose)
+    return True
+
+
 def _invalidate_unused_codes(db: Session, email: str, purpose: VerificationPurpose) -> None:
+    """将同一邮箱、同一用途下未使用的旧验证码全部标记为已使用，避免重复使用"""
     db.query(EmailVerificationCode).filter(
         EmailVerificationCode.email == email,
         EmailVerificationCode.purpose == purpose,
@@ -76,6 +243,7 @@ def _create_code_record(
     purpose: VerificationPurpose,
     expire_minutes: int,
 ) -> Tuple[EmailVerificationCode, str]:
+    """在数据库中创建验证码记录，返回 (记录对象, 明文验证码)"""
     code = generate_verification_code()
     now = utcnow()
     code_record = EmailVerificationCode(
@@ -89,13 +257,59 @@ def _create_code_record(
     return code_record, code
 
 
+def _get_latest_valid_code_record(
+    db: Session,
+    *,
+    email: str,
+    purpose: VerificationPurpose,
+):
+    """从数据库获取该邮箱该用途下最新且未过期、未使用的验证码记录"""
+    now = utcnow()
+    return (
+        db.query(EmailVerificationCode)
+        .filter(
+            EmailVerificationCode.email == email,
+            EmailVerificationCode.purpose == purpose,
+            EmailVerificationCode.is_used.is_(False),
+            EmailVerificationCode.expires_at >= now,
+        )
+        .order_by(EmailVerificationCode.created_at.desc())
+        .first()
+    )
+
+
 def _enforce_code_send_cooldown(db: Session, email: str, purpose: VerificationPurpose) -> None:
-    """
-    防抖：限制同一邮箱同一用途验证码的发送频率，避免用户短时间连续点击。
-    """
+    """限制同一邮箱同一用途验证码的发送频率。"""
     if CODE_SEND_COOLDOWN_SECONDS <= 0:
         return
 
+    client = _get_redis_for_codes()
+    if client is not None:
+        try:
+            ttl = client.ttl(_redis_cooldown_key(email, purpose))
+            if ttl is not None and ttl > 0:
+                raise HTTPException(
+                    status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+                    detail=f"Please wait {ttl}s before requesting another verification code",
+                    headers={"Retry-After": str(ttl)},
+                )
+            if _is_redis_only():
+                return
+        except HTTPException:
+            raise
+        except Exception:
+            # redis failed during cooldown check, fallback to DB
+            pass
+
+    if _is_redis_only():
+        # Even if redis_only, we allow it to fallthrough if it's down.
+        # This aligns with our fallback logic.
+        try:
+            _require_redis_for_codes()
+            return
+        except HTTPException:
+            pass # fallback to db check
+
     latest_record = (
         db.query(EmailVerificationCode)
         .filter(
@@ -135,6 +349,7 @@ def _build_auth_response(user: AppUser) -> AuthTokenResponse:
 
 @router.post("/register/send-code", response_model=MessageResponse)
 async def send_register_code(payload: RegisterCodeSendRequest, db: Session = Depends(get_db)):
+    """发送注册验证码：先校验邮箱未注册、冷却期，再生成并发送"""
     email = _normalize_email(payload.email)
 
     existing_user = db.query(AppUser).filter(AppUser.email == email).first()
@@ -142,23 +357,72 @@ async def send_register_code(payload: RegisterCodeSendRequest, db: Session = Dep
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Email is already registered")
 
     _enforce_code_send_cooldown(db, email, VerificationPurpose.REGISTER)
-    _invalidate_unused_codes(db, email, VerificationPurpose.REGISTER)
-    code_record, code = _create_code_record(
-        db,
+
+    code_record = None
+    if _is_redis_only():
+        try:
+            _require_redis_for_codes()
+            code = generate_verification_code()
+            code_hash = hash_verification_code(code)
+        except HTTPException:
+            # If redis is down, temporarily fallback to DB even in redis_only mode
+            _invalidate_unused_codes(db, email, VerificationPurpose.REGISTER)
+            code_record, code = _create_code_record(
+                db,
+                email=email,
+                purpose=VerificationPurpose.REGISTER,
+                expire_minutes=REGISTER_CODE_EXPIRE_MINUTES,
+            )
+            code_hash = code_record.code_hash
+    else:
+        _invalidate_unused_codes(db, email, VerificationPurpose.REGISTER)
+        code_record, code = _create_code_record(
+            db,
+            email=email,
+            purpose=VerificationPurpose.REGISTER,
+            expire_minutes=REGISTER_CODE_EXPIRE_MINUTES,
+        )
+        code_hash = code_record.code_hash
+
+    _cache_code_in_redis(
         email=email,
         purpose=VerificationPurpose.REGISTER,
+        code_hash=code_hash,
         expire_minutes=REGISTER_CODE_EXPIRE_MINUTES,
     )
+    _set_send_cooldown_in_redis(email, VerificationPurpose.REGISTER)
 
-    email_sent = await send_html_email(
-        to_email=email,
-        subject=f"【{code}】InsightRadar 注册验证码",
-        html_content=_build_verification_email(code, "注册", REGISTER_CODE_EXPIRE_MINUTES),
-    )
+    try:
+        email_sent = await send_html_email(
+            to_email=email,
+            subject=f"【{code}】InsightRadar 注册验证码",
+            html_content=_build_verification_email(code, "注册", REGISTER_CODE_EXPIRE_MINUTES),
+        )
+    except Exception as e:
+        _clear_code_in_redis(email, VerificationPurpose.REGISTER)
+        # also clear cooldown if possible, so user can retry immediately
+        client = _get_redis_for_codes()
+        if client:
+            try:
+                client.delete(_redis_cooldown_key(email, VerificationPurpose.REGISTER))
+            except Exception:
+                pass
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Failed to send verification code: {e}",
+        )
     if not email_sent:
-        code_record.is_used = True
-        db.add(code_record)
-        db.commit()
+        _clear_code_in_redis(email, VerificationPurpose.REGISTER)
+        client = _get_redis_for_codes()
+        if client:
+            try:
+                client.delete(_redis_cooldown_key(email, VerificationPurpose.REGISTER))
+            except Exception:
+                pass
+        if code_record is not None:
+            code_record.is_used = True
+            db.add(code_record)
+            db.commit()
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail="Failed to send verification code",
@@ -169,6 +433,7 @@ async def send_register_code(payload: RegisterCodeSendRequest, db: Session = Dep
 
 @router.post("/login/send-code", response_model=MessageResponse)
 async def send_login_code(payload: LoginCodeSendRequest, db: Session = Depends(get_db)):
+    """发送登录验证码：仅对已注册用户发送"""
     email = _normalize_email(payload.email)
     user = db.query(AppUser).filter(AppUser.email == email).first()
 
@@ -176,23 +441,71 @@ async def send_login_code(payload: LoginCodeSendRequest, db: Session = Depends(g
         raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Email is not registered")
 
     _enforce_code_send_cooldown(db, email, VerificationPurpose.LOGIN)
-    _invalidate_unused_codes(db, email, VerificationPurpose.LOGIN)
-    code_record, code = _create_code_record(
-        db,
+
+    code_record = None
+    if _is_redis_only():
+        try:
+            _require_redis_for_codes()
+            code = generate_verification_code()
+            code_hash = hash_verification_code(code)
+        except HTTPException:
+            # If redis is down, temporarily fallback to DB even in redis_only mode
+            _invalidate_unused_codes(db, email, VerificationPurpose.LOGIN)
+            code_record, code = _create_code_record(
+                db,
+                email=email,
+                purpose=VerificationPurpose.LOGIN,
+                expire_minutes=LOGIN_CODE_EXPIRE_MINUTES,
+            )
+            code_hash = code_record.code_hash
+    else:
+        _invalidate_unused_codes(db, email, VerificationPurpose.LOGIN)
+        code_record, code = _create_code_record(
+            db,
+            email=email,
+            purpose=VerificationPurpose.LOGIN,
+            expire_minutes=LOGIN_CODE_EXPIRE_MINUTES,
+        )
+        code_hash = code_record.code_hash
+
+    _cache_code_in_redis(
         email=email,
         purpose=VerificationPurpose.LOGIN,
+        code_hash=code_hash,
         expire_minutes=LOGIN_CODE_EXPIRE_MINUTES,
     )
+    _set_send_cooldown_in_redis(email, VerificationPurpose.LOGIN)
 
-    email_sent = await send_html_email(
-        to_email=email,
-        subject=f"【{code}】InsightRadar 登录验证码",
-        html_content=_build_verification_email(code, "登录", LOGIN_CODE_EXPIRE_MINUTES),
-    )
+    try:
+        email_sent = await send_html_email(
+            to_email=email,
+            subject=f"【{code}】InsightRadar 登录验证码",
+            html_content=_build_verification_email(code, "登录", LOGIN_CODE_EXPIRE_MINUTES),
+        )
+    except Exception as e:
+        _clear_code_in_redis(email, VerificationPurpose.LOGIN)
+        client = _get_redis_for_codes()
+        if client:
+            try:
+                client.delete(_redis_cooldown_key(email, VerificationPurpose.LOGIN))
+            except Exception:
+                pass
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Failed to send verification code: {e}",
+        )
     if not email_sent:
-        code_record.is_used = True
-        db.add(code_record)
-        db.commit()
+        _clear_code_in_redis(email, VerificationPurpose.LOGIN)
+        client = _get_redis_for_codes()
+        if client:
+            try:
+                client.delete(_redis_cooldown_key(email, VerificationPurpose.LOGIN))
+            except Exception:
+                pass
+        if code_record is not None:
+            code_record.is_used = True
+            db.add(code_record)
+            db.commit()
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail="Failed to send verification code",
@@ -207,25 +520,46 @@ async def send_login_code(payload: LoginCodeSendRequest, db: Session = Depends(g
     status_code=status.HTTP_201_CREATED,
 )
 async def register(payload: RegisterRequest, db: Session = Depends(get_db)):
+    """用户注册：校验验证码（Redis 优先，失败则回退数据库）后创建用户"""
     email = _normalize_email(payload.email)
     existing_user = db.query(AppUser).filter(AppUser.email == email).first()
     if existing_user:
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Email is already registered")
 
-    now = utcnow()
-    code_record = db.query(EmailVerificationCode).filter(
-        EmailVerificationCode.email == email,
-        EmailVerificationCode.purpose == VerificationPurpose.REGISTER,
-        EmailVerificationCode.is_used.is_(False),
-        EmailVerificationCode.expires_at >= now,
-    ).order_by(EmailVerificationCode.created_at.desc()).first()
+    redis_result = _verify_code_with_redis(
+        email,
+        VerificationPurpose.REGISTER,
+        payload.verification_code,
+        strict=False, # Never be strict so we can fallback to DB if redis is down
+    )
+    code_record = None
 
-    if not code_record:
-        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Verification code does not exist or expired")
-
-    if not verify_verification_code(payload.verification_code, code_record.code_hash):
+    if redis_result is False:
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid verification code")
 
+    if redis_result is None:
+        # 即使在 _is_redis_only() 模式下，也去数据库兜底查找
+        # 这样如果Redis挂了时代码回退到了DB，验证时也能从DB拿出来。
+        code_record = _get_latest_valid_code_record(
+            db,
+            email=email,
+            purpose=VerificationPurpose.REGISTER,
+        )
+        if not code_record:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Verification code does not exist or expired")
+
+        if not verify_verification_code(payload.verification_code, code_record.code_hash):
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid verification code")
+    else:
+        # Redis 成功时尽量同步消费 DB 里的最新验证码，保持一致性。
+        # 即使在 _is_redis_only()，如果先前发生了降级，这里也顺手清理掉。
+        code_record = _get_latest_valid_code_record(
+            db,
+            email=email,
+            purpose=VerificationPurpose.REGISTER,
+        )
+
+    now = utcnow()
     nickname = payload.nickname or email.split("@")[0]
     user = AppUser(
         email=email,
@@ -234,9 +568,11 @@ async def register(payload: RegisterRequest, db: Session = Depends(get_db)):
         metadata_={"email_verified_at": now.isoformat()},
     )
 
-    code_record.is_used = True
     db.add(user)
-    db.add(code_record)
+    if code_record is not None:
+        code_record.is_used = True
+        db.add(code_record)
+
     db.commit()
     db.refresh(user)
 
@@ -245,6 +581,7 @@ async def register(payload: RegisterRequest, db: Session = Depends(get_db)):
 
 @router.post("/login", response_model=AuthTokenResponse)
 async def login(payload: LoginRequest, db: Session = Depends(get_db)):
+    """密码登录"""
     email = _normalize_email(payload.email)
     user = db.query(AppUser).filter(AppUser.email == email).first()
 
@@ -259,28 +596,48 @@ async def login(payload: LoginRequest, db: Session = Depends(get_db)):
 
 @router.post("/login/code", response_model=AuthTokenResponse)
 async def login_with_code(payload: LoginWithCodeRequest, db: Session = Depends(get_db)):
+    """验证码登录：Redis 校验优先，失败则从数据库兜底"""
     email = _normalize_email(payload.email)
     user = db.query(AppUser).filter(AppUser.email == email).first()
 
     if not user:
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid email or verification code")
 
-    now = utcnow()
-    code_record = db.query(EmailVerificationCode).filter(
-        EmailVerificationCode.email == email,
-        EmailVerificationCode.purpose == VerificationPurpose.LOGIN,
-        EmailVerificationCode.is_used.is_(False),
-        EmailVerificationCode.expires_at >= now,
-    ).order_by(EmailVerificationCode.created_at.desc()).first()
+    redis_result = _verify_code_with_redis(
+        email,
+        VerificationPurpose.LOGIN,
+        payload.verification_code,
+        strict=False, # Never be strict so we can fallback to DB if redis is down
+    )
+    code_record = None
 
-    if not code_record:
-        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Verification code does not exist or expired")
-
-    if not verify_verification_code(payload.verification_code, code_record.code_hash):
+    if redis_result is False:
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid email or verification code")
 
-    code_record.is_used = True
-    db.add(code_record)
+    if redis_result is None:
+        code_record = _get_latest_valid_code_record(
+            db,
+            email=email,
+            purpose=VerificationPurpose.LOGIN,
+        )
+        if not code_record:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Verification code does not exist or expired")
+
+        if not verify_verification_code(payload.verification_code, code_record.code_hash):
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid email or verification code")
+    else:
+        # Redis 成功时尽量同步消费 DB 里的最新验证码，保持一致性。
+        # 即使在 _is_redis_only()，如果先前发生了降级，这里也顺手清理掉。
+        code_record = _get_latest_valid_code_record(
+            db,
+            email=email,
+            purpose=VerificationPurpose.LOGIN,
+        )
+
+    if code_record is not None:
+        code_record.is_used = True
+        db.add(code_record)
+
     db.commit()
 
     return _build_auth_response(user)