Coverage for src/lexigram/auth/web/middleware/throttle.py: 79%
77 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-26 00:58 +0800
1"""Rate-limiting middleware for authentication endpoints.
3Provides in-memory IP-based rate limiting targeted at brute-force-prone
4auth endpoints (``/auth/login``, ``/auth/register``, etc.).
6Typical usage::
8 app = RateLimitMiddleware(app, rate_limit="5/minute", block_duration=300)
9"""
11from __future__ import annotations
13from collections import defaultdict
14from typing import TYPE_CHECKING, Any
16from lexigram import serialization as _json
17from lexigram.primitives import clock as ambient_clock
19if TYPE_CHECKING:
20 from collections.abc import Callable
22# Endpoints to which rate limiting is applied.
23_RATE_LIMITED_PATHS: frozenset[str] = frozenset(
24 {
25 "/auth/login",
26 "/auth/register",
27 "/auth/token",
28 "/auth/refresh",
29 "/auth/password-reset",
30 }
31)
34class RateLimitExceededError(Exception):
35 """Raised when a client has exceeded the configured rate limit.
37 Carries ``retry_after`` (seconds) so callers can propagate the value in
38 a ``Retry-After`` response header.
39 """
41 _code = "LEX_ERR_SEC_012"
43 def __init__(self, retry_after: int) -> None:
44 super().__init__(f"Rate limit exceeded. Retry after {retry_after}s.")
45 self.retry_after = retry_after
48def _parse_rate_limit(rate_limit: str) -> tuple[int, int]:
49 """Parse a rate-limit string such as ``"5/minute"`` into ``(max, seconds)``.
51 Supported periods: ``second`` / ``seconds``, ``minute`` / ``minutes``,
52 ``hour`` / ``hours``.
54 Args:
55 rate_limit: A string in the form ``"<count>/<period>"``.
57 Returns:
58 A ``(max_requests, window_seconds)`` tuple.
60 Raises:
61 ValueError: If the string cannot be parsed.
62 """
63 try:
64 count_str, period = rate_limit.split("/", maxsplit=1)
65 count = int(count_str.strip())
66 period = period.strip().lower()
67 except (ValueError, AttributeError) as exc:
68 raise ValueError(
69 f"Invalid rate_limit format {rate_limit!r}. "
70 "Expected '<count>/<period>', e.g. '5/minute'."
71 ) from exc
73 _period_map: dict[str, int] = {
74 "second": 1,
75 "seconds": 1,
76 "minute": 60,
77 "minutes": 60,
78 "hour": 3600,
79 "hours": 3600,
80 }
81 if period not in _period_map:
82 raise ValueError(
83 f"Unknown period {period!r}. "
84 f"Supported values: {', '.join(sorted(_period_map))}."
85 )
86 return count, _period_map[period]
89class RateLimitMiddleware:
90 """ASGI middleware that rate-limits authentication endpoints.
92 Applies a sliding-window rate limit to requests whose ``path`` matches
93 a known auth endpoint. Clients that exceed the limit receive a ``429``
94 response and are blocked for *block_duration* seconds. Successful
95 responses (HTTP 2xx) clear the attempt history for the client.
97 Args:
98 app: The next ASGI application in the chain.
99 rate_limit: Rate-limit string, e.g. ``"5/minute"``.
100 block_duration: How long (seconds) a client remains blocked after
101 exceeding the limit. Defaults to ``60``.
102 paths: Optional set of URL paths to which rate limiting is applied.
103 Falls back to the built-in :data:`_RATE_LIMITED_PATHS` set.
104 """
106 def __init__(
107 self,
108 app: Any,
109 rate_limit: str = "5/minute",
110 block_duration: int = 60,
111 paths: frozenset[str] | None = None,
112 cache_service: Any | None = None,
113 ) -> None:
114 self.app = app
115 self.block_duration = block_duration
116 self._paths = paths if paths is not None else _RATE_LIMITED_PATHS
117 # Optional external cache (e.g. Redis) for distributed rate limiting.
118 # Not used in the built-in in-memory implementation; reserved for
119 # future Redis-backed backends.
120 self.cache_service = cache_service
122 max_requests, window_seconds = _parse_rate_limit(rate_limit)
123 self._max_requests = max_requests
124 self._window_seconds = window_seconds
126 # Mutable state exposed for tests
127 self.attempts: dict[str, list[float]] = defaultdict(list)
128 self.blocked: dict[str, float] = {}
130 # ------------------------------------------------------------------
131 # ASGI interface
132 # ------------------------------------------------------------------
134 async def __call__(
135 self,
136 scope: dict[str, Any],
137 receive: Any,
138 send: Any,
139 ) -> None:
140 """Process an ASGI request."""
141 if scope.get("type") != "http":
142 await self.app(scope, receive, send)
143 return
145 path: str = scope.get("path", "")
146 if path not in self._paths:
147 await self.app(scope, receive, send)
148 return
150 client = scope.get("client")
151 client_id: str = client[0] if client else "unknown"
153 # --- get current time ---
154 now = ambient_clock.monotonic()
156 # --- check / enforce block ---
157 if client_id in self.blocked:
158 unblock_at = self.blocked[client_id]
159 remaining = int(unblock_at - now)
160 if remaining > 0:
161 await self._send_429(send, remaining)
162 return
163 # Block expired — release
164 del self.blocked[client_id]
166 # --- sliding-window check ---
167 cutoff = now - self._window_seconds
168 self.attempts[client_id] = [t for t in self.attempts[client_id] if t > cutoff]
169 if len(self.attempts[client_id]) >= self._max_requests:
170 # Block the client and respond 429
171 self.blocked[client_id] = now + self.block_duration
172 await self._send_429(send, self.block_duration)
173 return
175 # --- record this attempt and delegate to next app ---
176 self.attempts[client_id].append(now)
178 # Wrap send so we can inspect the response status code
179 response_status: list[int] = []
181 async def intercepting_send(message: dict[str, Any]) -> None:
182 if message.get("type") == "http.response.start":
183 response_status.append(message.get("status", 0))
184 await send(message)
186 await self.app(scope, receive, intercepting_send)
188 # On a successful response, clear the attempt counter
189 if response_status and 200 <= response_status[0] < 300:
190 self.attempts[client_id] = []
192 # ------------------------------------------------------------------
193 # Helpers
194 # ------------------------------------------------------------------
196 @staticmethod
197 async def _send_429(send: Any, retry_after: int) -> None:
198 """Send a 429 Too Many Requests ASGI response."""
199 body = _json.dumps(
200 {
201 "detail": "Too many requests. Please try again later.",
202 "retry_after": retry_after,
203 }
204 )
205 retry_bytes = str(retry_after).encode()
206 await send(
207 {
208 "type": "http.response.start",
209 "status": 429,
210 "headers": [
211 (b"content-type", b"application/json"),
212 (b"retry-after", retry_bytes),
213 ],
214 }
215 )
216 await send({"type": "http.response.body", "body": body, "more_body": False})
219def throttle(
220 rate_limit: str = "5/minute",
221 block_duration: int = 60,
222 paths: frozenset[str] | None = None,
223) -> Callable[[Any], RateLimitMiddleware]:
224 """Decorator / factory that wraps an ASGI app with :class:`RateLimitMiddleware`.
226 Can be used as a decorator::
228 @throttle(rate_limit="10/minute")
229 async def app(scope, receive, send): ...
231 Args:
232 rate_limit: Rate-limit string, e.g. ``"5/minute"``.
233 block_duration: Block duration in seconds.
234 paths: URL paths to rate-limit (defaults to built-in auth paths).
236 Returns:
237 A callable that wraps an ASGI app with the configured middleware.
238 """
240 def decorator(app: Any) -> RateLimitMiddleware:
241 return RateLimitMiddleware(
242 app,
243 rate_limit=rate_limit,
244 block_duration=block_duration,
245 paths=paths,
246 )
248 return decorator
251__all__ = [
252 "RateLimitExceededError",
253 "RateLimitMiddleware",
254 "throttle",
255]