mirror of
https://github.com/fromchat-messenger/web.git
synced 2026-09-22 19:15:08 +03:00
232 lines
8.3 KiB
Python
232 lines
8.3 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from typing import Callable
|
|
from fastapi import Request
|
|
from slowapi import Limiter
|
|
from slowapi.util import get_remote_address
|
|
|
|
from backend.shared.utils import get_client_ip
|
|
|
|
logger = logging.getLogger("uvicorn.error")
|
|
|
|
def get_ip_key(request: Request) -> str:
|
|
"""Get rate limit key based on IP address."""
|
|
return get_client_ip(request) or get_remote_address(request)
|
|
|
|
# Initialize limiter with IP-based key function
|
|
# Note: We don't set default_limits to avoid affecting all users if one IP is attacked.
|
|
# Each endpoint should have an explicit rate limit based on its sensitivity.
|
|
# Rate limits automatically expire after the time window - IPs are not permanently blocked.
|
|
limiter = Limiter(
|
|
key_func=get_ip_key,
|
|
default_limits=[], # No global default - each endpoint must have explicit limits
|
|
storage_uri="memory://", # In-memory storage (can be changed to Redis later)
|
|
)
|
|
|
|
|
|
# Rate limit decorator for IP-based limiting
|
|
def rate_limit_per_ip(limit: str) -> Callable:
|
|
"""Rate limit based on IP address."""
|
|
return limiter.limit(limit, key_func=get_ip_key)
|
|
|
|
|
|
def _get_storage_dict(storage) -> dict | None:
|
|
"""Get the internal storage dictionary from slowapi's memory storage."""
|
|
if hasattr(storage, "_storage") and isinstance(storage._storage, dict):
|
|
return storage._storage
|
|
elif hasattr(storage, "storage") and isinstance(storage.storage, dict):
|
|
return storage.storage
|
|
return None
|
|
|
|
|
|
def reset_all_rate_limits() -> int:
|
|
"""
|
|
Reset all rate limits by clearing the storage.
|
|
This should be called on startup to ensure a clean state.
|
|
Returns the number of entries cleared.
|
|
"""
|
|
try:
|
|
# Access the private _storage attribute
|
|
storage = limiter._storage
|
|
storage_dict = _get_storage_dict(storage)
|
|
|
|
if storage_dict is None:
|
|
# Try using the storage's reset method if available
|
|
if hasattr(storage, "reset"):
|
|
try:
|
|
# Try reset() with no args first (clears all)
|
|
storage.reset()
|
|
logger.info("Reset all rate limits on startup using storage.reset()")
|
|
return 1 # Assume it worked
|
|
except TypeError:
|
|
# reset() might require arguments, try clearing differently
|
|
try:
|
|
# Some storage backends need explicit clearing
|
|
if hasattr(storage, "clear"):
|
|
storage.clear()
|
|
logger.info("Reset all rate limits on startup using storage.clear()")
|
|
return 1
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
logger.warning("Could not reset rate limits: storage dict not accessible and no reset method")
|
|
return 0
|
|
|
|
count = len(storage_dict)
|
|
if count > 0:
|
|
storage_dict.clear()
|
|
logger.info(f"Reset all rate limits on startup: cleared {count} entries")
|
|
return count
|
|
except Exception as e:
|
|
logger.warning(f"Failed to reset rate limits on startup: {e}")
|
|
return 0
|
|
|
|
|
|
def reset_rate_limit_for_ip(ip: str) -> bool:
|
|
"""
|
|
Manually reset rate limit for a specific IP address.
|
|
This clears all rate limit entries for the given IP.
|
|
Returns True if any entries were cleared, False otherwise.
|
|
"""
|
|
if not ip:
|
|
return False
|
|
|
|
try:
|
|
# Access the private _storage attribute
|
|
storage = limiter._storage
|
|
storage_dict = _get_storage_dict(storage)
|
|
|
|
if storage_dict is None:
|
|
# Try alternative methods
|
|
if hasattr(storage, "reset"):
|
|
try:
|
|
storage.reset(ip)
|
|
return True
|
|
except Exception:
|
|
pass
|
|
return False
|
|
|
|
cleared = False
|
|
# slowapi stores entries with keys like "LIMITER:{ip}:{endpoint}"
|
|
# We need to find all keys that contain this IP
|
|
# Also handle cases where IP might be in different positions
|
|
keys_to_remove = []
|
|
|
|
for key in list(storage_dict.keys()):
|
|
if isinstance(key, str):
|
|
# Check multiple patterns:
|
|
# - "LIMITER:{ip}:{endpoint}"
|
|
# - Keys containing the IP anywhere
|
|
# - Keys starting with the IP
|
|
if (key.startswith(f"LIMITER:{ip}:") or
|
|
key.startswith(f"LIMITER:{ip}") or
|
|
f":{ip}:" in key or
|
|
key.endswith(f":{ip}") or
|
|
(ip in key and "LIMITER" in key)):
|
|
keys_to_remove.append(key)
|
|
|
|
for key in keys_to_remove:
|
|
try:
|
|
del storage_dict[key]
|
|
cleared = True
|
|
logger.info(f"Cleared rate limit key: {key}")
|
|
except KeyError:
|
|
pass
|
|
|
|
if cleared:
|
|
logger.info(f"Successfully cleared rate limits for IP: {ip}")
|
|
else:
|
|
logger.warning(f"No rate limit entries found for IP: {ip}")
|
|
|
|
return cleared
|
|
except Exception as e:
|
|
logger.warning(f"Failed to reset rate limit for IP {ip}: {e}")
|
|
return False
|
|
|
|
|
|
def clear_all_rate_limits() -> int:
|
|
"""
|
|
Clear all rate limit entries. Use with caution - this affects all IPs.
|
|
Returns the number of entries cleared.
|
|
"""
|
|
try:
|
|
# Access the private _storage attribute
|
|
storage = limiter._storage
|
|
storage_dict = _get_storage_dict(storage)
|
|
|
|
if storage_dict is None:
|
|
return 0
|
|
|
|
count = len(storage_dict)
|
|
storage_dict.clear()
|
|
logger.warning(f"Cleared all {count} rate limit entries")
|
|
return count
|
|
except Exception as e:
|
|
logger.error(f"Failed to clear all rate limits: {e}")
|
|
return 0
|
|
|
|
|
|
def cleanup_expired_rate_limits() -> int:
|
|
"""
|
|
Clean up expired rate limit entries from memory storage.
|
|
This helps prevent rate limits from being stuck indefinitely.
|
|
Returns the number of entries cleaned up.
|
|
"""
|
|
try:
|
|
# Access the private _storage attribute
|
|
storage = limiter._storage
|
|
storage_dict = _get_storage_dict(storage)
|
|
|
|
if storage_dict is None:
|
|
return 0
|
|
|
|
# slowapi's memory storage stores entries as tuples: (count, reset_time)
|
|
# Entries should expire naturally, but we'll clean up any that are clearly expired
|
|
now = time.time()
|
|
cleaned = 0
|
|
keys_to_remove = []
|
|
|
|
for key, value in storage_dict.items():
|
|
if isinstance(value, (tuple, list)) and len(value) >= 2:
|
|
# Check if reset_time has passed (with some buffer)
|
|
reset_time = value[1] if isinstance(value[1], (int, float)) else 0
|
|
# Add 60 second buffer to ensure we don't remove active entries
|
|
if reset_time > 0 and now > (reset_time + 60):
|
|
keys_to_remove.append(key)
|
|
elif isinstance(value, dict):
|
|
# Some storage formats use dicts with 'expiry' or 'reset' fields
|
|
expiry = value.get("expiry") or value.get("reset") or value.get("reset_time")
|
|
if expiry and isinstance(expiry, (int, float)) and now > (expiry + 60):
|
|
keys_to_remove.append(key)
|
|
|
|
for key in keys_to_remove:
|
|
try:
|
|
del storage_dict[key]
|
|
cleaned += 1
|
|
except KeyError:
|
|
pass
|
|
|
|
if cleaned > 0:
|
|
logger.info(f"Cleaned up {cleaned} expired rate limit entries")
|
|
|
|
return cleaned
|
|
except Exception as e:
|
|
logger.warning(f"Failed to cleanup expired rate limits: {e}")
|
|
return 0
|
|
|
|
|
|
async def start_rate_limit_cleanup_task() -> None:
|
|
"""Start a background task to periodically clean up expired rate limit entries."""
|
|
while True:
|
|
try:
|
|
await asyncio.sleep(300) # Run every 5 minutes
|
|
cleanup_expired_rate_limits()
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"Error in rate limit cleanup task: {e}")
|
|
await asyncio.sleep(60) # Wait 1 minute before retrying |