diff --git a/docs/index.md b/docs/index.md index 3a7697c..37d5e8a 100644 --- a/docs/index.md +++ b/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 diff --git a/slowapi/extension.py b/slowapi/extension.py index 2e6da1f..fa16222 100644 --- a/slowapi/extension.py +++ b/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 diff --git a/tests/test_fastapi_extension.py b/tests/test_fastapi_extension.py index 2994f02..eb3f189 100644 --- a/tests/test_fastapi_extension.py +++ b/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""" + )