File size: 7,326 Bytes
6f9bb4a | 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 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 | #!/usr/bin/env python3
"""
Per-user rate limiting for the FhirFlame demo.
Protects shared AI resources from excessive usage.
"""
import os
import hmac
from dataclasses import dataclass
from typing import Optional, Any
from database import db_manager
@dataclass
class RateLimitResult:
allowed: bool
message: str
remaining_hour: int = 0
remaining_day: int = 0
class RateLimiter:
"""Tracks and enforces per-user request limits."""
def __init__(self):
self.enabled = os.getenv("RATE_LIMIT_ENABLED", "true").lower() == "true"
self.per_hour = int(os.getenv("RATE_LIMIT_PER_HOUR", "12"))
self.per_day = int(os.getenv("RATE_LIMIT_PER_DAY", "30"))
self.batch_per_hour = int(os.getenv("RATE_LIMIT_BATCH_PER_HOUR", "2"))
self.access_password = os.getenv("DEMO_ACCESS_PASSWORD", "")
self.password_request_url = self._safe_request_url(
os.getenv("DEMO_PASSWORD_REQUEST_URL", "https://bartoszlenart.com")
)
def check(
self,
user_id: str,
action: str = "process",
access_password: Optional[str] = None,
) -> RateLimitResult:
"""Check whether a user may perform an action without recording it."""
if not self.enabled:
return RateLimitResult(True, "", self.per_hour, self.per_day)
user_id = self._normalize_user_id(user_id)
hour_count = db_manager.count_rate_limit_events(user_id, action=None, since_seconds=3600)
day_count = db_manager.count_rate_limit_events(user_id, action=None, since_seconds=86400)
if action == "batch":
batch_hour_count = db_manager.count_rate_limit_events(
user_id, action="batch", since_seconds=3600
)
if batch_hour_count >= self.batch_per_hour:
if self._valid_access_password(access_password):
return RateLimitResult(True, "", 0, max(0, self.per_day - day_count))
return RateLimitResult(
allowed=False,
message=self._batch_limit_message(self.batch_per_hour, self.password_request_url),
remaining_hour=max(0, self.batch_per_hour - batch_hour_count),
remaining_day=max(0, self.per_day - day_count),
)
if hour_count >= self.per_hour:
if self._valid_access_password(access_password):
return RateLimitResult(True, "", 0, max(0, self.per_day - day_count))
return RateLimitResult(
allowed=False,
message=self._hour_limit_message(self.per_hour, self.password_request_url),
remaining_hour=0,
remaining_day=max(0, self.per_day - day_count),
)
if day_count >= self.per_day:
if self._valid_access_password(access_password):
return RateLimitResult(True, "", max(0, self.per_hour - hour_count), 0)
return RateLimitResult(
allowed=False,
message=self._day_limit_message(self.per_day, self.password_request_url),
remaining_hour=max(0, self.per_hour - hour_count),
remaining_day=0,
)
return RateLimitResult(
allowed=True,
message="",
remaining_hour=max(0, self.per_hour - hour_count - 1),
remaining_day=max(0, self.per_day - day_count - 1),
)
def record(self, user_id: str, action: str = "process") -> None:
"""Record a completed rate-limited action."""
if not self.enabled:
return
user_id = self._normalize_user_id(user_id)
db_manager.record_rate_limit_event(user_id, action)
db_manager.cleanup_rate_limit_events(older_than_hours=48)
def check_and_record(
self,
user_id: str,
action: str = "process",
access_password: Optional[str] = None,
) -> RateLimitResult:
"""Atomically check limits and record the request if allowed."""
result = self.check(user_id, action, access_password)
if result.allowed:
self.record(user_id, action)
return result
def _valid_access_password(self, candidate: Optional[str]) -> bool:
"""Validate a configured limit-override password in constant time."""
if not self.access_password or not candidate:
return False
return hmac.compare_digest(str(candidate), self.access_password)
@staticmethod
def _safe_request_url(url: str) -> str:
"""Allow only web links in the rate-limit message."""
normalized = str(url).strip()
if normalized.startswith(("https://", "http://")):
return normalized
return "https://bartoszlenart.com"
@staticmethod
def _normalize_user_id(user_id: Optional[str]) -> str:
if not user_id or not str(user_id).strip():
return "anonymous"
return str(user_id).strip()[:255]
@staticmethod
def _hour_limit_message(limit: int, request_url: str) -> str:
return (
f"🔒 **Rate limit reached**\n\n"
f"This demo allows **{limit} processing requests per hour** per user "
f"to protect shared AI resources.\n\n"
f"Enter the access password below to continue, or "
f"[request a password]({request_url})."
)
@staticmethod
def _day_limit_message(limit: int, request_url: str) -> str:
return (
f"🔒 **Daily limit reached**\n\n"
f"This demo allows **{limit} processing requests per day** per user.\n\n"
f"Enter the access password below to continue, or "
f"[request a password]({request_url})."
)
@staticmethod
def _batch_limit_message(limit: int, request_url: str) -> str:
return (
f"🔒 **Batch demo limit reached**\n\n"
f"Live batch processing is limited to **{limit} runs per hour** per user.\n\n"
f"Enter the access password below to continue, or "
f"[request a password]({request_url})."
)
rate_limiter = RateLimiter()
def get_user_id_from_request(request: Any) -> str:
"""Extract a stable per-user identifier from a Gradio or HTTP request."""
if request is None:
return "anonymous"
headers = getattr(request, "headers", None) or {}
if hasattr(headers, "get"):
forwarded = headers.get("x-forwarded-for") or headers.get("X-Forwarded-For")
if forwarded:
return forwarded.split(",")[0].strip()
client = getattr(request, "client", None)
if client and getattr(client, "host", None):
return client.host
return "anonymous"
def enforce_rate_limit(
user_id: Optional[str],
action: str = "process",
access_password: Optional[str] = None,
) -> Optional[str]:
"""
Check rate limits for a user action.
Returns an error message if blocked, otherwise None.
"""
result = rate_limiter.check_and_record(user_id, action, access_password)
return None if result.allowed else result.message
|