Fix rate limiting

This commit is contained in:
2025-12-01 18:29:54 +03:00
Unverified
parent ebf558720c
commit 8bda2220c6
4 changed files with 314 additions and 2 deletions
+29
View File
@@ -276,6 +276,29 @@ class AdminCLI:
table.add_row(entry) table.add_row(entry)
self.console.print(table) self.console.print(table)
def cmd_unblock_ip(self, args: List[str]) -> None:
if not args:
raise CLIError("Usage: unblock-ip <ip_address>")
self._require_auth()
ip = args[0].strip()
if not ip:
raise CLIError("IP address cannot be empty")
response = self._request("POST", "moderation/unblock-ip", json={"ip": ip})
data = response.json()
message = data.get("message", "IP unblocked")
self.console.print(f"[bold green]{message}[/]")
def cmd_clear_all_rate_limits(self) -> None:
"""Clear all rate limit entries. Use with caution."""
self._require_auth()
if not self._confirm("Clear ALL rate limit entries? This affects all IPs."):
self.console.print("[yellow]Operation cancelled.[/]")
return
response = self._request("POST", "moderation/clear-all-rate-limits")
data = response.json()
message = data.get("message", "Rate limits cleared")
self.console.print(f"[bold green]{message}[/]")
def cmd_help(self) -> None: def cmd_help(self) -> None:
cmds = { cmds = {
"login [username]": "Authenticate as owner/admin.", "login [username]": "Authenticate as owner/admin.",
@@ -287,6 +310,8 @@ class AdminCLI:
"block-word <words>": "Add words/phrases to chat filter.", "block-word <words>": "Add words/phrases to chat filter.",
"unblock-word <words>": "Remove words/phrases from filter.", "unblock-word <words>": "Remove words/phrases from filter.",
"blocklist": "Show current blocklist.", "blocklist": "Show current blocklist.",
"unblock-ip <ip>": "Unblock an IP address from rate limiting.",
"clear-all-rate-limits": "Clear all rate limit entries (use with caution).",
"list": "List all users.", "list": "List all users.",
"user <user>": "Show detailed user information.", "user <user>": "Show detailed user information.",
"whoami": "Display current session context.", "whoami": "Display current session context.",
@@ -347,6 +372,10 @@ class AdminCLI:
self.cmd_unblock_word(args) self.cmd_unblock_word(args)
elif command == "blocklist": elif command == "blocklist":
self.cmd_list_blocklist() self.cmd_list_blocklist()
elif command == "unblock-ip":
self.cmd_unblock_ip(args)
elif command == "clear-all-rate-limits":
self.cmd_clear_all_rate_limits()
elif command == "verify": elif command == "verify":
self.cmd_verify(args) self.cmd_verify(args)
elif command == "unverify": elif command == "unverify":
+27 -1
View File
@@ -1,3 +1,4 @@
import asyncio
import time import time
from fastapi import FastAPI, Request from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
@@ -68,9 +69,34 @@ async def lifespan(app: FastAPI):
except Exception as e: except Exception as e:
logger.error(f"Failed to start messaging cleanup task: {e}") logger.error(f"Failed to start messaging cleanup task: {e}")
# Reset all rate limits on startup to ensure clean state
# This prevents rate limits from persisting across restarts
try:
from security.rate_limit import reset_all_rate_limits
cleared = reset_all_rate_limits()
if cleared > 0:
logger.info(f"Cleared {cleared} rate limit entries on startup")
except Exception as e:
logger.warning(f"Failed to reset rate limits on startup: {e}")
# Start the rate limit cleanup task
try:
from security.rate_limit import start_rate_limit_cleanup_task
cleanup_task = asyncio.create_task(start_rate_limit_cleanup_task())
logger.info("Rate limit cleanup task started")
except Exception as e:
logger.error(f"Failed to start rate limit cleanup task: {e}")
cleanup_task = None
yield yield
# Shutdown (if needed in the future) # Shutdown - cancel cleanup task if it exists
if cleanup_task:
cleanup_task.cancel()
try:
await cleanup_task
except asyncio.CancelledError:
pass
# Инициализация FastAPI # Инициализация FastAPI
app = FastAPI(title="FromChat", lifespan=lifespan) app = FastAPI(title="FromChat", lifespan=lifespan)
+52
View File
@@ -7,12 +7,17 @@ from dependencies import get_current_user
from models import User from models import User
from security.audit import log_security from security.audit import log_security
from security.profanity import add_to_blocklist, get_blocklist, remove_from_blocklist from security.profanity import add_to_blocklist, get_blocklist, remove_from_blocklist
from security.rate_limit import reset_rate_limit_for_ip, clear_all_rate_limits
class BlocklistUpdateRequest(BaseModel): class BlocklistUpdateRequest(BaseModel):
words: List[str] = Field(default_factory=list, min_items=1) words: List[str] = Field(default_factory=list, min_items=1)
class UnblockIPRequest(BaseModel):
ip: str = Field(..., min_length=1)
router = APIRouter(prefix="/moderation", tags=["moderation"]) router = APIRouter(prefix="/moderation", tags=["moderation"])
@@ -59,4 +64,51 @@ def delete_from_blocklist(
return {"removed": removed, "words": updated} return {"removed": removed, "words": updated}
@router.post("/unblock-ip")
def unblock_ip(
request: UnblockIPRequest,
current_user: User = Depends(get_current_user)
):
"""Unblock an IP address from rate limiting."""
_ensure_owner(current_user)
ip = request.ip.strip()
if not ip:
raise HTTPException(status_code=400, detail="IP address is required")
cleared = reset_rate_limit_for_ip(ip)
log_security(
"rate_limit_unblock",
actor=current_user.username,
actor_id=current_user.id,
ip=ip,
success=cleared,
)
if cleared:
return {"status": "success", "message": f"Rate limit cleared for IP: {ip}"}
else:
return {"status": "success", "message": f"No rate limit entries found for IP: {ip}"}
@router.post("/clear-all-rate-limits")
def clear_all_rate_limits_endpoint(
current_user: User = Depends(get_current_user)
):
"""Clear all rate limit entries. Use with caution."""
_ensure_owner(current_user)
cleared = clear_all_rate_limits()
log_security(
"rate_limit_clear_all",
actor=current_user.username,
actor_id=current_user.id,
entries_cleared=cleared,
)
return {"status": "success", "message": f"Cleared {cleared} rate limit entries"}
+206 -1
View File
@@ -1,5 +1,8 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import logging
import time
from typing import Callable from typing import Callable
from fastapi import Request from fastapi import Request
from slowapi import Limiter from slowapi import Limiter
@@ -7,6 +10,8 @@ from slowapi.util import get_remote_address
from utils import get_client_ip from utils import get_client_ip
logger = logging.getLogger("uvicorn.error")
def get_ip_key(request: Request) -> str: def get_ip_key(request: Request) -> str:
"""Get rate limit key based on IP address.""" """Get rate limit key based on IP address."""
return get_client_ip(request) or get_remote_address(request) return get_client_ip(request) or get_remote_address(request)
@@ -14,6 +19,7 @@ def get_ip_key(request: Request) -> str:
# Initialize limiter with IP-based key function # Initialize limiter with IP-based key function
# Note: We don't set default_limits to avoid affecting all users if one IP is attacked. # 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. # 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( limiter = Limiter(
key_func=get_ip_key, key_func=get_ip_key,
default_limits=[], # No global default - each endpoint must have explicit limits default_limits=[], # No global default - each endpoint must have explicit limits
@@ -24,4 +30,203 @@ limiter = Limiter(
# Rate limit decorator for IP-based limiting # Rate limit decorator for IP-based limiting
def rate_limit_per_ip(limit: str) -> Callable: def rate_limit_per_ip(limit: str) -> Callable:
"""Rate limit based on IP address.""" """Rate limit based on IP address."""
return limiter.limit(limit, key_func=get_ip_key) 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