diff --git a/docs/examples.md b/docs/examples.md index 6fe3518..1ee072e 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -101,4 +101,33 @@ or ```python app = Starlette() # or FastAPI() app.add_middleware(SlowAPIASGIMiddleware) -``` \ No newline at end of file +``` + +## Use view function's name instead of full endpoint as part of the storage key + +Let's use this route as an example: +```python +@app.route("/some_route/{some_param}") +def my_func(some_param): + ... +``` + +```python +limiter = Limiter(key_func=lambda: "mock", default_limits=["1/minute"], key_style="url") +``` + +When initializing the Limiter object with `key_style="url"`, it will use the full endpoint url as part of the storage key. + +When calling the `/some_route/my_param` endpoint would result with a key shaped like: `LIMITER/mock//some_route/my_param/1/1/minute`. + +> This means, that if the route contains some URL parameter, calling the endpoint with different parameters won't share the limitations. + +```python +limiter = Limiter(key_func=lambda: "mock", default_limits=["1/minute"], key_style="endpoint") +``` + +When initializing the Limiter object with `key_style="endpoint"`, it will use the function name as part of the storage key. + +When calling the `/some_route/my_param` endpoint would result with a key shaped like: `LIMITER/mock/{module}.my_func/1/1/minute` + +> This means, that if the route contains some URL parameter, calling the endpoint with different parameters will still share the limitations, since the view function is the same. diff --git a/slowapi/extension.py b/slowapi/extension.py index 059ba00..78bad1b 100644 --- a/slowapi/extension.py +++ b/slowapi/extension.py @@ -32,6 +32,7 @@ from starlette.config import Config from starlette.datastructures import MutableHeaders from starlette.requests import Request from starlette.responses import JSONResponse, Response +from typing_extensions import Literal from .errors import RateLimitExceeded from .wrappers import Limit, LimitGroup @@ -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( @@ -546,18 +550,20 @@ class Limiter: """ Determine if the request is within limits """ - endpoint = request["path"] or "" - # view_func = current_app.view_functions.get(endpoint, None) + endpoint_url = request["path"] or "" view_func = endpoint_func - name = "%s.%s" % (view_func.__module__, view_func.__name__) if view_func else "" + endpoint_func_name = ( + f"{view_func.__module__}.{view_func.__name__}" if view_func else "" + ) + _endpoint_key = endpoint_url if self._key_style == "url" else endpoint_func_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 - or name in self._exempt_routes + or endpoint_func_name in self._exempt_routes or any(fn() for fn in self._request_filters) ): return @@ -565,23 +571,27 @@ class Limiter: dynamic_limits: List[Limit] = [] if not in_middleware: - limits = self._route_limits[name] if name in self._route_limits else [] + limits = ( + self._route_limits[endpoint_func_name] + if endpoint_func_name in self._route_limits + else [] + ) dynamic_limits = [] - if name in self._dynamic_route_limits: - for lim in self._dynamic_route_limits[name]: + if endpoint_func_name in self._dynamic_route_limits: + for lim in self._dynamic_route_limits[endpoint_func_name]: try: dynamic_limits.extend(list(lim.with_request(request))) except ValueError as e: self.logger.error( "failed to load ratelimit for view function %s (%s)", - name, + endpoint_func_name, e, ) try: all_limits: List[Limit] = [] if self._storage_dead and self._fallback_limiter: - if in_middleware and name in self.__marked_for_limiting: + if in_middleware and endpoint_func_name in self.__marked_for_limiting: pass else: if self.__should_check_backend() and self._storage.check(): @@ -603,12 +613,15 @@ class Limiter: ) if ( not route_limits - and not (in_middleware and name in self.__marked_for_limiting) + and not ( + in_middleware + and endpoint_func_name in self.__marked_for_limiting + ) or combined_defaults ): 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..42e6322 100644 --- a/tests/test_fastapi_extension.py +++ b/tests/test_fastapi_extension.py @@ -335,3 +335,37 @@ 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", + ["url", "endpoint"], + ) + def test_key_style(self, build_fastapi_app, key_style): + app, limiter = build_fastapi_app(key_func=lambda: "mock", key_style=key_style) + + @app.get("/t1/{my_param}") + @limiter.limit("1/minute") + async def t1_func(my_param: str, request: Request): + return PlainTextResponse("test") + + client = TestClient(app) + client.get("/t1/param_one") + second_call = client.get("/t1/param_two") + # with the "url" key_style, since the `my_param` value changed, the storage key is different + # meaning it should not raise any RateLimitExceeded error. + if key_style == "url": + assert second_call.status_code == 200 + assert limiter._storage.get("LIMITER/mock//t1/param_one/1/1/minute") == 1 + assert limiter._storage.get("LIMITER/mock//t1/param_two/1/1/minute") == 1 + # However, with the `endpoint` key_style, it will use the function name (e.g: "t1_func") + # meaning it will raise a RateLimitExceeded error, because no matter the parameter value + # it will share the limitations. + elif key_style == "endpoint": + assert second_call.status_code == 429 + # check that we counted 2 requests, even though we had a different value for "my_param" + assert ( + limiter._storage.get( + "LIMITER/mock/tests.test_fastapi_extension.t1_func/1/1/minute" + ) + == 2 + ) diff --git a/tests/test_starlette_extension.py b/tests/test_starlette_extension.py index 1e97723..7f21c1d 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,38 @@ class TestDecorators(TestSlowapi): assert response.text == "test" else: assert "error" in response.json() + + @pytest.mark.parametrize( + "key_style", + ["url", "endpoint"], + ) + def test_key_style(self, build_starlette_app, key_style): + app, limiter = 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/{my_param}", t1_func) + + client = TestClient(app) + client.get("/t1/param_one") + second_call = client.get("/t1/param_two") + # with the "url" key_style, since the `my_param` value changed, the storage key is different + # meaning it should not raise any RateLimitExceeded error. + if key_style == "url": + assert second_call.status_code == 200 + assert limiter._storage.get("LIMITER/mock//t1/param_one/1/1/minute") == 1 + assert limiter._storage.get("LIMITER/mock//t1/param_two/1/1/minute") == 1 + # However, with the `endpoint` key_style, it will use the function name (e.g: "t1_func") + # meaning it will raise a RateLimitExceeded error, because no matter the parameter value + # it will share the limitations. + elif key_style == "endpoint": + assert second_call.status_code == 429 + # check that we counted 2 requests, even though we had a different value for "my_param" + assert ( + limiter._storage.get( + "LIMITER/mock/tests.test_starlette_extension.t1_func/1/1/minute" + ) + == 2 + )