From 17a514faef2bd0f4199e06f534ebf8c5aa5958cb Mon Sep 17 00:00:00 2001 From: thentgesMindee Date: Thu, 3 Nov 2022 16:46:30 +0100 Subject: [PATCH 1/8] new: :sparkles: add key_style parameter to choose between endpoint or url --- slowapi/extension.py | 10 +++++++--- tests/test_fastapi_extension.py | 24 ++++++++++++++++++++++++ tests/test_starlette_extension.py | 27 +++++++++++++++++++++++++++ 3 files changed, 58 insertions(+), 3 deletions(-) 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 From 031dc042d5a73f10ed038b88f771392cf8932fae Mon Sep 17 00:00:00 2001 From: thentgesMindee Date: Thu, 3 Nov 2022 17:13:48 +0100 Subject: [PATCH 2/8] chg: :pencil: add key_style example in docs --- docs/examples.md | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) 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. From b0d7bb03b725e593426f8bec730def1e648b2567 Mon Sep 17 00:00:00 2001 From: thentgesMindee Date: Mon, 7 Nov 2022 17:37:21 +0100 Subject: [PATCH 3/8] chg: :twisted_rightwards_arrows: update tests after rebase --- tests/test_fastapi_extension.py | 4 ++-- tests/test_starlette_extension.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_fastapi_extension.py b/tests/test_fastapi_extension.py index 034890f..e27c96f 100644 --- a/tests/test_fastapi_extension.py +++ b/tests/test_fastapi_extension.py @@ -346,8 +346,8 @@ class TestDecorators(TestSlowapi): ), ], ) - def test_key_style(self, key_style, expected_key): - app, limiter = self.build_fastapi_app( + def test_key_style(self, build_fastapi_app, key_style, expected_key): + app, limiter = build_fastapi_app( key_func=lambda: "mock", key_style=key_style ) diff --git a/tests/test_starlette_extension.py b/tests/test_starlette_extension.py index 615e440..54fcc86 100644 --- a/tests/test_starlette_extension.py +++ b/tests/test_starlette_extension.py @@ -333,8 +333,8 @@ class TestDecorators(TestSlowapi): ), ], ) - def test_key_style(self, key_style, expected_key): - app, limiter = self.build_starlette_app( + def test_key_style(self, build_starlette_app, key_style, expected_key): + app, limiter = build_starlette_app( key_func=lambda: "mock", key_style=key_style ) From 2ba9ed0630fdc5183951993553d166163edfaec9 Mon Sep 17 00:00:00 2001 From: thentgesMindee Date: Mon, 7 Nov 2022 17:47:56 +0100 Subject: [PATCH 4/8] fix: :art: apply black on tests files --- tests/test_fastapi_extension.py | 4 +--- tests/test_starlette_extension.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/test_fastapi_extension.py b/tests/test_fastapi_extension.py index e27c96f..cdebc90 100644 --- a/tests/test_fastapi_extension.py +++ b/tests/test_fastapi_extension.py @@ -347,9 +347,7 @@ class TestDecorators(TestSlowapi): ], ) def test_key_style(self, build_fastapi_app, key_style, expected_key): - app, limiter = build_fastapi_app( - key_func=lambda: "mock", key_style=key_style - ) + app, limiter = build_fastapi_app(key_func=lambda: "mock", key_style=key_style) @app.get("/t1") @limiter.limit("1/minute") diff --git a/tests/test_starlette_extension.py b/tests/test_starlette_extension.py index 54fcc86..5e2f336 100644 --- a/tests/test_starlette_extension.py +++ b/tests/test_starlette_extension.py @@ -334,9 +334,7 @@ class TestDecorators(TestSlowapi): ], ) def test_key_style(self, build_starlette_app, key_style, expected_key): - app, limiter = build_starlette_app( - key_func=lambda: "mock", key_style=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): From 39a18133d411d041484e2b074a7212b7f3ff2d62 Mon Sep 17 00:00:00 2001 From: thentgesMindee Date: Mon, 7 Nov 2022 17:58:44 +0100 Subject: [PATCH 5/8] fix: use Literal from typing_extension for python3.7 support --- slowapi/extension.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/slowapi/extension.py b/slowapi/extension.py index 34a0e13..938c949 100644 --- a/slowapi/extension.py +++ b/slowapi/extension.py @@ -14,7 +14,6 @@ from typing import ( Any, Callable, Dict, - Literal, List, Optional, Set, @@ -33,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 From 7f87a2620166a98314943fa5abd356d14790661a Mon Sep 17 00:00:00 2001 From: thentgesMindee Date: Mon, 7 Nov 2022 18:39:01 +0100 Subject: [PATCH 6/8] chg: :white_check_mark: update tests to reflect behavior explained in docs --- tests/test_fastapi_extension.py | 23 ++++++++++++++++++----- tests/test_starlette_extension.py | 22 +++++++++++++++++----- 2 files changed, 35 insertions(+), 10 deletions(-) diff --git a/tests/test_fastapi_extension.py b/tests/test_fastapi_extension.py index cdebc90..7009a22 100644 --- a/tests/test_fastapi_extension.py +++ b/tests/test_fastapi_extension.py @@ -339,7 +339,7 @@ class TestDecorators(TestSlowapi): @pytest.mark.parametrize( "key_style, expected_key", [ - ("url", "LIMITER/mock//t1/1/1/minute"), + ("url", "LIMITER/mock//t1/param_one/1/1/minute"), ( "endpoint", "LIMITER/mock/tests.test_fastapi_extension.t1_func/1/1/minute", @@ -349,11 +349,24 @@ class TestDecorators(TestSlowapi): def test_key_style(self, build_fastapi_app, key_style, expected_key): app, limiter = build_fastapi_app(key_func=lambda: "mock", key_style=key_style) - @app.get("/t1") + @app.get("/t1/{my_param}") @limiter.limit("1/minute") - async def t1_func(request: Request): + async def t1_func(my_param: str, request: Request): return PlainTextResponse("test") client = TestClient(app) - client.get("/t1", headers={"foo": "10"}) - assert limiter._storage.get(expected_key) == 1 + 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 + # also assert that we counted only one request on the expected key + assert limiter._storage.get(expected_key) == 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(expected_key) == 2 diff --git a/tests/test_starlette_extension.py b/tests/test_starlette_extension.py index 5e2f336..bde151f 100644 --- a/tests/test_starlette_extension.py +++ b/tests/test_starlette_extension.py @@ -326,7 +326,7 @@ class TestDecorators(TestSlowapi): @pytest.mark.parametrize( "key_style, expected_key", [ - ("url", "LIMITER/mock//t1/1/1/minute"), + ("url", "LIMITER/mock//t1/param_one/1/1/minute"), ( "endpoint", "LIMITER/mock/tests.test_starlette_extension.t1_func/1/1/minute", @@ -340,9 +340,21 @@ class TestDecorators(TestSlowapi): async def t1_func(request: Request): return PlainTextResponse("test") - app.add_route("/t1", t1_func) + app.add_route("/t1/{my_param}", t1_func) client = TestClient(app) - client.get("/t1", headers={"foo": "10"}) - - assert limiter._storage.get(expected_key) == 1 + 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 + # also assert that we counted only one request on the expected key + assert limiter._storage.get(expected_key) == 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(expected_key) == 2 From 1d9badb4a0ecab899f158ea983dd5b0bf8b1e99f Mon Sep 17 00:00:00 2001 From: thentgesMindee Date: Mon, 7 Nov 2022 18:47:25 +0100 Subject: [PATCH 7/8] chg: _check_request_limit | rename name, endpoint > endpoint_func_name, endpoint_url --- slowapi/extension.py | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/slowapi/extension.py b/slowapi/extension.py index 938c949..78bad1b 100644 --- a/slowapi/extension.py +++ b/slowapi/extension.py @@ -550,18 +550,20 @@ class Limiter: """ Determine if the request is within limits """ - endpoint = request["path"] or "" + endpoint_url = request["path"] or "" 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 + 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_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 @@ -569,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(): @@ -607,7 +613,10 @@ 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)) From 3ae7154271b6015dafc26a5fdb3392b3eb77c348 Mon Sep 17 00:00:00 2001 From: thentgesMindee Date: Tue, 8 Nov 2022 17:58:11 +0100 Subject: [PATCH 8/8] chg: :white_check_mark: minor change in testing --- tests/test_fastapi_extension.py | 23 +++++++++++------------ tests/test_starlette_extension.py | 23 +++++++++++------------ 2 files changed, 22 insertions(+), 24 deletions(-) diff --git a/tests/test_fastapi_extension.py b/tests/test_fastapi_extension.py index 7009a22..42e6322 100644 --- a/tests/test_fastapi_extension.py +++ b/tests/test_fastapi_extension.py @@ -337,16 +337,10 @@ class TestDecorators(TestSlowapi): assert response.status_code == 200 if i < 6 else 429 @pytest.mark.parametrize( - "key_style, expected_key", - [ - ("url", "LIMITER/mock//t1/param_one/1/1/minute"), - ( - "endpoint", - "LIMITER/mock/tests.test_fastapi_extension.t1_func/1/1/minute", - ), - ], + "key_style", + ["url", "endpoint"], ) - def test_key_style(self, build_fastapi_app, key_style, expected_key): + 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}") @@ -361,12 +355,17 @@ class TestDecorators(TestSlowapi): # meaning it should not raise any RateLimitExceeded error. if key_style == "url": assert second_call.status_code == 200 - # also assert that we counted only one request on the expected key - assert limiter._storage.get(expected_key) == 1 + 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(expected_key) == 2 + 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 bde151f..7f21c1d 100644 --- a/tests/test_starlette_extension.py +++ b/tests/test_starlette_extension.py @@ -324,16 +324,10 @@ class TestDecorators(TestSlowapi): assert "error" in response.json() @pytest.mark.parametrize( - "key_style, expected_key", - [ - ("url", "LIMITER/mock//t1/param_one/1/1/minute"), - ( - "endpoint", - "LIMITER/mock/tests.test_starlette_extension.t1_func/1/1/minute", - ), - ], + "key_style", + ["url", "endpoint"], ) - def test_key_style(self, build_starlette_app, key_style, expected_key): + 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") @@ -349,12 +343,17 @@ class TestDecorators(TestSlowapi): # meaning it should not raise any RateLimitExceeded error. if key_style == "url": assert second_call.status_code == 200 - # also assert that we counted only one request on the expected key - assert limiter._storage.get(expected_key) == 1 + 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(expected_key) == 2 + assert ( + limiter._storage.get( + "LIMITER/mock/tests.test_starlette_extension.t1_func/1/1/minute" + ) + == 2 + )