34 lines
1.0 KiB
Python
34 lines
1.0 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
PerToolBox Server - 限流中间件
|
|
Copyright (C) 2024 Sea Network Technology Studio
|
|
Author: Canglan <admin@sea-studio.top>
|
|
License: AGPL v3
|
|
"""
|
|
|
|
from fastapi import Request, HTTPException
|
|
from slowapi import Limiter, _rate_limit_exceeded_handler
|
|
from slowapi.util import get_remote_address
|
|
from slowapi.errors import RateLimitExceeded
|
|
from ..config import settings
|
|
|
|
# 创建限流器
|
|
limiter = Limiter(key_func=get_remote_address)
|
|
|
|
def setup_rate_limit(app):
|
|
"""配置限流"""
|
|
if settings.RATE_LIMIT_ENABLED:
|
|
app.state.limiter = limiter
|
|
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
|
|
|
|
def rate_limit(requests: int = None, period: int = None):
|
|
"""限流装饰器工厂"""
|
|
if not settings.RATE_LIMIT_ENABLED:
|
|
return lambda func: func
|
|
|
|
req = requests or settings.RATE_LIMIT_REQUESTS
|
|
per = period or settings.RATE_LIMIT_PERIOD
|
|
|
|
# 直接返回 slowapi 的限流装饰器
|
|
return limiter.limit(f"{req}/{per} seconds") |