Coverage for src/lexigram/web/middleware/rate_limit.py: 18%
192 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 04:37 +0800
1"""Rate limiting middleware using Redis or CacheBackendProtocol.
3Implements sliding window algorithm with per-user and per-IP limits.
4Supports Redis (atomic Lua script), CacheBackendProtocol (fixed-window via get/set),
5and in-memory (process-local, dev/test only).
6"""
8from __future__ import annotations
10import asyncio
11from collections import deque
12from collections.abc import Callable
13from datetime import timedelta
14from functools import wraps
15from typing import TYPE_CHECKING, Any, Literal
17from starlette.requests import Request
18from starlette.responses import JSONResponse
19from starlette.types import ASGIApp, Receive, Scope, Send
21from lexigram.contracts.exceptions import RateLimitError
22from lexigram.contracts.web import WebRateLimiterProtocol
23from lexigram.logging import get_logger
24from lexigram.primitives import clock as ambient_clock
26if TYPE_CHECKING:
27 from lexigram.contracts.infra.cache import CacheBackendProtocol
29logger = get_logger(__name__)
32def _rate_limit_429_response(exc: Exception) -> JSONResponse:
33 """Build the standard 429 response for a rate-limit breach.
35 Args:
36 exc: The :class:`~lexigram.contracts.exceptions.RateLimitError`.
38 Returns:
39 A 429 ``JSONResponse`` with a ``Retry-After`` header when the
40 error carries a ``retry_after`` detail.
41 """
43 details = getattr(exc, "details", {}) or {}
44 retry_after = details.get("retry_after")
45 return JSONResponse(
46 {"error": "rate_limit_exceeded", "message": str(exc)},
47 status_code=429,
48 headers={"Retry-After": str(retry_after)} if retry_after else {},
49 )
52class MemoryLimiter:
53 """In-memory sliding window rate limiter for request tracking.
55 This is a local implementation to avoid cross-package imports.
56 """
58 def __init__(
59 self,
60 window_size: float,
61 max_requests: int,
62 ) -> None:
63 self.window_size = window_size
64 self.max_requests = max_requests
65 self._requests: deque[float] = deque()
66 self._lock = asyncio.Lock()
68 def _cleanup_old_requests(self) -> None:
69 """Remove requests outside the current window."""
70 now = ambient_clock.monotonic()
71 cutoff = now - self.window_size
72 while self._requests and self._requests[0] < cutoff:
73 self._requests.popleft()
75 async def try_acquire(self) -> bool:
76 """Try to acquire permission without blocking."""
77 async with self._lock:
78 self._cleanup_old_requests()
79 if len(self._requests) < self.max_requests:
80 self._requests.append(ambient_clock.monotonic())
81 return True
82 return False
84 @property
85 def current_requests(self) -> int:
86 """Get current number of requests in the window."""
87 self._cleanup_old_requests()
88 return len(self._requests)
91class CacheBackendLimiter:
92 """Fixed-window rate limiter backed by any CacheBackendProtocol.
94 Uses a fixed-window counter pattern: ``get`` → increment → ``set``.
95 This is not atomic (race between get and set), but the window is short
96 and the error margin is bounded to at most ``max_requests`` extra
97 requests per process per window. For strict distributed enforcement
98 use the Redis-backed limiter instead.
100 Args:
101 cache: Any ``CacheBackendProtocol`` implementation (Redis, Memcached, …).
102 window_size: Window duration in seconds.
103 max_requests: Maximum requests allowed per window.
104 """
106 def __init__(
107 self,
108 cache: CacheBackendProtocol,
109 window_size: float,
110 max_requests: int,
111 ) -> None:
112 self._cache = cache
113 self.window_size = window_size
114 self.max_requests = max_requests
116 async def try_acquire(self, key: str) -> tuple[bool, int]:
117 """Attempt to acquire a rate-limit slot.
119 Returns:
120 Tuple of (allowed, remaining) where remaining is tokens left in
121 the current window.
122 """
123 current_time = ambient_clock.timestamp()
124 bucket = int(current_time // self.window_size)
125 cache_key = f"rl:{key}:{bucket}"
126 raw_result = await self._cache.get(cache_key)
127 raw = raw_result.unwrap_or(None) if raw_result.is_ok() else None
128 count: int = int(raw) if raw is not None else 0
129 if count >= self.max_requests:
130 return False, 0
131 # Increment and reset TTL; not atomic but bounded error
132 new_count = count + 1
133 await self._cache.set(cache_key, new_count, ttl=int(self.window_size) + 1)
134 return True, self.max_requests - new_count
137class RateLimiter:
138 """Token bucket rate limiter using Redis.
140 Implements token bucket algorithm:
141 - Bucket has max capacity (burst limit)
142 - Tokens refill at steady rate
143 - Each request consumes 1 token
144 - Request rejected if no tokens available
146 Features:
147 - Per-user limits (authenticated requests)
148 - Per-IP limits (anonymous requests)
149 - Per-endpoint limits
150 - Sliding window for accurate rate calculation
152 Example:
153 >>> limiter = RateLimiter(redis_client)
154 >>>
155 >>> @app.post("/api/expensive")
156 >>> async def expensive_endpoint(request: Request):
157 ... await limiter.check_rate_limit(
158 ... request=request,
159 ... max_requests=10,
160 ... window_seconds=60,
161 ... )
162 ... return await do_expensive_work()
163 """
165 def __init__(
166 self,
167 redis_client: Any | None = None,
168 trusted_proxies: frozenset[str] | None = None,
169 cache_backend: CacheBackendProtocol | None = None,
170 ) -> None:
171 """Initialize rate limiter.
173 Args:
174 redis_client: Redis client for atomic Lua-script based sliding
175 window limiting. Preferred for production.
176 trusted_proxies: Set of trusted proxy IP addresses whose
177 ``X-Forwarded-For`` header is honoured. When *None*
178 (the default) the header is **ignored** and the direct
179 peer address from ``request.client.host`` is used
180 instead. Set to a non-empty ``frozenset`` when the
181 application runs behind a known reverse proxy.
182 cache_backend: ``CacheBackendProtocol`` used for fixed-window distributed
183 rate limiting when no ``redis_client`` is provided. Shared
184 across all workers that share the same cache; slightly less
185 precise than the Redis Lua path but requires no raw Redis
186 dependency.
187 """
188 self.redis = redis_client
189 self._cache_backend = cache_backend
190 self._trusted_proxies = trusted_proxies
191 self._memory_limiters: dict[str, MemoryLimiter] = {}
192 self._cache_limiters: dict[str, CacheBackendLimiter] = {}
193 if not redis_client and not cache_backend:
194 logger.warning(
195 "rate_limiter_using_in_memory_store",
196 extra={
197 "reason": "No Redis client or CacheBackendProtocol provided. Rate limiting state is "
198 "process-local and NOT shared across workers or instances. "
199 "Provide a redis_client or cache_backend for production multi-process deployments."
200 },
201 )
203 async def check_rate_limit(
204 self,
205 request: Request,
206 *,
207 max_requests: int,
208 window_seconds: int,
209 scope: Literal["user", "ip", "endpoint"] = "user",
210 ) -> None:
211 """Check if request is within rate limit.
213 Args:
214 request: Starlette request.
215 max_requests: Maximum requests per window.
216 window_seconds: Time window in seconds.
217 scope: Rate limit scope (user, ip, or endpoint).
219 Raises:
220 RateLimitError: If rate limit exceeded.
222 Example:
223 >>> # Allow 100 requests per minute per user
224 >>> await limiter.check_rate_limit(
225 ... request=request,
226 ... max_requests=100,
227 ... window_seconds=60,
228 ... scope="user",
229 ... )
230 """
231 # Get rate limit key based on scope
232 key = self._get_rate_limit_key(request, scope)
234 # Check using sliding window
235 now = ambient_clock.now()
236 window_start = now - timedelta(seconds=window_seconds)
238 # Redis key for this rate limit
239 redis_key = f"rate_limit:{key}"
241 # Use Lua script for atomic operation
242 lua_script = """
243 local key = KEYS[1]
244 local now = tonumber(ARGV[1])
245 local window_start = tonumber(ARGV[2])
246 local max_requests = tonumber(ARGV[3])
247 local window_seconds = tonumber(ARGV[4])
249 -- Remove old entries outside window
250 redis.call('ZREMRANGEBYSCORE', key, 0, window_start)
252 -- Count requests in current window
253 local current_count = redis.call('ZCARD', key)
255 if current_count < max_requests then
256 -- Add current request
257 redis.call('ZADD', key, now, now)
258 redis.call('EXPIRE', key, window_seconds)
259 return {1, max_requests - current_count - 1}
260 else
261 -- Rate limit exceeded
262 local oldest = redis.call('ZRANGE', key, 0, 0, 'WITHSCORES')
263 local retry_after = math.ceil(oldest[2] + window_seconds - now)
264 return {0, retry_after}
265 end
266 """
268 if self.redis:
269 result = await self.redis.eval(
270 lua_script,
271 1,
272 redis_key,
273 now.timestamp(),
274 window_start.timestamp(),
275 max_requests,
276 window_seconds,
277 )
278 allowed = result[0]
279 remaining = result[1]
280 elif self._cache_backend is not None:
281 # CacheBackendProtocol-backed fixed-window distributed counter
282 limiter_key = f"{key}:{max_requests}:{window_seconds}"
283 if limiter_key not in self._cache_limiters:
284 self._cache_limiters[limiter_key] = CacheBackendLimiter(
285 cache=self._cache_backend,
286 window_size=window_seconds,
287 max_requests=max_requests,
288 )
289 cache_limiter = self._cache_limiters[limiter_key]
290 allowed, remaining = await cache_limiter.try_acquire(key)
291 result = [0, window_seconds] if not allowed else [1, remaining]
292 else:
293 # In-memory fallback
294 limiter_key = f"{key}:{max_requests}:{window_seconds}"
295 if limiter_key not in self._memory_limiters:
296 self._memory_limiters[limiter_key] = MemoryLimiter(
297 window_size=window_seconds,
298 max_requests=max_requests,
299 )
301 limiter = self._memory_limiters[limiter_key]
302 allowed = await limiter.try_acquire()
303 remaining = max_requests - limiter.current_requests
305 if not allowed:
306 retry_after = window_seconds
307 result = [0, retry_after]
308 else:
309 result = [1, remaining]
311 if not allowed:
312 retry_after = int(result[1])
314 logger.warning(
315 "security.rate_limit_exceeded",
316 key=key,
317 max_requests=max_requests,
318 window_seconds=window_seconds,
319 retry_after=retry_after,
320 )
322 raise RateLimitError(details={"retry_after": retry_after})
324 remaining = result[1]
326 # Add rate limit headers
327 request.state.rate_limit_remaining = remaining
328 request.state.rate_limit_limit = max_requests
329 request.state.rate_limit_reset = int(
330 (now + timedelta(seconds=window_seconds)).timestamp()
331 )
333 logger.debug(
334 "rate_limit.check_passed",
335 key=key,
336 remaining=remaining,
337 limit=max_requests,
338 )
340 def _get_rate_limit_key(
341 self,
342 request: Request,
343 scope: Literal["user", "ip", "endpoint"],
344 ) -> str:
345 """Get rate limit key based on scope.
347 Args:
348 request: Starlette request.
349 scope: Rate limit scope.
351 Returns:
352 Rate limit key.
353 """
354 if scope == "user":
355 # Use user ID if authenticated
356 user_id = getattr(request.state, "user_id", None)
357 if user_id:
358 return f"user:{user_id}"
359 # Fall back to IP
360 return f"ip:{self._get_client_ip(request)}"
362 if scope == "ip":
363 return f"ip:{self._get_client_ip(request)}"
365 if scope == "endpoint":
366 # Per-endpoint limit
367 path = request.url.path
368 method = request.method
369 return f"endpoint:{method}:{path}"
371 raise ValueError(f"Invalid scope: {scope}")
373 def _get_client_ip(self, request: Request) -> str:
374 """Get client IP address.
376 The ``X-Forwarded-For`` header is read **only** when the direct
377 peer address (``request.client.host``) is in the instance's
378 ``trusted_proxies`` set. Trusting the header unconditionally
379 makes IP-based rate limits trivially spoofable by any client
380 that sets the header itself.
382 Args:
383 request: Starlette request.
385 Returns:
386 Client IP address (string).
387 """
388 direct_ip = request.client.host if request.client else None
390 # Only honour X-Forwarded-For when the request arrives from a
391 # known trusted proxy and the header is actually present.
392 if self._trusted_proxies is not None and direct_ip in self._trusted_proxies:
393 forwarded = request.headers.get("X-Forwarded-For")
394 if forwarded:
395 return forwarded.split(",")[0].strip()
397 return direct_ip or "unknown"
400def rate_limit(
401 max_requests: int,
402 window_seconds: int,
403 scope: Literal["user", "ip", "endpoint"] = "user",
404 limiter: WebRateLimiterProtocol | None = None,
405) -> Callable:
406 """Decorator for rate limiting endpoints.
408 Enforces rate limits using the provided ``limiter`` instance (constructor injection).
409 The ``limiter`` must be explicitly passed; no runtime container resolution is performed.
411 If no limiter is provided, the endpoint executes without rate limiting and a warning
412 is logged.
414 Args:
415 max_requests: Maximum requests per window.
416 window_seconds: Time window in seconds.
417 scope: Rate limit scope (user, ip, or endpoint).
418 limiter: Required pre-constructed :class:`WebRateLimiterProtocol` instance.
419 Must be injected at decoration time.
421 Returns:
422 Decorator function.
424 Example:
425 >>> my_limiter = RateLimiter(redis_client=redis)
426 >>> @app.post("/api/upload")
427 >>> @rate_limit(max_requests=5, window_seconds=60, limiter=my_limiter)
428 >>> async def upload_file(request: Request):
429 ... return await process_upload()
430 """
431 _injected_limiter = limiter # captured once at decoration time
433 def decorator(func: Callable) -> Callable:
434 @wraps(func)
435 async def wrapper(*args: Any, **kwargs: Any) -> Any:
436 # ── 1. Locate the request object ───────────────────────────────
437 request: Request | None = None
438 for arg in args:
439 if (
440 hasattr(arg, "scope")
441 and hasattr(arg, "app")
442 and "method" in arg.scope
443 ):
444 request = arg
445 break
446 if request is None:
447 request = kwargs.get("request")
449 if request is None:
450 logger.warning(
451 "rate_limit_skipped_no_request",
452 endpoint=func.__name__,
453 )
454 return await func(*args, **kwargs)
456 # ── 2. Resolve the limiter (only via explicit injection) ─
457 resolved_limiter: WebRateLimiterProtocol | None = _injected_limiter
459 if resolved_limiter is None:
460 logger.warning(
461 "rate_limit_skipped_no_limiter",
462 endpoint=func.__name__,
463 reason="no limiter explicitly injected; use rate_limit(limiter=...) to provide one",
464 )
465 return await func(*args, **kwargs)
467 # ── 3. Enforce the rate limit ───────────────────────────────────
468 await resolved_limiter.check_rate_limit(
469 request=request,
470 max_requests=max_requests,
471 window_seconds=window_seconds,
472 scope=scope,
473 )
475 return await func(*args, **kwargs)
477 return wrapper
479 return decorator
482class RateLimitMiddleware:
483 """Middleware to add rate limit headers to all responses.
485 Example:
486 >>> app.add_middleware(RateLimitMiddleware)
487 """
489 def __init__(
490 self,
491 app: ASGIApp,
492 *,
493 rate_limiter: RateLimiter | None = None,
494 config: Any | None = None,
495 ) -> None:
496 """Initialize middleware.
498 Args:
499 app: ASGI application.
500 rate_limiter: Rate limiter instance (optional).
501 config: ``RateLimitConfig`` driving per-path rules and defaults.
502 When ``enabled`` and a limiter is present, ``__call__``
503 enforces the matched rule (or the default limit) instead of
504 only stamping headers.
505 """
506 self.app = app
507 self.rate_limiter = rate_limiter
508 self.config = config
510 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
511 """Process request and add rate limit headers.
513 Args:
514 scope: ASGI scope
515 receive: ASGI receive
516 send: ASGI send
517 """
518 if scope["type"] != "http":
519 await self.app(scope, receive, send)
520 return
522 # ── Enforce before forwarding ─────────────────────────────────────
523 # The existing check_rate_limit() algorithm (Redis Lua / cache /
524 # in-memory) is the single enforcement engine. A breached request
525 # gets the standard 429 + Retry-After response (RateLimitError is
526 # caught here because middleware raises bypass Starlette's inner
527 # ExceptionMiddleware on supported versions and would surface as 500).
528 if (
529 self.rate_limiter is not None
530 and self.config is not None
531 and self.config.enabled
532 ):
533 request = Request(scope, receive)
534 if (
535 self.config.whitelist_ips
536 and request.client
537 and request.client.host in self.config.whitelist_ips
538 ):
539 pass # whitelisted — skip enforcement (D2)
540 else:
541 rule = self.config.get_rule(request.url.path)
542 if rule is not None:
543 max_requests, window_seconds = rule.requests, rule.window
544 else:
545 max_requests = self.config.default_limit
546 window_seconds = self.config.default_window
547 try:
548 await self.rate_limiter.check_rate_limit(
549 request,
550 max_requests=max_requests,
551 window_seconds=window_seconds,
552 scope="user",
553 )
554 except RateLimitError as exc:
555 response = _rate_limit_429_response(exc)
556 await response(scope, receive, send)
557 return
559 async def send_with_headers(message: Any) -> None:
560 if message["type"] == "http.response.start":
561 headers = list(message.get("headers", []))
563 # Add rate limit headers if available (stored in request.state)
564 # scope["state"] is a plain dict in Starlette; State is a wrapper
565 state = scope.get("state")
566 if isinstance(state, dict):
567 remaining = state.get("rate_limit_remaining")
568 limit = state.get("rate_limit_limit")
569 reset = state.get("rate_limit_reset")
570 elif state is not None:
571 remaining = getattr(state, "rate_limit_remaining", None)
572 limit = getattr(state, "rate_limit_limit", None)
573 reset = getattr(state, "rate_limit_reset", None)
574 else:
575 remaining = limit = reset = None
577 if remaining is not None:
578 headers.append((b"x-ratelimit-remaining", str(remaining).encode()))
579 if limit is not None:
580 headers.append((b"x-ratelimit-limit", str(limit).encode()))
581 if reset is not None:
582 headers.append((b"x-ratelimit-reset", str(reset).encode()))
584 message["headers"] = headers
586 await send(message)
588 await self.app(scope, receive, send_with_headers)