Fix rate limiting

This commit is contained in:
2025-11-12 16:31:55 +03:00
Unverified
parent 22d667093d
commit 019414608f
4 changed files with 37 additions and 49 deletions
+8 -5
View File
@@ -18,7 +18,7 @@ import os
from security.audit import log_security
from security.profanity import contains_profanity
from security.user_agent_blocklist import is_user_agent_blocked
from security.rate_limit import rate_limit_per_ip, rate_limit_per_user
from security.rate_limit import rate_limit_per_ip
router = APIRouter()
_FAILED_ATTEMPT_WINDOW_SECONDS = 300
@@ -445,7 +445,7 @@ def logout(
@router.post("/change-password")
@rate_limit_per_user("5/hour")
@rate_limit_per_ip("5/hour")
def change_password(
request: Request,
password_request: ChangePasswordRequest,
@@ -487,7 +487,8 @@ def change_password(
@router.get("/users")
def list_users(current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
@rate_limit_per_ip("30/minute") # Per-IP limit to prevent abuse
def list_users(request: Request, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
users = db.query(User).order_by(User.username.asc()).all()
return {
"users": [
@@ -497,13 +498,15 @@ def list_users(current_user: User = Depends(get_current_user), db: Session = Dep
@router.get("/crypto/public-key/of/{user_id}")
def get_public_key_of(user_id: int, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
@rate_limit_per_ip("100/minute") # Per-IP limit to prevent abuse
def get_public_key_of(request: Request, user_id: int, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
row = db.query(CryptoPublicKey).filter(CryptoPublicKey.user_id == user_id).first()
return {"publicKey": row.public_key_b64 if row else None}
@router.get("/users/search")
def search_users(q: str, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
@rate_limit_per_ip("60/minute") # Per-IP limit to prevent abuse
def search_users(request: Request, q: str, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
if len(q.strip()) < 2:
return {"users": []}
+14 -10
View File
@@ -27,7 +27,7 @@ import json
from better_profanity import profanity as _bp
from security.audit import log_access, log_dm, log_public_chat, log_security
from security.profanity import censor_text
from security.rate_limit import rate_limit_per_user
from security.rate_limit import rate_limit_per_ip
router = APIRouter()
logger = logging.getLogger("uvicorn.error")
@@ -386,7 +386,7 @@ async def _send_message_internal(
@router.post("/send_message")
@rate_limit_per_user("30/minute")
@rate_limit_per_ip("30/minute")
async def send_message(
request: Request,
message_request: SendMessageRequest | None = None,
@@ -414,7 +414,8 @@ async def send_message(
@router.get("/get_messages")
async def get_messages(db: Session = Depends(get_db)):
@rate_limit_per_ip("60/minute") # Per-IP limit to prevent abuse
async def get_messages(request: Request, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
messages = db.query(Message).order_by(Message.timestamp.asc()).all()
messages_data = []
@@ -428,7 +429,7 @@ async def get_messages(db: Session = Depends(get_db)):
@router.post("/dm/send")
@rate_limit_per_user("20/minute")
@rate_limit_per_ip("20/minute")
async def dm_send(
request: Request,
payload: dict | None = None,
@@ -578,7 +579,8 @@ def convert_envelopes(envs: list[DMEnvelope]):
}
@router.get("/dm/fetch")
async def dm_fetch(since: int | None = None, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
@rate_limit_per_ip("60/minute") # Per-IP limit to prevent abuse
async def dm_fetch(request: Request, since: int | None = None, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
q = db.query(DMEnvelope).filter(DMEnvelope.recipient_id == current_user.id)
if since:
q = q.filter(DMEnvelope.id > since)
@@ -586,7 +588,8 @@ async def dm_fetch(since: int | None = None, current_user: User = Depends(get_cu
@router.get("/dm/history/{other_user_id}")
async def dm_history(other_user_id: int, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
@rate_limit_per_ip("60/minute") # Per-IP limit to prevent abuse
async def dm_history(request: Request, other_user_id: int, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
return convert_envelopes(
db.query(DMEnvelope)
.filter(
@@ -599,7 +602,8 @@ async def dm_history(other_user_id: int, current_user: User = Depends(get_curren
@router.get("/dm/conversations")
async def get_dm_conversations(current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
@rate_limit_per_ip("60/minute") # Per-IP limit to prevent abuse
async def get_dm_conversations(request: Request, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
# Get all DM conversations where current user is involved
conversations_query = db.query(DMEnvelope).filter(
(DMEnvelope.sender_id == current_user.id) | (DMEnvelope.recipient_id == current_user.id)
@@ -641,7 +645,7 @@ async def get_dm_conversations(current_user: User = Depends(get_current_user), d
@router.put("/edit_message/{message_id}")
@rate_limit_per_user("20/minute")
@rate_limit_per_ip("20/minute")
async def edit_message(
request: Request,
message_id: int,
@@ -718,7 +722,7 @@ async def delete_message(
@router.post("/add_reaction")
@rate_limit_per_user("50/minute")
@rate_limit_per_ip("50/minute")
async def add_reaction(
request: Request,
reaction_request: ReactionRequest,
@@ -787,7 +791,7 @@ async def add_reaction(
@router.post("/dm/add_reaction")
@rate_limit_per_user("50/minute")
@rate_limit_per_ip("50/minute")
async def add_dm_reaction(
request: Request,
reaction_request: DMReactionRequest,
+4 -4
View File
@@ -17,7 +17,7 @@ from similarity import is_user_similar_to_verified
from .messaging import messagingManager
from security.audit import log_security
from security.profanity import contains_profanity
from security.rate_limit import rate_limit_per_user
from security.rate_limit import rate_limit_per_ip
router = APIRouter()
@@ -41,7 +41,7 @@ PROFILE_PICTURES_DIR = Path("data/uploads/pfp")
os.makedirs(PROFILE_PICTURES_DIR, exist_ok=True)
@router.post("/upload-profile-picture")
@rate_limit_per_user("10/minute")
@rate_limit_per_ip("10/minute")
async def upload_profile_picture(
request: Request,
profile_picture: UploadFile = File(...),
@@ -167,7 +167,7 @@ async def list_users(
}
@router.put("/user/profile")
@rate_limit_per_user("10/minute")
@rate_limit_per_ip("10/minute")
async def update_user_profile(
request: Request,
update_request: UpdateProfileRequest,
@@ -245,7 +245,7 @@ async def update_user_profile(
@router.put("/user/bio")
@rate_limit_per_user("10/minute")
@rate_limit_per_ip("10/minute")
async def update_user_bio(
request: Request,
bio_request: UpdateBioRequest,
+10 -29
View File
@@ -4,43 +4,24 @@ from typing import Callable
from fastapi import Request
from slowapi import Limiter
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
from utils import get_client_ip
# Initialize limiter with IP-based key function
limiter = Limiter(
key_func=lambda request: get_client_ip(request) or get_remote_address(request),
default_limits=["1000/hour"], # Global default limit
storage_uri="memory://", # In-memory storage (can be changed to Redis later)
)
def get_user_id_key(request: Request) -> str:
"""Get rate limit key based on authenticated user ID."""
user = getattr(getattr(request, "state", None), "current_user", None)
if user and hasattr(user, "id"):
return f"user:{user.id}"
# Fallback to IP if not authenticated
return get_client_ip(request) or get_remote_address(request)
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.
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 decorators for different endpoint types
# 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 rate_limit_per_user(limit: str) -> Callable:
"""Rate limit based on authenticated user ID, fallback to IP.
Note: The user must be authenticated (get_current_user dependency must run first).
The user will be available in request.state.current_user after authentication.
"""
return limiter.limit(limit, key_func=get_user_id_key)