42 lines
1.8 KiB
Python
42 lines
1.8 KiB
Python
# core/rate_limit.py
|
|
|
|
from fastapi import Request, WebSocket, HTTPException, status
|
|
from starlette.requests import HTTPConnection
|
|
from db.redis import redis_client
|
|
|
|
class RateLimiter:
|
|
def __init__(self, requests: int, window_seconds: int, scope: str = "global"):
|
|
"""
|
|
:param requests: number of requests allowed
|
|
:param window_seconds: window time in seconds
|
|
:param scope: type of limit ("global" for all site, "endpoint" for a specific path)
|
|
"""
|
|
self.requests = requests
|
|
self.window_seconds = window_seconds
|
|
self.scope = scope
|
|
|
|
async def __call__(self, connection: HTTPConnection):
|
|
client_ip = connection.client.host if connection.client else "127.0.0.1"
|
|
|
|
real_ip = connection.headers.get("x-real-ip", connection.headers.get("x-forwarded-for", client_ip))
|
|
real_ip = real_ip.split(",")[0].strip()
|
|
if self.scope == "global":
|
|
key = f"rate_limit:global:{real_ip}"
|
|
else:
|
|
path = connection.scope["path"]
|
|
key = f"rate_limit:endpoint:{real_ip}:{path}"
|
|
|
|
current_count = await redis_client.incr(key)
|
|
|
|
if current_count == 1:
|
|
await redis_client.expire(key, self.window_seconds)
|
|
|
|
if current_count > self.requests:
|
|
await redis_client.expire(key, self.window_seconds)
|
|
|
|
# If it's a WebSocket connection, we might want to raise WebSocketException
|
|
# But Starlette's HTTPException is also handled by FastAPI for WebSockets by closing the connection.
|
|
raise HTTPException(
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
detail="تعداد درخواستهای شما از حد مجاز گذشته است. لطفاً بعداً دوباره تلاش کنید."
|
|
) |