diff --git a/slowapi/extension.py b/slowapi/extension.py index 059ba00..34a0e13 100644 --- a/slowapi/extension.py +++ b/slowapi/extension.py @@ -14,6 +14,7 @@ from typing import ( Any, Callable, Dict, + Literal, List, Optional, Set, @@ -123,6 +124,7 @@ class Limiter: * **enabled**: set to False to deactivate the limiter (default: True) * **config_filename**: name of the config file for Starlette from which to load settings for the rate limiter. Defaults to ".env". + * **key_style**: set to "url" to use the url, "endpoint" to use the view_func """ def __init__( @@ -143,6 +145,7 @@ class Limiter: key_prefix: str = "", enabled: bool = True, config_filename: Optional[str] = None, + key_style: Literal["endpoint", "url"] = "url", ) -> None: """ Configure the rate limiter at app level @@ -177,6 +180,7 @@ class Limiter: self._key_func = key_func self._key_prefix = key_prefix + self._key_style = key_style for limit in set(default_limits): self._default_limits.extend( @@ -547,13 +551,13 @@ class Limiter: Determine if the request is within limits """ endpoint = request["path"] or "" - # view_func = current_app.view_functions.get(endpoint, None) view_func = endpoint_func name = "%s.%s" % (view_func.__module__, view_func.__name__) if view_func else "" + _endpoint_key = endpoint if self._key_style == "url" else name # cases where we don't need to check the limits if ( - not endpoint + not _endpoint_key or not self.enabled # or we are sending a static file # or view_func == current_app.send_static_file @@ -608,7 +612,7 @@ class Limiter: ): all_limits += list(itertools.chain(*self._default_limits)) # actually check the limits, so far we've only computed the list of limits to check - self.__evaluate_limits(request, endpoint, all_limits) + self.__evaluate_limits(request, _endpoint_key, all_limits) except Exception as e: # no qa if isinstance(e, RateLimitExceeded): raise diff --git a/tests/test_fastapi_extension.py b/tests/test_fastapi_extension.py index 6ab72be..034890f 100644 --- a/tests/test_fastapi_extension.py +++ b/tests/test_fastapi_extension.py @@ -335,3 +335,27 @@ class TestDecorators(TestSlowapi): response = client.get("/t2", headers={"foo": "5"}) assert response.status_code == 200 if i < 6 else 429 + + @pytest.mark.parametrize( + "key_style, expected_key", + [ + ("url", "LIMITER/mock//t1/1/1/minute"), + ( + "endpoint", + "LIMITER/mock/tests.test_fastapi_extension.t1_func/1/1/minute", + ), + ], + ) + def test_key_style(self, key_style, expected_key): + app, limiter = self.build_fastapi_app( + key_func=lambda: "mock", key_style=key_style + ) + + @app.get("/t1") + @limiter.limit("1/minute") + async def t1_func(request: Request): + return PlainTextResponse("test") + + client = TestClient(app) + client.get("/t1", headers={"foo": "10"}) + assert limiter._storage.get(expected_key) == 1 diff --git a/tests/test_starlette_extension.py b/tests/test_starlette_extension.py index 1e97723..615e440 100644 --- a/tests/test_starlette_extension.py +++ b/tests/test_starlette_extension.py @@ -1,6 +1,7 @@ import time import hiro # type: ignore +import pytest # type: ignore from starlette.requests import Request from starlette.responses import PlainTextResponse from starlette.testclient import TestClient @@ -321,3 +322,29 @@ class TestDecorators(TestSlowapi): assert response.text == "test" else: assert "error" in response.json() + + @pytest.mark.parametrize( + "key_style, expected_key", + [ + ("url", "LIMITER/mock//t1/1/1/minute"), + ( + "endpoint", + "LIMITER/mock/tests.test_starlette_extension.t1_func/1/1/minute", + ), + ], + ) + def test_key_style(self, key_style, expected_key): + app, limiter = self.build_starlette_app( + key_func=lambda: "mock", key_style=key_style + ) + + @limiter.limit("1/minute") + async def t1_func(request: Request): + return PlainTextResponse("test") + + app.add_route("/t1", t1_func) + + client = TestClient(app) + client.get("/t1", headers={"foo": "10"}) + + assert limiter._storage.get(expected_key) == 1