from __future__ import annotations import ipaddress from datetime import datetime, timedelta from typing import Set from pymongo import AsyncMongoClient, ReturnDocument from micropie import HttpMiddleware def _valid_ip(value: str | None) -> str | None: try: return str(ipaddress.ip_address(value.strip())) except Exception: return None class MongoRateLimitMiddleware(HttpMiddleware): """ Global MongoDB-based rate limiter (Heroku + Cloudflare safe). - One document per client IP - Fixed window counter - Escalating temporary blocks - Permanent block based on 24h violation history - Fully atomic (single DB op per request) - PyMongo Async API (no Motor) Requires: - MongoDB 4.2+ (aggregation pipeline updates) """ # --- rate config --- MAX_REQUESTS = 50 WINDOW_SECONDS = 60 BLOCK_AFTER_VIOLATIONS = 3 BLOCK_FOR_SECONDS = 900 PERMA_WINDOW_HOURS = 24 PERMA_BLOCK_AFTER = 10 def __init__( self, mongo_uri: str, db_name: str, collection_name: str = "rate_limits_global", *, allowed_hosts: Set[str] | None = None, trust_proxy_headers: bool = True, require_cf_ray: bool = True, ): self.client = AsyncMongoClient(mongo_uri) self.db = self.client[db_name] self.collection = self.db[collection_name] # Security / proxy config self.allowed_hosts = allowed_hosts or set() self.trust_proxy_headers = trust_proxy_headers self.require_cf_ray = require_cf_ray # --------------------------------------------------------- # Real client IP resolution (for Cloudflare + Heroku or similar setups) # --------------------------------------------------------- def _client_ip(self, request) -> str: headers = getattr(request, "headers", {}) or {} # 1) Optional host allow-list (prevents origin bypass) if self.allowed_hosts: host = (headers.get("host") or "").split(":", 1)[0].lower() if host and host not in self.allowed_hosts: return "unknown" # 2) Only trust proxy headers if allowed can_trust = self.trust_proxy_headers if can_trust and self.require_cf_ray: can_trust = bool(headers.get("cf-ray")) if can_trust: # Cloudflare headers (best) for h in ("cf-connecting-ip", "true-client-ip"): ip = _valid_ip(headers.get(h)) if ip: return ip # Standard proxy chain xff = headers.get("x-forwarded-for") if isinstance(xff, str): ip = _valid_ip(xff.split(",", 1)[0]) if ip: return ip # Fallback proxy header ip = _valid_ip(headers.get("x-real-ip")) if ip: return ip # 3) ASGI scope fallback (Heroku router) client = request.scope.get("client") or ("unknown", 0) return _valid_ip(client[0]) or "unknown" # --------------------------------------------------------- # Middleware hook # --------------------------------------------------------- async def before_request(self, request): client_ip = self._client_ip(request) now = datetime.utcnow() window_start_cutoff = now - timedelta(seconds=self.WINDOW_SECONDS) perma_window_cutoff = now - timedelta(hours=self.PERMA_WINDOW_HOURS) key = client_ip doc = await self.collection.find_one_and_update( {"_id": key}, [ # 1) Baseline fields { "$set": { "_id": key, "ip": client_ip, "count": {"$ifNull": ["$count", 0]}, "window_start": {"$ifNull": ["$window_start", now]}, "violations": {"$ifNull": ["$violations", 0]}, "blocked_until": {"$ifNull": ["$blocked_until", None]}, "permanent_blocked": {"$ifNull": ["$permanent_blocked", False]}, "permanent_blocked_at": { "$ifNull": ["$permanent_blocked_at", None] }, "violation_events": {"$ifNull": ["$violation_events", []]}, } }, # 2) Prune old violation events { "$set": { "violation_events": { "$filter": { "input": "$violation_events", "as": "t", "cond": {"$gte": ["$$t", perma_window_cutoff]}, } } } }, # 3) Are we currently blocked? { "$set": { "_blocked_now": { "$or": [ "$permanent_blocked", { "$and": [ {"$ne": ["$blocked_until", None]}, {"$gt": ["$blocked_until", now]}, ] }, ] } } }, # 4) Update window/count atomically (only if not blocked) { "$set": { "_window_expired": { "$cond": [ "$_blocked_now", False, {"$lt": ["$window_start", window_start_cutoff]}, ] } } }, { "$set": { "window_start": { "$cond": [ "$_blocked_now", "$window_start", {"$cond": ["$_window_expired", now, "$window_start"]}, ] }, "count": { "$cond": [ "$_blocked_now", "$count", { "$cond": [ "$_window_expired", 1, {"$add": ["$count", 1]}, ] }, ] }, } }, # 5) Over limit? { "$set": { "_over_limit": { "$and": [ {"$not": "$_blocked_now"}, {"$gt": ["$count", self.MAX_REQUESTS]}, ] } } }, # 6) Record violation if over limit { "$set": { "violations": { "$cond": [ "$_over_limit", {"$add": ["$violations", 1]}, "$violations", ] }, "violation_events": { "$cond": [ "$_over_limit", {"$concatArrays": ["$violation_events", [now]]}, "$violation_events", ] }, } }, # 7) Temporary block escalation { "$set": { "blocked_until": { "$cond": [ { "$and": [ "$_over_limit", { "$gte": [ "$violations", self.BLOCK_AFTER_VIOLATIONS, ] }, ] }, now + timedelta(seconds=self.BLOCK_FOR_SECONDS), "$blocked_until", ] } } }, # 8) Permanent block escalation {"$set": {"_events_24h": {"$size": "$violation_events"}}}, { "$set": { "permanent_blocked": { "$cond": [ { "$and": [ "$_over_limit", { "$gte": [ "$_events_24h", self.PERMA_BLOCK_AFTER, ] }, ] }, True, "$permanent_blocked", ] }, "permanent_blocked_at": { "$cond": [ { "$and": [ "$_over_limit", { "$gte": [ "$_events_24h", self.PERMA_BLOCK_AFTER, ] }, {"$eq": ["$permanent_blocked_at", None]}, ] }, now, "$permanent_blocked_at", ] }, } }, # 9) Cleanup temp fields { "$unset": [ "_blocked_now", "_window_expired", "_over_limit", "_events_24h", ] }, ], upsert=True, return_document=ReturnDocument.AFTER, projection={"count": 1, "blocked_until": 1, "permanent_blocked": 1}, ) doc = doc or {} # --- responses --- if doc.get("permanent_blocked"): return { "status_code": 403, "body": f"Access permanently blocked for IP {client_ip}.", "headers": [], } blocked_until = doc.get("blocked_until") if isinstance(blocked_until, datetime) and now < blocked_until: retry_after = max(0, int((blocked_until - now).total_seconds())) return { "status_code": 429, "body": f"Too many requests from {client_ip}. Temporarily blocked.", "headers": [("Retry-After", str(retry_after))], } if int(doc.get("count", 0)) > self.MAX_REQUESTS: return { "status_code": 429, "body": f"Rate limit exceeded for IP {client_ip}.", "headers": [], } return None async def after_request(self, request, status_code, response_body, extra_headers): return None