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