fix: restore rate limiting under FastAPI 0.138 and await the async limiter

slowapi/extension.py

@@ -24,6 +24,7 @@
     Union,
 )
 
+import anyio
 from limits import RateLimitItem  # type: ignore
 from limits.errors import ConfigurationError  # type: ignore
 from limits.storage import storage_from_string  # type: ignore
@@ -74,15 +75,17 @@ class HEADERS:
 MAX_BACKEND_CHECKS = 5
 
 
-def _rate_limit_exceeded_handler(request: Request, exc: RateLimitExceeded) -> Response:
+async def _rate_limit_exceeded_handler(
+    request: Request, exc: RateLimitExceeded
+) -> Response:
     """
     Build a simple JSON response that includes the details of the rate limit
     that was hit. If no limit is hit, the countdown is added to headers.
     """
     response = JSONResponse(
         {"error": f"Rate limit exceeded: {exc.detail}"}, status_code=429
     )
-    response = request.app.state.limiter._inject_headers(
+    response = await request.app.state.limiter._inject_headers(
         response, request.state.view_rate_limit
     )
     return response
@@ -376,7 +379,7 @@ def limiter(self) -> RateLimiter:
         else:
             return self._limiter
 
-    def _inject_headers(
+    async def _inject_headers(
         self, response: Response, current_limit: Tuple[RateLimitItem, List[str]]
     ) -> Response:
         if self.enabled and self._headers_enabled and current_limit is not None:
@@ -385,7 +388,7 @@ def _inject_headers(
                     "parameter `response` must be an instance of starlette.responses.Response"
                 )
             try:
-                window_stats: Tuple[int, int] = self.limiter.get_window_stats(
+                window_stats: Tuple[int, int] = await self.limiter.get_window_stats(
                     current_limit[0], *current_limit[1]
                 )
                 reset_in = 1 + window_stats[0]
@@ -420,7 +423,7 @@ def _inject_headers(
                         " in-memory storage"
                     )
                     self._storage_dead = True
-                    response = self._inject_headers(response, current_limit)
+                    response = await self._inject_headers(response, current_limit)
                 if self._swallow_errors:
                     self.logger.exception(
                         "Failed to update rate limit headers. Swallowing error"
@@ -750,18 +753,20 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Response:
                         if self._auto_check and not getattr(
                             request.state, "_rate_limiting_complete", False
                         ):
-                            self._check_request_limit(request, func, False)
+                            await self._check_request_limit(request, func, False)
                             request.state._rate_limiting_complete = True
                     response = await func(*args, **kwargs)  # type: ignore
-                    if self.enabled:
+                    # `_inject_headers` is a no-op when headers are disabled, and
+                    # reaching `view_rate_limit` is not always safe, so skip it.
+                    if self.enabled and self._headers_enabled:
                         if not isinstance(response, Response):
                             # get the response object from the decorated endpoint function
-                            self._inject_headers(
+                            await self._inject_headers(
                                 kwargs.get("response"),  # type: ignore
                                 request.state.view_rate_limit,
                             )
                         else:
-                            self._inject_headers(
+                            await self._inject_headers(
                                 response, request.state.view_rate_limit
                             )
                     return response
@@ -783,19 +788,29 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Response:
                         if self._auto_check and not getattr(
                             request.state, "_rate_limiting_complete", False
                         ):
-                            self._check_request_limit(request, func, False)
+                            # Rate limiting is asynchronous, and a sync endpoint runs
+                            # in an anyio worker thread, so hand the coroutine back to
+                            # the event loop driving the request.
+                            anyio.from_thread.run(
+                                self._check_request_limit, request, func, False
+                            )
                             request.state._rate_limiting_complete = True
                     response = func(*args, **kwargs)
-                    if self.enabled:
+                    # Same as above, and here the skipped call also saves a second
+                    # round trip to the event loop.
+                    if self.enabled and self._headers_enabled:
                         if not isinstance(response, Response):
                             # get the response object from the decorated endpoint function
-                            self._inject_headers(
+                            anyio.from_thread.run(
+                                self._inject_headers,
                                 kwargs.get("response"),
                                 request.state.view_rate_limit,  # type: ignore
                             )
                         else:
-                            self._inject_headers(
-                                response, request.state.view_rate_limit
+                            anyio.from_thread.run(
+                                self._inject_headers,
+                                response,
+                                request.state.view_rate_limit,
                             )
                     return response
 

slowapi/middleware.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import Callable, Iterable, Optional, Tuple
+from typing import Any, Callable, Iterable, Iterator, Optional, Tuple
 
 from starlette.applications import Starlette
 from starlette.datastructures import MutableHeaders
@@ -15,14 +15,50 @@
 from slowapi import Limiter, _rate_limit_exceeded_handler
 
 
+def _flatten_routes(routes: Iterable[BaseRoute]) -> Iterator[Any]:
+    """
+    Yield the matchable routes of an app.
+
+    Since FastAPI 0.138, `include_router` no longer copies the sub routes into the
+    parent app: it leaves a single `fastapi.routing._IncludedRouter` node holding
+    the original router. Such a node has no `endpoint`, so we expand it into its
+    effective (prefix-aware) route contexts, which expose both `matches()` and the
+    original route. Anything else is yielded as is.
+
+    Upstream bug: https://github.com/laurentS/slowapi/issues/281 (unfixed as of
+    0.1.10; PRs 282, 285 and 286 are open). Unlike those, this also covers plain
+    Starlette routes added to an included router, whose context carries no
+    `endpoint` of its own.
+    """
+    for route in routes:
+        effective_route_contexts = getattr(route, "effective_route_contexts", None)
+        if callable(effective_route_contexts):
+            yield from effective_route_contexts()
+        else:
+            yield route
+
+
+def _route_endpoint(route: Any) -> Optional[Callable]:
+    """
+    The endpoint function of a matchable route, or None if it has no endpoint
+    (a `Mount` for instance).
+    """
+    # `_EffectiveRouteContext` wraps the route it was built from, and only carries
+    # an `endpoint` of its own for API routes.
+    route = getattr(route, "original_route", route)
+    return getattr(route, "endpoint", None)
+
+
 def _find_route_handler(
     routes: Iterable[BaseRoute], scope: Scope
 ) -> Optional[Callable]:
     handler = None
-    for route in routes:
+    for route in _flatten_routes(routes):
         match, _ = route.matches(scope)
-        if match == Match.FULL and hasattr(route, "endpoint"):
-            handler = route.endpoint  # type: ignore
+        if match == Match.FULL:
+            endpoint = _route_endpoint(route)
+            if endpoint is not None:
+                handler = endpoint
     return handler
 
 
@@ -56,27 +92,6 @@ async def _check_limits(
     return None, False, None
 
 
-def sync_check_limits(
-    limiter: Limiter, request: Request, handler: Optional[Callable], app: Starlette
-) -> Tuple[Optional[Response], bool]:
-    """
-    Returns a `Response` object if an error occurred, as well as a boolean to know
-    whether we should inject headers or not.
-    Used in our WSGI middleware, it only supports synchronous exception_handler.
-    This will fallback on _rate_limit_exceeded_handler otherwise.
-    """
-    exception_handler, _bool, exc = _check_limits(limiter, request, handler, app)
-    if not exception_handler or not exc:
-        return None, _bool
-
-    # cannot execute asynchronous code in a synchronous middleware,
-    # -> fallback on default exception handler
-    if inspect.iscoroutinefunction(exception_handler):
-        exception_handler = _rate_limit_exceeded_handler
-
-    return exception_handler(request, exc), _bool  # type: ignore
-
-
 async def async_check_limits(
     limiter: Limiter, request: Request, handler: Optional[Callable], app: Starlette
 ) -> Tuple[Optional[Response], bool]:
@@ -127,7 +142,7 @@ async def dispatch(
         if _should_exempt(limiter, handler):
             return await call_next(request)
 
-        error_response, should_inject_headers = sync_check_limits(
+        error_response, should_inject_headers = await async_check_limits(
             limiter, request, handler, app
         )
         if error_response is not None:
@@ -137,7 +152,7 @@ async def dispatch(
         if should_inject_headers:
             view_rate_limit = getattr(request.state, "view_rate_limit", None)
             if view_rate_limit is not None:
-                response = limiter._inject_headers(response, view_rate_limit)
+                response = await limiter._inject_headers(response, view_rate_limit)
         return response
 
 

tests/__init__.py

@@ -15,7 +15,7 @@
 
 async def _async_rate_limit_exceeded_handler(request: Request, exc: RateLimitExceeded):
     await asyncio.sleep(0)
-    return _rate_limit_exceeded_handler(request, exc)
+    return await _rate_limit_exceeded_handler(request, exc)
 
 
 class TestSlowapi:
@@ -31,6 +31,7 @@ def _factory(config={}, **limiter_args):
             middleware, exception_handler = request.param
 
             limiter_args.setdefault("key_func", get_remote_address)
+            limiter_args.setdefault("storage_uri", "async+memory://")
             limiter = Limiter(**limiter_args)
             app = Starlette(debug=True)
             app.state.limiter = limiter
@@ -55,6 +56,7 @@ def build_fastapi_app(self, request):
         def _factory(config={}, **limiter_args):
             middleware, exception_handler = request.param
             limiter_args.setdefault("key_func", get_remote_address)
+            limiter_args.setdefault("storage_uri", "async+memory://")
             limiter = Limiter(**limiter_args)
             app = FastAPI()
             app.state.limiter = limiter

tests/test_fastapi_extension.py

@@ -1,9 +1,13 @@
 import hiro  # type: ignore
 import pytest  # type: ignore
+from fastapi import APIRouter, FastAPI
 from starlette.requests import Request
 from starlette.responses import PlainTextResponse, Response
 from starlette.testclient import TestClient
 
+from slowapi.errors import RateLimitExceeded
+from slowapi.extension import Limiter, _rate_limit_exceeded_handler
+from slowapi.middleware import SlowAPIASGIMiddleware
 from slowapi.util import get_ipaddr
 from tests import TestSlowapi
 
@@ -369,3 +373,136 @@ async def t1_func(my_param: str, request: Request):
                 )
                 == 2
             )
+
+
+class TestIncludedRouters:
+    """
+    Routes added through `include_router` must still be found by the middleware.
+
+    Since FastAPI 0.138, `app.routes` holds an opaque `_IncludedRouter` node per
+    `include_router` call instead of the flattened sub routes. A middleware that
+    only looks at the top level finds no handler, and slowapi then treats every
+    request as exempt.
+
+    These tests only cover `SlowAPIASGIMiddleware`: `SlowAPIMiddleware` does not
+    await `_check_limits` and is broken for unrelated reasons.
+    """
+
+    def build_app(self, **limiter_args):
+        limiter_args.setdefault("key_func", lambda: "mock")
+        limiter_args.setdefault("storage_uri", "async+memory://")
+        limiter = Limiter(**limiter_args)
+        app = FastAPI()
+        app.state.limiter = limiter
+        app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+        app.add_middleware(SlowAPIASGIMiddleware)
+        return app, limiter
+
+    def test_default_limits_apply_to_included_router(self):
+        app, limiter = self.build_app(default_limits=["2/minute"])
+        router = APIRouter()
+
+        @router.get("/t1")
+        async def t1(request: Request):
+            return PlainTextResponse("test")
+
+        app.include_router(router, prefix="/v0")
+
+        client = TestClient(app)
+        assert [client.get("/v0/t1").status_code for _ in range(3)] == [200, 200, 429]
+
+    def test_default_limits_apply_to_nested_included_router(self):
+        app, limiter = self.build_app(default_limits=["2/minute"])
+        inner = APIRouter()
+
+        @inner.get("/t1/{my_param}")
+        async def t1(my_param: str, request: Request):
+            return PlainTextResponse("test")
+
+        outer = APIRouter()
+        outer.include_router(inner, prefix="/inner")
+        app.include_router(outer, prefix="/v0")
+
+        client = TestClient(app)
+        assert [client.get("/v0/inner/t1/p").status_code for _ in range(3)] == [
+            200,
+            200,
+            429,
+        ]
+
+    def test_default_limits_apply_to_route_added_directly(self):
+        app, limiter = self.build_app(default_limits=["2/minute"])
+
+        @app.get("/t1")
+        async def t1(request: Request):
+            return PlainTextResponse("test")
+
+        client = TestClient(app)
+        assert [client.get("/t1").status_code for _ in range(3)] == [200, 200, 429]
+
+    def test_plain_starlette_route_in_included_router(self):
+        app, limiter = self.build_app(default_limits=["2/minute"])
+        router = APIRouter()
+
+        async def t1(request: Request):
+            return PlainTextResponse("test")
+
+        router.add_route("/t1", t1, methods=["GET"])
+        app.include_router(router, prefix="/v0")
+
+        client = TestClient(app)
+        assert [client.get("/v0/t1").status_code for _ in range(3)] == [200, 200, 429]
+
+    def test_exempt_route_in_included_router(self):
+        app, limiter = self.build_app(default_limits=["2/minute"])
+        router = APIRouter()
+
+        @router.get("/t1")
+        @limiter.exempt
+        async def t1(request: Request):
+            return PlainTextResponse("test")
+
+        app.include_router(router, prefix="/v0")
+
+        client = TestClient(app)
+        assert [client.get("/v0/t1").status_code for _ in range(3)] == [200, 200, 200]
+
+    def test_unknown_route_is_not_limited(self):
+        app, limiter = self.build_app(default_limits=["2/minute"])
+        router = APIRouter()
+
+        @router.get("/t1")
+        async def t1(request: Request):
+            return PlainTextResponse("test")
+
+        app.include_router(router, prefix="/v0")
+
+        client = TestClient(app)
+        assert [client.get("/nope").status_code for _ in range(3)] == [404, 404, 404]
+
+
+class TestHeaderInjection(TestSlowapi):
+    """
+    Header injection reads window stats from the storage, which is async in this
+    fork. A missing `await` there turns every breached request into a 500, and
+    silently marks the storage dead on the way out.
+    """
+
+    def test_middleware_429_goes_through_the_default_handler(self, build_fastapi_app):
+        app, limiter = build_fastapi_app(
+            key_func=lambda: "mock",
+            default_limits=["2/minute"],
+            headers_enabled=True,
+            in_memory_fallback_enabled=True,
+        )
+
+        @app.get("/t1")
+        async def t1(request: Request):
+            return PlainTextResponse("test")
+
+        client = TestClient(app)
+        responses = [client.get("/t1") for _ in range(3)]
+
+        assert [r.status_code for r in responses] == [200, 200, 429]
+        assert responses[0].headers["X-RateLimit-Limit"] == "2"
+        assert not limiter._storage_dead