Spaces:
Running
Running
File size: 6,164 Bytes
64f7965 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | import time
from collections import defaultdict
from typing import Dict, List, Tuple
from threading import Timer, Lock
from fastapi import HTTPException
# In-Memory Email Rate Limits per action (3 emails per 1 hour window, 60s per-send cooldown)
EMAIL_ACTION_LIMITS = {
"email_verification": {"limit": 3, "window_seconds": 3600, "cooldown_seconds": 60},
"forgot_password": {"limit": 3, "window_seconds": 3600, "cooldown_seconds": 60},
"change_password_confirmation": {"limit": 3, "window_seconds": 3600, "cooldown_seconds": 60},
}
# In-memory stores & thread lock
_store: Dict[str, Dict[str, List[float]]] = defaultdict(lambda: defaultdict(list))
_cooldowns: Dict[str, Dict[str, float]] = defaultdict(dict)
_lock = Lock()
def _cleanup_store():
"""Periodically purges expired entries from _store and _cooldowns to prevent memory leaks."""
now = time.time()
with _lock:
for action, users in list(_store.items()):
config = EMAIL_ACTION_LIMITS.get(action, {})
window = config.get("window_seconds", 3600)
cutoff = now - window
for identifier, timestamps in list(users.items()):
filtered = [ts for ts in timestamps if ts > cutoff]
if filtered:
_store[action][identifier] = filtered
else:
del _store[action][identifier]
for action, users in list(_cooldowns.items()):
for identifier, until in list(users.items()):
if now >= until:
del _cooldowns[action][identifier]
# Schedule next cleanup run in 30 minutes
t = Timer(1800, _cleanup_store)
t.daemon = True
t.start()
# Start background cleanup timer (daemon thread so it doesn't block process exit)
_cleanup_timer = Timer(1800, _cleanup_store)
_cleanup_timer.daemon = True
_cleanup_timer.start()
def check_and_record_email_rate_limit(
action: str, identifier: str, cooldown_seconds: int = 60
) -> Tuple[bool, int, int]:
"""
Check and record an in-memory rate limit and per-send cooldown for a specific email action.
:param action: 'email_verification', 'forgot_password', or 'change_password_confirmation'
:param identifier: Email address or user ID
:param cooldown_seconds: Minimum seconds between individual sends (default 60s)
:return: (is_allowed, remaining_attempts, retry_after_seconds)
"""
if action not in EMAIL_ACTION_LIMITS:
raise ValueError(f"Unknown action '{action}'. Allowed actions: {list(EMAIL_ACTION_LIMITS.keys())}")
config = EMAIL_ACTION_LIMITS[action]
limit = config["limit"]
window = config["window_seconds"]
cooldown = config.get("cooldown_seconds", cooldown_seconds)
now = time.time()
clean_id = identifier.lower().strip()
with _lock:
# 1. Check per-send cooldown first
cooldown_until = _cooldowns[action].get(clean_id, 0)
if now < cooldown_until:
wait = int(cooldown_until - now) + 1
cutoff = now - window
timestamps = [ts for ts in _store[action][clean_id] if ts > cutoff]
remaining = max(0, limit - len(timestamps))
return False, remaining, max(1, wait)
# 2. Check 1-hour window limit
cutoff = now - window
timestamps = [ts for ts in _store[action][clean_id] if ts > cutoff]
if len(timestamps) >= limit:
oldest = timestamps[0]
retry_after = int(oldest + window - now) + 1
_store[action][clean_id] = timestamps
return False, 0, max(1, retry_after)
# Allowed: record timestamp and set next cooldown
timestamps.append(now)
_store[action][clean_id] = timestamps
_cooldowns[action][clean_id] = now + cooldown
remaining = limit - len(timestamps)
return True, remaining, 0
def enforce_email_rate_limit(action: str, identifier: str) -> int:
"""
Enforces in-memory rate limit and per-send cooldown. Raises HTTP 429 if violated.
:return: Number of remaining attempts in the current window.
"""
allowed, remaining, retry_after = check_and_record_email_rate_limit(action, identifier)
if not allowed:
action_name = action.replace("_", " ").title()
if retry_after <= 60:
msg = f"Please wait {retry_after} seconds before requesting another {action_name} email."
else:
minutes = (retry_after + 59) // 60
msg = f"Rate limit exceeded for {action_name}. Maximum 3 emails per hour allowed. Please try again in {minutes} minute(s)."
raise HTTPException(
status_code=429,
detail=msg,
headers={"Retry-After": str(retry_after)}
)
return remaining
def get_email_rate_limit_status(action: str, identifier: str) -> Dict[str, int]:
"""
Get current usage and cooldown status without recording a new attempt.
"""
if action not in EMAIL_ACTION_LIMITS:
raise ValueError(f"Unknown action '{action}'. Allowed actions: {list(EMAIL_ACTION_LIMITS.keys())}")
config = EMAIL_ACTION_LIMITS[action]
limit = config["limit"]
window = config["window_seconds"]
cooldown = config.get("cooldown_seconds", 60)
now = time.time()
cutoff = now - window
clean_id = identifier.lower().strip()
with _lock:
timestamps = [ts for ts in _store[action][clean_id] if ts > cutoff]
used = len(timestamps)
remaining = max(0, limit - used)
cooldown_until = _cooldowns[action].get(clean_id, 0)
cooldown_remaining = max(0, int(cooldown_until - now) + 1) if now < cooldown_until else 0
window_retry_after = 0
if used >= limit and timestamps:
window_retry_after = max(1, int(timestamps[0] + window - now) + 1)
retry_after = max(cooldown_remaining, window_retry_after)
return {
"limit": limit,
"used": used,
"remaining": remaining,
"window_seconds": window,
"cooldown_seconds": cooldown,
"cooldown_remaining_seconds": cooldown_remaining,
"retry_after_seconds": retry_after
}
|