fix: Address issue with decorator

The value returned by an endpoint may not be an instance or subclass
of `Response`, in this case FastAPI builds the actual Response
at an upper level in the middleware stack.
This patch allows the decorator to inspect the endpoint to retrieve the
associated `Response` to allow headers to be injected.

docs/index.md

@@ -51,6 +51,11 @@ The above app will have a route `t1` that will accept up to 5 requests per minut
     @limiter.limit("5/minute")
     async def homepage(request: Request):
         return PlainTextResponse("test")
+
+    @app.get("/mars")
+    @limiter.limit("5/minute")
+    async def homepage(request: Request, response: Response):
+        return {"key": "value"}
 ```
 
 This will provide the same result, but with a FastAPI app.
@@ -85,6 +90,17 @@ and not:
         pass
 ```
 
+  * Similarly, if the returned response is not an instance of `Response` and
+will be built at an upper level in the middleware stack, you'll need to provide
+the response object explicitly if you want the `Limiter` to modify the headers
+(`headers_enabled=True`):
+
+    ```python
+    @limiter.limit("5/minute")
+    async def myendpoint(request: Request, response: Response)
+        return {"key": "value"}
+    ```
+
   * `websocket` endpoints are not supported yet.
 
 # Developing and contributing

slowapi/extension.py

@@ -354,6 +354,10 @@ class Limiter:
         self, response: Response, current_limit: Tuple[RateLimitItem, List[str]]
     ) -> Response:
         if self.enabled and self._headers_enabled and current_limit is not None:
+            if not isinstance(response, Response):
+                raise Exception(
+                    "parameter `response` must be an instance of starlette.responses.Response"
+                )
             try:
                 window_stats: Tuple[int, int] = self.limiter.get_window_stats(
                     current_limit[0], *current_limit[1]
@@ -624,7 +628,11 @@ class Limiter:
                         self._check_request_limit(request, func, False)
                         request.state._rate_limiting_complete = True
                     response = await func(*args, **kwargs)  # type: ignore
-                    self._inject_headers(response, request.state.view_rate_limit)
+                    if self._headers_enabled and not isinstance(response, Response):
+                        # get the response object from the decorated endpoint function
+                        self._inject_headers(kwargs.get("response"), request.state.view_rate_limit)
+                    else:
+                        self._inject_headers(response, request.state.view_rate_limit)
                     return response
 
                 return async_wrapper
@@ -646,7 +654,11 @@ class Limiter:
                         self._check_request_limit(request, func, False)
                         request.state._rate_limiting_complete = True
                     response = func(*args, **kwargs)
-                    self._inject_headers(response, request.state.view_rate_limit)
+                    if not isinstance(response, Response):
+                        # get the response object from the decorated endpoint function
+                        self._inject_headers(kwargs.get("response"), request.state.view_rate_limit)
+                    else:
+                        self._inject_headers(response, request.state.view_rate_limit)
                     return response
 
                 return sync_wrapper

tests/test_fastapi_extension.py

@@ -1,7 +1,7 @@
 import hiro  # type: ignore
 import pytest  # type: ignore
 from starlette.requests import Request
-from starlette.responses import PlainTextResponse
+from starlette.responses import PlainTextResponse, Response
 from starlette.testclient import TestClient
 
 from slowapi.util import get_ipaddr
@@ -22,6 +22,50 @@ class TestDecorators(TestSlowapi):
             response = client.get("/t1")
             assert response.status_code == 200 if i < 5 else 429
 
+
+    def test_single_decorator_with_headers(self):
+        app, limiter = self.build_fastapi_app(key_func=get_ipaddr, headers_enabled=True)
+
+        @app.get("/t1")
+        @limiter.limit("5/minute")
+        async def t1(request: Request):
+            return PlainTextResponse("test")
+
+        client = TestClient(app)
+        for i in range(0, 10):
+            response = client.get("/t1")
+            assert response.status_code == 200 if i < 5 else 429
+            assert response.headers.get('X-RateLimit-Limit') is not None if i < 5 else True
+            assert response.headers.get('Retry-After') is not None if i < 5 else True
+
+    def test_single_decorator_not_response(self):
+        app, limiter = self.build_fastapi_app(key_func=get_ipaddr)
+
+        @app.get("/t1")
+        @limiter.limit("5/minute")
+        async def t1(request: Request, response: Response):
+            return {"key": "value"}
+
+        client = TestClient(app)
+        for i in range(0, 10):
+            response = client.get("/t1")
+            assert response.status_code == 200 if i < 5 else 429
+
+    def test_single_decorator_not_response_with_headers(self):
+        app, limiter = self.build_fastapi_app(key_func=get_ipaddr, headers_enabled=True)
+
+        @app.get("/t1")
+        @limiter.limit("5/minute")
+        async def t1(request: Request):
+            return PlainTextResponse("test")
+
+        client = TestClient(app)
+        for i in range(0, 10):
+            response = client.get("/t1")
+            assert response.status_code == 200 if i < 5 else 429
+            assert response.headers.get('X-RateLimit-Limit') is not None if i < 5 else True
+            assert response.headers.get('Retry-After') is not None if i < 5 else True
+
     def test_multiple_decorators(self):
         app, limiter = self.build_fastapi_app(key_func=get_ipaddr)
 
@@ -47,6 +91,31 @@ class TestDecorators(TestSlowapi):
                 == 429
             )
 
+    def test_multiple_decorators_not_response(self):
+        app, limiter = self.build_fastapi_app(key_func=get_ipaddr)
+
+        @app.get("/t1")
+        @limiter.limit(
+            "100 per minute", lambda: "test"
+        )  # effectively becomes a limit for all users
+        @limiter.limit("50/minute")  # per ip as per default key_func
+        async def t1(request: Request, response: Response):
+            return {"key":"value"}
+
+        with hiro.Timeline().freeze() as timeline:
+            cli = TestClient(app)
+            for i in range(0, 100):
+                response = cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.2"})
+                assert response.status_code == 200 if i < 50 else 429
+            for i in range(50):
+                assert cli.get("/t1").status_code == 200
+
+            assert cli.get("/t1").status_code == 429
+            assert (
+                cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
+                == 429
+            )
+
     def test_endpoint_missing_request_param(self):
         app, limiter = self.build_fastapi_app(key_func=get_ipaddr)
 
@@ -90,6 +159,21 @@ class TestDecorators(TestSlowapi):
             r"""parameter `request` must be an instance of starlette.requests.Request"""
         )
 
+    def test_endpoint_response_param_invalid(self):
+        app, limiter = self.build_fastapi_app(key_func=get_ipaddr, headers_enabled=True)
+
+        @app.get("/t4")
+        @limiter.limit("5/minute")
+        async def t4(request: Request, response: str = None):
+            return {"key": "value"}
+
+        with pytest.raises(Exception) as exc_info:
+            client = TestClient(app)
+            client.get("/t4")
+        assert exc_info.match(
+            r"""parameter `response` must be an instance of starlette.responses.Response"""
+        )
+
     def test_endpoint_request_param_invalid_sync(self):
         app, limiter = self.build_fastapi_app(key_func=get_ipaddr)
 
@@ -104,3 +188,18 @@ class TestDecorators(TestSlowapi):
         assert exc_info.match(
             r"""parameter `request` must be an instance of starlette.requests.Request"""
         )
+
+    def test_endpoint_response_param_invalid_sync(self):
+        app, limiter = self.build_fastapi_app(key_func=get_ipaddr, headers_enabled=True)
+
+        @app.get("/t5")
+        @limiter.limit("5/minute")
+        def t5(request: Request, response: str = None):
+            return {"key": "value"}
+
+        with pytest.raises(Exception) as exc_info:
+            client = TestClient(app)
+            client.get("/t5")
+        assert exc_info.match(
+            r"""parameter `response` must be an instance of starlette.responses.Response"""
+        )