|
1 | 1 | """ |
2 | 2 | Rate limiting shared by the app and its routers. |
3 | 3 |
|
4 | | -The Limiter lives in its own module because main.py imports the routers; a |
5 | | -router importing it from main.py would be a circular import. |
| 4 | +Built on the `limits` library: a fixed window per client IP and per route, |
| 5 | +counted in this process's memory. The Limiter lives in its own module because |
| 6 | +main.py imports the routers; a router importing it from main.py would be a |
| 7 | +circular import. |
6 | 8 |
|
7 | 9 | Limits are per client IP, as uvicorn sees it (behind a reverse proxy, make sure |
8 | 10 | uvicorn trusts the proxy's forwarded headers), and per process: with several |
9 | 11 | workers each keeps its own count, so the effective limit is multiplied. |
10 | 12 |
|
11 | | -A limited route answers with X-RateLimit-Limit, X-RateLimit-Remaining and |
12 | | -X-RateLimit-Reset (Unix time), so a client can pace itself before a 429. slowapi |
13 | | -writes them onto the route's `response: Response` parameter, which is why the |
14 | | -limited routes declare one they never touch. |
| 13 | +A limited route declares `Depends(limiter.limit("10/minute"))`. The dependency |
| 14 | +counts the request, writes X-RateLimit-Limit, X-RateLimit-Remaining and |
| 15 | +X-RateLimit-Reset (Unix time) onto the response so a client can pace itself |
| 16 | +before a 429, and raises RateLimitExceeded over the limit; main.py maps that to |
| 17 | +the 429 built by rate_limit_exceeded_handler. |
15 | 18 | """ |
16 | 19 |
|
17 | | -from fastapi import Request |
| 20 | +import math |
| 21 | +from collections.abc import Awaitable, Callable |
| 22 | + |
| 23 | +from fastapi import Request, Response |
18 | 24 | from fastapi.responses import JSONResponse |
19 | | -from limits import parse_many |
20 | | -from slowapi import Limiter |
21 | | -from slowapi.errors import RateLimitExceeded |
22 | | -from slowapi.util import get_remote_address |
| 25 | +from limits import RateLimitItem, parse_many |
| 26 | +from limits.storage import MemoryStorage |
| 27 | +from limits.strategies import FixedWindowRateLimiter |
23 | 28 |
|
24 | 29 | from .config import RATE_LIMIT_LOGIN, RATE_LIMIT_REGISTER |
25 | 30 |
|
|
33 | 38 | except ValueError as e: |
34 | 39 | raise RuntimeError(f"{_name}={_value!r} is not a valid rate limit: {e}") |
35 | 40 |
|
36 | | -limiter = Limiter(key_func=get_remote_address, headers_enabled=True) |
| 41 | + |
| 42 | +class RateLimitExceeded(Exception): |
| 43 | + """ |
| 44 | + A request went over one of its route's limits. |
| 45 | +
|
| 46 | + Attributes: |
| 47 | + limit: The limit that was hit. |
| 48 | + reset_time: Unix time at which that limit's window resets. |
| 49 | + """ |
| 50 | + |
| 51 | + def __init__(self, limit: RateLimitItem, reset_time: float) -> None: |
| 52 | + super().__init__(f"rate limit exceeded: {limit}") |
| 53 | + self.limit = limit |
| 54 | + self.reset_time = reset_time |
| 55 | + |
| 56 | + |
| 57 | +def client_address(request: Request) -> str: |
| 58 | + """ |
| 59 | + Key requests by client IP. |
| 60 | +
|
| 61 | + Args: |
| 62 | + request: The incoming request. |
| 63 | +
|
| 64 | + Returns: |
| 65 | + The peer address uvicorn reports, or "unknown" when there is none (an |
| 66 | + in-process test client, for instance). |
| 67 | + """ |
| 68 | + return request.client.host if request.client else "unknown" |
| 69 | + |
| 70 | + |
| 71 | +def route_scope(request: Request) -> str: |
| 72 | + """ |
| 73 | + Name the counter a request belongs to: its method and route path. |
| 74 | +
|
| 75 | + The route's path template (`/api/users/`) rather than the request URL, so a |
| 76 | + route with path parameters still has one counter. |
| 77 | +
|
| 78 | + Args: |
| 79 | + request: The incoming request, already matched to a route. |
| 80 | +
|
| 81 | + Returns: |
| 82 | + A string such as "POST /api/users/". |
| 83 | + """ |
| 84 | + route = request.scope.get("route") |
| 85 | + path = getattr(route, "path", request.url.path) |
| 86 | + return f"{request.method} {path}" |
| 87 | + |
| 88 | + |
| 89 | +class Limiter: |
| 90 | + """ |
| 91 | + Fixed-window counters per client and per route, in this process's memory. |
| 92 | +
|
| 93 | + Attributes: |
| 94 | + enabled: False makes every limit a no-op that writes no headers. The |
| 95 | + test suite turns it off because every test client shares one IP. |
| 96 | + key_func: Maps a request to its client key; client_address by default. |
| 97 | + """ |
| 98 | + |
| 99 | + def __init__( |
| 100 | + self, |
| 101 | + key_func: Callable[[Request], str] = client_address, |
| 102 | + *, |
| 103 | + enabled: bool = True, |
| 104 | + ) -> None: |
| 105 | + self.enabled = enabled |
| 106 | + self.key_func = key_func |
| 107 | + self._storage = MemoryStorage() |
| 108 | + self._strategy = FixedWindowRateLimiter(self._storage) |
| 109 | + |
| 110 | + def reset(self) -> None: |
| 111 | + """Forget every count.""" |
| 112 | + self._storage.reset() |
| 113 | + |
| 114 | + def limit(self, value: str) -> Callable[[Request, Response], Awaitable[None]]: |
| 115 | + """ |
| 116 | + Build the dependency that enforces `value` on the route declaring it. |
| 117 | +
|
| 118 | + Args: |
| 119 | + value: One or more limits in `limits` syntax ("10/minute", |
| 120 | + "100/hour;1000/day"), each of which must hold. |
| 121 | +
|
| 122 | + Returns: |
| 123 | + A dependency for the route's `dependencies=[Depends(...)]`. It |
| 124 | + raises RateLimitExceeded when a limit is hit; otherwise it writes |
| 125 | + the X-RateLimit headers of the limit with the least budget left. |
| 126 | + """ |
| 127 | + items = parse_many(value) |
| 128 | + |
| 129 | + async def enforce(request: Request, response: Response) -> None: |
| 130 | + if not self.enabled: |
| 131 | + return |
| 132 | + key = self.key_func(request) |
| 133 | + scope = route_scope(request) |
| 134 | + for item in items: |
| 135 | + if not self._strategy.hit(item, key, scope): |
| 136 | + stats = self._strategy.get_window_stats(item, key, scope) |
| 137 | + raise RateLimitExceeded(item, stats.reset_time) |
| 138 | + item, stats = min( |
| 139 | + ((i, self._strategy.get_window_stats(i, key, scope)) for i in items), |
| 140 | + key=lambda pair: pair[1].remaining, |
| 141 | + ) |
| 142 | + response.headers["X-RateLimit-Limit"] = str(item.amount) |
| 143 | + response.headers["X-RateLimit-Remaining"] = str(stats.remaining) |
| 144 | + response.headers["X-RateLimit-Reset"] = str(math.ceil(stats.reset_time)) |
| 145 | + |
| 146 | + return enforce |
| 147 | + |
| 148 | + |
| 149 | +limiter = Limiter() |
37 | 150 |
|
38 | 151 |
|
39 | 152 | def rate_limit_exceeded_handler(request: Request, exc: Exception) -> JSONResponse: |
40 | 153 | """ |
41 | 154 | Answer a limit hit in the API's usual error shape, with a Retry-After. |
42 | 155 |
|
43 | | - slowapi's default handler returns {"error": ...}; clients read {"detail": |
44 | | - ...} like every other error. The wait is the limit's window, the longest a |
45 | | - client has to wait under the default fixed-window strategy. |
| 156 | + Clients read {"detail": ...} like every other error. The wait is the |
| 157 | + limit's window, the longest a client has to wait under the fixed-window |
| 158 | + strategy; the X-RateLimit headers say when the window actually resets. |
46 | 159 |
|
47 | 160 | Args: |
48 | 161 | request: The request that hit the limit. |
49 | | - exc: The RateLimitExceeded raised by slowapi (typed as Exception |
| 162 | + exc: The RateLimitExceeded raised by the limiter (typed as Exception |
50 | 163 | because that is Starlette's handler signature). |
51 | 164 |
|
52 | 165 | Returns: |
53 | | - A 429 JSON response with a Retry-After header in seconds. |
| 166 | + A 429 JSON response with a Retry-After header in seconds and the |
| 167 | + X-RateLimit headers with no budget left. |
54 | 168 | """ |
55 | 169 | assert isinstance(exc, RateLimitExceeded) |
56 | | - wrapper = exc.limit |
57 | | - retry_after = wrapper.limit.get_expiry() if wrapper is not None else 60 |
| 170 | + retry_after = exc.limit.get_expiry() |
58 | 171 | return JSONResponse( |
59 | 172 | status_code=429, |
60 | 173 | content={"detail": f"Too many requests. Try again in {retry_after} seconds."}, |
61 | | - headers={"Retry-After": str(retry_after)}, |
| 174 | + headers={ |
| 175 | + "Retry-After": str(retry_after), |
| 176 | + "X-RateLimit-Limit": str(exc.limit.amount), |
| 177 | + "X-RateLimit-Remaining": "0", |
| 178 | + "X-RateLimit-Reset": str(math.ceil(exc.reset_time)), |
| 179 | + }, |
62 | 180 | ) |
0 commit comments