"""跨 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()