"""Author RAG Chatbot SaaS — Rate Limiting Middleware (Pure ASGI). Implements sliding window rate limiting using Redis. - Chat endpoint: 60 requests/minute per IP (main POST only) - Auth endpoint: 10 requests/minute per IP - Visitor chat: 60 messages/hour per visitor fingerprint (configurable) Uses pure ASGI to avoid BaseHTTPMiddleware request body consumption issues. Fixes: B10 — Visitor limit targets only POST /api/chat/{slug} (not session/init etc.). B11 — Blocked 429 responses do not consume quota (decr after reject). B12 — IP chat limit applies to main chat POST only, not init/track-click. """ import hashlib import json import re import structlog from starlette.types import ASGIApp, Receive, Scope, Send from app.config import get_settings logger = structlog.get_logger(__name__) cfg = get_settings() # Path-to-limit mapping: (path_prefix, limit_per_minute) _RATE_LIMIT_RULES: list[tuple[str, int]] = [ ("/api/auth", cfg.RATE_LIMIT_AUTH_PER_MINUTE), ("/api/admin", 120), # Admin API: generous but bounded ] # Regex for the MAIN chat endpoint only (not /session/init, /track-click etc.) # Pattern: /api/chat/{author_slug} — exactly two path segments after /api/chat _MAIN_CHAT_PATTERN = re.compile(r"^/api/chat/[^/]+$") # Visitor rate limit: messages/hour per fingerprint _VISITOR_LIMIT = cfg.RATE_LIMIT_VISITOR_PER_HOUR _VISITOR_WINDOW = 3600 # 1 hour in seconds _CHAT_IP_LIMIT = cfg.RATE_LIMIT_CHAT_PER_MINUTE class RateLimitMiddleware: """Sliding window rate limiter backed by Redis (pure ASGI). Two independent limits on the main chat POST only: 1. IP-based limit: per-minute (general abuse prevention) 2. Visitor fingerprint limit: per-hour (budget protection) """ def __init__(self, app: ASGIApp) -> None: self.app = app async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if scope["type"] != "http": await self.app(scope, receive, send) return path = scope.get("path", "") method = scope.get("method", "GET") try: from app.dependencies import get_redis redis = await get_redis() except Exception as e: logger.error("Rate limit check failed (Redis error)", error=str(e)) await self.app(scope, receive, send) return client_ip = _get_client_ip(scope) # ── Main chat POST: IP + visitor limits ─────────────────────────────── if method == "POST" and _MAIN_CHAT_PATTERN.match(path): ip_key = f"ratelimit:chat:{client_ip}" ip_count = await redis.incr(ip_key) if ip_count == 1: await redis.expire(ip_key, 60) if ip_count > _CHAT_IP_LIMIT: await redis.decr(ip_key) logger.warning("IP rate limit exceeded", ip=client_ip, path=path, count=ip_count) await _send_429(send, f"Too many requests. Limit: {_CHAT_IP_LIMIT}/minute.") return headers = dict(scope.get("headers", [])) ua = headers.get(b"user-agent", b"").decode() accept_lang = headers.get(b"accept-language", b"").decode() fp = _make_fingerprint(client_ip, ua, accept_lang) visitor_key = f"vrl:{fp}" v_count = await redis.incr(visitor_key) if v_count == 1: await redis.expire(visitor_key, _VISITOR_WINDOW) if v_count > _VISITOR_LIMIT: await redis.decr(visitor_key) logger.warning( "Visitor rate limit exceeded", fp=fp[:8], path=path, count=v_count, ) await _send_429( send, "You've sent a lot of messages! Take a short break and come back in a bit.", ) return else: # ── Other routes: prefix-based IP limits (auth, admin) ────────── ip_limit = _get_limit_for_path(path) if ip_limit is not None: key_segment = path.split("/")[2] if len(path.split("/")) > 2 else "unknown" ip_key = f"ratelimit:{key_segment}:{client_ip}" count = await redis.incr(ip_key) if count == 1: await redis.expire(ip_key, 60) if count > ip_limit: await redis.decr(ip_key) logger.warning("IP rate limit exceeded", ip=client_ip, path=path, count=count) await _send_429(send, f"Too many requests. Limit: {ip_limit}/minute.") return await self.app(scope, receive, send) def _get_limit_for_path(path: str) -> int | None: """Find the IP-based rate limit for a given path prefix.""" for prefix, limit in _RATE_LIMIT_RULES: if path.startswith(prefix): return limit return None def _get_client_ip(scope: Scope) -> str: """Extract client IP — checks X-Forwarded-For first (proxy/Docker aware).""" headers = dict(scope.get("headers", [])) forwarded = headers.get(b"x-forwarded-for", b"").decode() if forwarded: return forwarded.split(",")[0].strip() client = scope.get("client") return client[0] if client else "unknown" def _make_fingerprint(ip: str, ua: str, accept_lang: str) -> str: """Build a short anonymized fingerprint from non-PII signals.""" raw = f"{ip}|{ua}|{accept_lang}" return hashlib.sha256(raw.encode()).hexdigest()[:32] async def _send_429(send: Send, message: str) -> None: """Send a 429 Too Many Requests response with a human-readable message.""" body = json.dumps({"detail": message}).encode() await send({ "type": "http.response.start", "status": 429, "headers": [ (b"content-type", b"application/json"), (b"retry-after", b"60"), (b"content-length", str(len(body)).encode()), ], }) await send({ "type": "http.response.body", "body": body, })