|
|
| """
|
| 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
|
|
|
