rate_limit.py

"""跨 worker API 限流：已验证用户按账号计数，匿名请求按客户端 IP 计数。"""

from datetime import datetime
from functools import wraps
from inspect import iscoroutinefunction
import hashlib
import hmac
import logging
import secrets
import threading
import time

from fastapi import Request
from fastapi import HTTPException
from auth_io import run_identity_io
from owned_thread import OwnedThreadGroup
from sqlalchemy import case, delete
from sqlalchemy.dialects.mysql import insert as mysql_insert
from sqlalchemy.dialects.postgresql import insert as postgresql_insert
from sqlalchemy.dialects.sqlite import insert as sqlite_insert

from 安全认证 import get_account_from_token
from database_sql import SessionLocal, engine
from models_sql import ApiRateLimitBucket, VerificationCode, VerificationContactCooldown
from 安全认证 import JWT_SECRET

logger = logging.getLogger("ComfyUI-Ranking.RateLimit")
_cleanup_lock = threading.Lock()
_last_cleanup_at = 0.0

MAX_SHARED_WORK = 40
_shared_work = OwnedThreadGroup(MAX_SHARED_WORK,
    lambda: HTTPException(status_code=503, detail='限流服务暂不可用'))
run_shared_work = _shared_work.run
start_shared_work = _shared_work.start
stop_shared_work = _shared_work.stop
close_shared_work = _shared_work.close


def get_remote_address(request: Request) -> str:
    """仅信任 ASGI transport 提供的客户端地址，不直接信任可伪造的代理头。"""
    client = getattr(request, "client", None)
    return getattr(client, "host", None) or "unknown"


def get_rate_limit_key(request: Request) -> str:
    """返回稳定的限流身份键；只读取并验证 Bearer Token，不记录或返回 Token。"""
    authorization = request.headers.get("Authorization", "")
    scheme, _, token = authorization.partition(" ")
    if scheme.lower() == "bearer" and token.strip():
        account = get_account_from_token(token.strip())
        if account:
            return f"account:{account.casefold()}"
    return f"ip:{get_remote_address(request)}"


def _consume_shared_limit(identity: str, scope: str, limit: int) -> None:
    """使用数据库原子 UPSERT 计数；数据库不可用时 fail closed。"""
    now = int(time.time())
    window_start = now - (now % 60)
    bucket_key = hashlib.sha256(f"{identity}\0{scope}".encode("utf-8")).hexdigest()
    session = None
    try:
        session = SessionLocal()
        dialect = engine.dialect.name
        values = {
            "bucket_key": bucket_key,
            "window_start": window_start,
            "request_count": 1,
            "updated_at": datetime.utcnow(),
        }
        if dialect == "sqlite":
            statement = sqlite_insert(ApiRateLimitBucket).values(**values)
            statement = statement.on_conflict_do_update(
                index_elements=[ApiRateLimitBucket.bucket_key],
                set_={
                    "window_start": window_start,
                    "request_count": case(
                        (ApiRateLimitBucket.window_start == window_start,
                         case((ApiRateLimitBucket.request_count < limit + 1,
                               ApiRateLimitBucket.request_count + 1), else_=limit + 1)),
                        else_=1,
                    ),
                    "updated_at": datetime.utcnow(),
                },
            )
        elif dialect == "postgresql":
            statement = postgresql_insert(ApiRateLimitBucket).values(**values)
            statement = statement.on_conflict_do_update(
                index_elements=[ApiRateLimitBucket.bucket_key],
                set_={
                    "window_start": window_start,
                    "request_count": case(
                        (ApiRateLimitBucket.window_start == window_start,
                         case((ApiRateLimitBucket.request_count < limit + 1,
                               ApiRateLimitBucket.request_count + 1), else_=limit + 1)),
                        else_=1,
                    ),
                    "updated_at": datetime.utcnow(),
                },
            )
        elif dialect == "mysql":
            statement = mysql_insert(ApiRateLimitBucket).values(**values)
            statement = statement.on_duplicate_key_update(
                [
                    # MySQL evaluates ON DUPLICATE KEY UPDATE assignments left-to-right.
                    # Read the old window before replacing it, or a new minute would
                    # incorrectly continue the prior minute's count.
                    ("request_count", case(
                        (ApiRateLimitBucket.window_start == window_start,
                         case((ApiRateLimitBucket.request_count < limit + 1,
                               ApiRateLimitBucket.request_count + 1), else_=limit + 1)),
                        else_=1,
                    )),
                    ("window_start", window_start),
                    ("updated_at", datetime.utcnow()),
                ]
            )
        else:
            raise RuntimeError(f"Unsupported SQL dialect for shared rate limiting: {dialect}")

        session.execute(statement)
        bucket = session.get(ApiRateLimitBucket, bucket_key)
        if not bucket:
            raise RuntimeError("Rate limit bucket upsert did not produce a row")
        request_count = bucket.request_count
        session.commit()
        if request_count > limit:
            raise HTTPException(status_code=429, detail=f"请求过于频繁，限流 {limit}/分钟，请稍后再试")
        _cleanup_old_buckets_if_due()
    except HTTPException:
        if session is not None:
            try: session.rollback()
            except Exception: logger.error("SHARED_RATE_LIMIT_ROLLBACK_FAIL")
        raise
    except Exception as exc:
        if session is not None:
            try: session.rollback()
            except Exception: logger.error("SHARED_RATE_LIMIT_ROLLBACK_FAIL")
        logger.error("SHARED_RATE_LIMIT_DB_FAIL | scope=%s | error_type=%s", scope, type(exc).__name__)
        raise HTTPException(status_code=503, detail="限流服务暂不可用") from exc
    finally:
        if session is not None:
            try: session.close()
            except Exception as error:
                logger.error("SHARED_RATE_LIMIT_CLOSE_FAIL | error_type=%s", type(error).__name__)
                raise HTTPException(status_code=503, detail="限流服务暂不可用") from None


def _cleanup_old_buckets_if_due() -> None:
    global _last_cleanup_at
    now = time.time()
    if now - _last_cleanup_at < 900:
        return
    if not _cleanup_lock.acquire(blocking=False):
        return
    try:
        if now - _last_cleanup_at < 900:
            return
        session = SessionLocal()
        try:
            session.execute(delete(ApiRateLimitBucket).where(ApiRateLimitBucket.window_start < int(now) - 86_400))
            session.execute(delete(VerificationContactCooldown).where(
                VerificationContactCooldown.last_sent_at < int(now) - 86_400
            ))
            session.execute(delete(VerificationCode).where(VerificationCode.expires_at < int(now)))
            session.commit()
            _last_cleanup_at = now
        except Exception:
            session.rollback()
            logger.warning("SHARED_RATE_LIMIT_CLEANUP_FAIL")
        finally:
            session.close()
    finally:
        _cleanup_lock.release()


def shared_limit(limit: int):
    """FastAPI endpoint decorator backed by the shared SQL minute-window table."""
    if not isinstance(limit, int) or limit < 1:
        raise ValueError("limit must be a positive integer")

    def decorate(endpoint):
        scope = f"{endpoint.__module__}.{endpoint.__name__}"
        is_async = iscoroutinefunction(endpoint)

        @wraps(endpoint)
        async def wrapped(*args, **kwargs):
            request = kwargs.get("request")
            if request is None:
                request = next((arg for arg in args if isinstance(arg, Request)), None)
            if request is None:
                raise RuntimeError(f"{scope} requires a FastAPI Request parameter")
            identity = await run_identity_io(get_rate_limit_key, request)
            await run_shared_work(_consume_shared_limit, identity, scope, limit)
            if is_async:
                return await endpoint(*args, **kwargs)
            return await run_shared_work(endpoint, *args, **kwargs)

        return wrapped

    return decorate


def reserve_contact_cooldown(contact: str, cooldown_seconds: int) -> int:
    """原子预留联系方式冷却；返回 0 表示成功，否则返回剩余秒数。"""
    normalized = str(contact or "").strip().casefold()
    contact_key = hashlib.sha256(normalized.encode("utf-8")).hexdigest()
    now = int(time.time())
    cutoff = now - cooldown_seconds
    claim_id = secrets.token_hex(16)
    session = SessionLocal()
    try:
        values = {"contact_key": contact_key, "last_sent_at": now, "claim_id": claim_id}
        dialect = engine.dialect.name
        if dialect == "sqlite":
            statement = sqlite_insert(VerificationContactCooldown).values(**values)
            statement = statement.on_conflict_do_update(
                index_elements=[VerificationContactCooldown.contact_key],
                set_={
                    "last_sent_at": case((VerificationContactCooldown.last_sent_at <= cutoff, now),
                                         else_=VerificationContactCooldown.last_sent_at),
                    "claim_id": case((VerificationContactCooldown.last_sent_at <= cutoff, claim_id),
                                     else_=VerificationContactCooldown.claim_id),
                },
            )
        elif dialect == "postgresql":
            statement = postgresql_insert(VerificationContactCooldown).values(**values)
            statement = statement.on_conflict_do_update(
                index_elements=[VerificationContactCooldown.contact_key],
                set_={
                    "last_sent_at": case((VerificationContactCooldown.last_sent_at <= cutoff, now),
                                         else_=VerificationContactCooldown.last_sent_at),
                    "claim_id": case((VerificationContactCooldown.last_sent_at <= cutoff, claim_id),
                                     else_=VerificationContactCooldown.claim_id),
                },
            )
        elif dialect == "mysql":
            statement = mysql_insert(VerificationContactCooldown).values(**values)
            statement = statement.on_duplicate_key_update(
                [
                    # Evaluate eligibility against the old timestamp before advancing it.
                    ("claim_id", case((VerificationContactCooldown.last_sent_at <= cutoff, claim_id),
                                      else_=VerificationContactCooldown.claim_id)),
                    ("last_sent_at", case((VerificationContactCooldown.last_sent_at <= cutoff, now),
                                          else_=VerificationContactCooldown.last_sent_at)),
                ]
            )
        else:
            raise RuntimeError(f"Unsupported SQL dialect for contact cooldown: {dialect}")

        session.execute(statement)
        record = session.get(VerificationContactCooldown, contact_key)
        if not record:
            raise RuntimeError("Contact cooldown upsert did not produce a row")
        allowed = record.claim_id == claim_id
        last_sent_at = record.last_sent_at
        session.commit()
        return 0 if allowed else max(1, cooldown_seconds - (now - last_sent_at))
    except Exception as exc:
        session.rollback()
        logger.error("CONTACT_COOLDOWN_DB_FAIL | error_type=%s", type(exc).__name__)
        raise HTTPException(status_code=503, detail="验证码发送频控暂不可用") from exc
    finally:
        session.close()


def _verification_key(key: str) -> str:
    return hashlib.sha256(str(key).strip().casefold().encode("utf-8")).hexdigest()


def _verification_digest(key: str, code: str) -> str:
    return hmac.new(JWT_SECRET.encode("utf-8"), f"{_verification_key(key)}\0{code}".encode("utf-8"), hashlib.sha256).hexdigest()


def store_verification_code(key: str, code: str, expires_at: int) -> None:
    """存储跨 worker 可用的一次性验证码，不保存验证码或联系方式明文。"""
    session = SessionLocal()
    try:
        session.merge(VerificationCode(contact_key=_verification_key(key),
                                        code_digest=_verification_digest(key, code),
                                        expires_at=int(expires_at)))
        session.commit()
    except Exception as exc:
        session.rollback()
        logger.error("VERIFICATION_CODE_STORE_FAIL | error_type=%s", type(exc).__name__)
        raise HTTPException(status_code=503, detail="验证码服务暂不可用") from exc
    finally:
        session.close()


def consume_verification_code(key: str, code: str, session=None) -> bool:
    """正确且未过期时原子删除验证码；可传入业务事务 session 一并提交。"""
    owns_session = session is None
    session = session or SessionLocal()
    try:
        result = session.execute(delete(VerificationCode).where(
            VerificationCode.contact_key == _verification_key(key),
            VerificationCode.code_digest == _verification_digest(key, code),
            VerificationCode.expires_at >= int(time.time()),
        ))
        consumed = result.rowcount == 1
        if owns_session:
            session.commit()
        return consumed
    except Exception as exc:
        if owns_session:
            session.rollback()
        logger.error("VERIFICATION_CODE_CONSUME_FAIL | error_type=%s", type(exc).__name__)
        raise HTTPException(status_code=503, detail="验证码服务暂不可用") from exc
    finally:
        if owns_session:
            session.close()