wikm/Neda
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: fix rate_limit.py

Browse Source
This commit is contained in:
roai_linux 2026-03-28 16:44:44 +03:30
parent 21e19ed6a9
commit 230c2460c9
1 changed files with 8 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
Back/core/rate_limit.py
View File

@ -1,6 +1,7 @@
# core/rate_limit.py # core/rate_limit.py
from fastapi import Request, HTTPException, status from fastapi import Request, WebSocket, HTTPException, status
from starlette.requests import HTTPConnection
from db.redis import redis_client from db.redis import redis_client
class RateLimiter: class RateLimiter:
@ -14,15 +15,15 @@ class RateLimiter:
self.window_seconds = window_seconds self.window_seconds = window_seconds
self.scope = scope self.scope = scope
async def __call__(self, request: Request): async def __call__(self, connection: HTTPConnection):
client_ip = request.client.host if request.client else "127.0.0.1" client_ip = connection.client.host if connection.client else "127.0.0.1"
real_ip = request.headers.get("x-real-ip", request.headers.get("x-forwarded-for", client_ip)) real_ip = connection.headers.get("x-real-ip", connection.headers.get("x-forwarded-for", client_ip))
real_ip = real_ip.split(",")[0].strip() real_ip = real_ip.split(",")[0].strip()
if self.scope == "global": if self.scope == "global":
key = f"rate_limit:global:{real_ip}" key = f"rate_limit:global:{real_ip}"
else: else:
path = request.scope["path"] path = connection.scope["path"]
key = f"rate_limit:endpoint:{real_ip}:{path}" key = f"rate_limit:endpoint:{real_ip}:{path}"
current_count = await redis_client.incr(key) current_count = await redis_client.incr(key)
@ -33,6 +34,8 @@ class RateLimiter:
if current_count > self.requests: if current_count > self.requests:
await redis_client.expire(key, self.window_seconds) 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( raise HTTPException(
status_code=status.HTTP_429_TOO_MANY_REQUESTS, status_code=status.HTTP_429_TOO_MANY_REQUESTS,
detail="Too many requests. Please try again later." detail="Too many requests. Please try again later."