python/slowapi
1
0
Fork 0
mirror of https://github.com/laurentS/slowapi.git synced 2026-03-13 09:10:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'master' into publish-from-ci

Browse Source
This commit is contained in:
Laurent Savaete
2023-03-30 11:29:29 +02:00
committed by GitHub
parent b5c1468264 f2bf93d890
commit be74f2a403
17 changed files with 1693 additions and 857 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
.github/ISSUE_TEMPLATE/bug-report---something-is-not-working.md vendored Normal file
View File

@@ -0,0 +1,28 @@
---
name: Bug report - Something is not working
about: You found a bug, or your code is not working as you expect
title: ''
labels: ''
assignees: ''
---
**Describe the bug**
A clear and concise description of what the bug is.
**To Reproduce**
Show us your code, only copy the relevant parts, the shorter it is, the easier it is to help you.
**Expected behavior**
A clear and concise description of what you expected to happen.
**Screenshots**
If applicable, add screenshots to help explain your problem.
**Your app (please complete the following information):**
- fastapi or starlette?
- Version?
- slowapi version (have you tried with the latest version)?
**Additional context**
Add any other context about the problem here.

8
.github/workflows/python-package.yml vendored
View File

@@ -15,12 +15,12 @@ jobs:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ['3.7', '3.8', '3.9', '3.10']
python-version: ['3.7', '3.8', '3.9', '3.10', '3.11']
steps:
- uses: actions/checkout@master
- uses: actions/checkout@v3
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v3
uses: actions/setup-python@v4
with:
python-version: ${{ matrix.python-version }}
- name: Install build dependencies for requests in python 3.9
@@ -30,7 +30,7 @@ jobs:
uses: snok/install-poetry@v1
with:
# Version of Poetry to use
version: 1.1.15
version: 1.2.2
- name: Install dependencies
run: |
poetry install

8
CHANGELOG.md
View File

@@ -1,5 +1,13 @@
# Change Log
## [0.1.7] - 2022-11-09
### Added
- Added ASGI middleware alternative (thanks @thentgesMindee)
- Added support for custom cost per hit (thanks @nootr)
- Added `key_style` parameter to choose between endpoint or url (thanks @thentgesMindee)
## [0.1.6] - 2022-08-20
### Added

3
README.md
View File

@@ -2,7 +2,8 @@
A rate limiting library for Starlette and FastAPI adapted from [flask-limiter](http://github.com/alisaifee/flask-limiter).
Note: this is alpha quality code still, the API may change, and things may fall apart while you try it.
This package is used in various production setups, handling millions of requests per month, and seems to behave as expected.
There might be some API changes when changing the code to be fully `async`, but we will notify users via appropriate `semver` version changes.
The documentation is on [read the docs](https://slowapi.readthedocs.io/en/latest/).

66
docs/examples.md
View File

@@ -65,3 +65,69 @@ limiter = Limiter(key_func=get_remote_address, storage_uri="redis://<host>:<port
where the /n in the redis url is the database number. To use the default one, just drop the /n from the url.
There are more examples in the [limits docs](https://limits.readthedocs.io/en/stable/storage.html) which is the library slowapi uses to manage storage.
## Set a custom cost per hit
Setting a custom cost per hit is useful to throttle requests based on something else than the request count.
Define a function which takes a request as parameter and returns a cost and pass it to the `limit` decorator:
```python
def get_hit_cost(request: Request) -> int:
return len(request)
@app.route("/someroute")
@limiter.limit("100/minute", cost=get_hit_cost)
def t(request: Request):
return PlainTextResponse("I'm limited by the request size")
```
## WSGI vs ASGI Middleware
`SlowAPIMiddleware` inheriting from Starlette's BaseHTTPMiddleware, you can find an alternative ASGI Middleware `SlowAPIASGIMiddleware`.
A few reasons to choose the ASGI middleware over the HTTP one are:
- Starlette [is probably going to deprecate BaseHTTPMiddleware](https://github.com/encode/starlette/issues/1678)
- ASGI middlewares [are more performant than WSGI ones](https://github.com/tiangolo/fastapi/issues/2241)
- built-in support for asynchronous exception handlers
- ...
Both middlewares are added to your application the same way:
```python
app = Starlette() # or FastAPI()
app.add_middleware(SlowAPIMiddleware)
```
or
```python
app = Starlette() # or FastAPI()
app.add_middleware(SlowAPIASGIMiddleware)
```
## 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.

5
docs/index.md
View File

@@ -72,10 +72,11 @@ Most feature are coming from (will come from) FlaskLimiter and the underlying [l
Supported now:
- Single and multiple `limit` decorator on endpoint functions to apply limits
- redis, memcached and memory backends to track your limits (memory as a fallback)
- support for sync and async HTTP endpoints
- Redis, memcached and memory backends to track your limits (memory as a fallback)
- Support for sync and async HTTP endpoints
- Support for shared limits across a set of routes
- Support for default global limit
- Support for a custom cost per hit
# Limitations and known issues

10
docs/requirements.txt
View File

@@ -1,14 +1,14 @@
appdirs==1.4.3
atomicwrites==1.3.0; sys_platform == "win32"
attrs==19.3.0
black==19.10b0
certifi==2019.11.28
black==23.3.0
certifi==2022.12.7
chardet==3.0.4
click==7.1.1
colorama==0.4.3; sys_platform == "win32"
dataclasses==0.6; python_version < "3.7"
fastapi==0.65.2
future==0.18.2
future==0.18.3
hiro==0.5.1
idna==2.9
importlib-metadata==1.5.0; python_version < "3.8"
@@ -36,11 +36,11 @@ pydantic==1.6.2
pyparsing==2.4.6
pytest==5.3.5
pyyaml==5.4
redis==3.4.1
redis==4.3.6
regex==2020.2.20
requests==2.23.0
six==1.14.0
starlette==0.13.2
starlette==0.25.0
toml==0.10.0
tornado==6.0.4
tqdm==4.50.0; python_version > "2.7"

1671
poetry.lock generated
View File

File diff suppressed because it is too large Load Diff

15
pyproject.toml
View File

@@ -1,6 +1,6 @@
[tool.poetry]
name = "slowapi"
version = "0.1.6"
version = "0.1.7"
description = "A rate limiting extension for Starlette and Fastapi"
authors = ["Laurent Savaete <laurent@where.tf>"]
license = "MIT"
@@ -15,15 +15,16 @@ include = ["slowapi/py.typed"]
[tool.poetry.dependencies]
python = ">=3.7,<4.0"
limits = "^1.5"
limits = "^2.3"
redis = {version = "^3.4.1", optional = true}
[tool.poetry.dev-dependencies]
isort = "^4.3.21"
mypy = "^0.910"
black = "^22.3.0"
fastapi = "^0.61.0"
black = "^23.0.0"
fastapi = "^0.89.0"
lxml = "^4.9.1"
starlette = "^0.13.6"
starlette = "^0.22.0"
mock = "^4.0.1"
hiro = "^0.5.1"
requests = "^2.22.0"
@@ -33,6 +34,8 @@ mkautodoc = "^0.1.0"
types-redis = "^3.5.6"
coverage = "^6.3"
flake8 = "^4.0.1"
setuptools = "^65.5.0"
httpx = "^0.23.3"
[tool.black]
line-length = 88
@@ -50,4 +53,4 @@ requires = ["poetry-core>=1.0.0"]
build-backend = "poetry.core.masonry.api"
[tool.poetry.extras]
redis = ["redis^3.4.1"]
redis = ["redis"]

6
renovate.json Normal file
View File

@@ -0,0 +1,6 @@
{
"$schema": "https://docs.renovatebot.com/renovate-schema.json",
"extends": [
"config:base"
]
}

145
slowapi/extension.py
View File

@@ -25,11 +25,12 @@ from typing import (
from limits import RateLimitItem # type: ignore
from limits.errors import ConfigurationError # type: ignore
from limits.storage import MemoryStorage, storage_from_string # type: ignore
from limits.storage import Storage # type: ignore
from limits.strategies import STRATEGIES, RateLimiter # type: ignore
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
@@ -121,6 +122,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__(
@@ -141,6 +143,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
@@ -175,12 +178,13 @@ 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(
[
LimitGroup(
limit, self._key_func, None, False, None, None, None, False
limit, self._key_func, None, False, None, None, None, 1, False
)
]
)
@@ -188,7 +192,15 @@ class Limiter:
self._application_limits.extend(
[
LimitGroup(
limit, self._key_func, "global", False, None, None, None, False
limit,
self._key_func,
"global",
False,
None,
None,
None,
1,
False,
)
]
)
@@ -196,7 +208,7 @@ class Limiter:
self._in_memory_fallback.extend(
[
LimitGroup(
limit, self._key_func, None, False, None, None, None, False
limit, self._key_func, None, False, None, None, None, 1, False
)
]
)
@@ -223,7 +235,7 @@ class Limiter:
C.HEADERS_ENABLED, False
)
self._storage_options.update(self.get_app_config(C.STORAGE_OPTIONS, {}))
self._storage: Storage = storage_from_string(
self._storage = storage_from_string(
self._storage_uri or self.get_app_config(C.STORAGE_URL, "memory://"),
**self._storage_options,
)
@@ -261,7 +273,15 @@ class Limiter:
if not self._application_limits and app_limits:
self._application_limits = [
LimitGroup(
app_limits, self._key_func, "global", False, None, None, None, False
app_limits,
self._key_func,
"global",
False,
None,
None,
None,
1,
False,
)
]
@@ -271,7 +291,7 @@ class Limiter:
if not self._default_limits and conf_limits:
self._default_limits = [
LimitGroup(
conf_limits, self._key_func, None, False, None, None, None, False
conf_limits, self._key_func, None, False, None, None, None, 1, False
)
]
fallback_enabled = self.get_app_config(C.IN_MEMORY_FALLBACK_ENABLED, False)
@@ -288,6 +308,7 @@ class Limiter:
None,
None,
None,
1,
False,
)
]
@@ -311,7 +332,11 @@ class Limiter:
"""
Place holder until we find a better way to load config from app
"""
return self.app_config(key, default=default_value, cast=type(default_value))
return (
self.app_config(key, default=default_value, cast=type(default_value))
if default_value
else self.app_config(key, default=default_value)
)
def __should_check_backend(self) -> bool:
if self.__check_backend_count > MAX_BACKEND_CHECKS:
@@ -338,6 +363,9 @@ class Limiter:
The backend that keeps track of consumption of endpoints vs limits
"""
if self._storage_dead and self._in_memory_fallback_enabled:
assert (
self._fallback_limiter
), "Fallback limiter is needed when in memory fallback is enabled"
return self._fallback_limiter
else:
return self._limiter
@@ -395,6 +423,58 @@ class Limiter:
raise
return response
def _inject_asgi_headers(
self, headers: MutableHeaders, current_limit: Tuple[RateLimitItem, List[str]]
) -> MutableHeaders:
"""
Injects 'X-RateLimit-Reset', 'X-RateLimit-Remaining', 'X-RateLimit-Limit'
and 'Retry-After' headers into :headers parameter if needed.
Basically the same as _inject_headers, but without access to the Response object.
-> supports ASGI Middlewares.
"""
if self.enabled and self._headers_enabled and current_limit is not None:
try:
window_stats: Tuple[int, int] = self.limiter.get_window_stats(
current_limit[0], *current_limit[1]
)
reset_in = 1 + window_stats[0]
headers[self._header_mapping[HEADERS.LIMIT]] = str(
current_limit[0].amount
)
headers[self._header_mapping[HEADERS.REMAINING]] = str(window_stats[1])
headers[self._header_mapping[HEADERS.RESET]] = str(reset_in)
# response may have an existing retry after
existing_retry_after_header = headers.get("Retry-After")
if existing_retry_after_header is not None:
reset_in = max(
self._determine_retry_time(existing_retry_after_header),
reset_in,
)
headers[self._header_mapping[HEADERS.RETRY_AFTER]] = (
formatdate(reset_in)
if self._retry_after == "http-date"
else str(int(reset_in - time.time()))
)
except Exception:
if self._in_memory_fallback and not self._storage_dead:
self.logger.warning(
"Rate limit storage unreachable - falling back to"
" in-memory storage"
)
self._storage_dead = True
headers = self._inject_asgi_headers(headers, current_limit)
if self._swallow_errors:
self.logger.exception(
"Failed to update rate limit headers. Swallowing error"
)
else:
raise
return headers
def __evaluate_limits(
self, request: Request, endpoint: str, limits: List[Limit]
) -> None:
@@ -420,7 +500,9 @@ class Limiter:
args = [self._key_prefix] + args
if not limit_for_header or lim.limit < limit_for_header[0]:
limit_for_header = (lim.limit, args)
if not self.limiter.hit(lim.limit, *args):
cost = lim.cost(request) if callable(lim.cost) else lim.cost
if not self.limiter.hit(lim.limit, *args, cost=cost):
self.logger.warning(
"ratelimit %s (%s) exceeded at endpoint: %s",
lim.limit,
@@ -464,24 +546,26 @@ class Limiter:
def _check_request_limit(
self,
request: Request,
endpoint_func: Callable[..., Any],
endpoint_func: Optional[Callable[..., Any]],
in_middleware: bool = True,
) -> None:
"""
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
@@ -489,23 +573,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():
@@ -527,12 +615,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
@@ -559,9 +650,9 @@ class Limiter:
methods: Optional[List[str]] = None,
error_message: Optional[str] = None,
exempt_when: Optional[Callable[..., bool]] = None,
cost: Union[int, Callable[..., int]] = 1,
override_defaults: bool = True,
) -> Callable[..., Any]:
_scope = scope if shared else None
def decorator(func: Callable[..., Response]):
@@ -578,6 +669,7 @@ class Limiter:
methods,
error_message,
exempt_when,
cost,
override_defaults,
)
else:
@@ -591,6 +683,7 @@ class Limiter:
methods,
error_message,
exempt_when,
cost,
override_defaults,
)
)
@@ -691,6 +784,7 @@ class Limiter:
methods: Optional[List[str]] = None,
error_message: Optional[str] = None,
exempt_when: Optional[Callable[..., bool]] = None,
cost: Union[int, Callable[..., int]] = 1,
override_defaults: bool = True,
) -> Callable:
"""
@@ -708,6 +802,7 @@ class Limiter:
error message used in the response.
* **exempt_when**: function returning a boolean indicating whether to exempt
the route from the limit
* **cost**: integer (or callable that returns one) which is the cost of a hit
* **override_defaults**: whether to override the default limits (default: True)
"""
return self.__limit_decorator(
@@ -717,6 +812,7 @@ class Limiter:
methods=methods,
error_message=error_message,
exempt_when=exempt_when,
cost=cost,
override_defaults=override_defaults,
)
@@ -727,6 +823,7 @@ class Limiter:
key_func: Optional[Callable[..., str]] = None,
error_message: Optional[str] = None,
exempt_when: Optional[Callable[..., bool]] = None,
cost: Union[int, Callable[..., int]] = 1,
override_defaults: bool = True,
) -> Callable:
"""
@@ -746,6 +843,7 @@ class Limiter:
error message used in the response.
* **exempt_when**: function returning a boolean indicating whether to exempt
the route from the limit
* **cost**: integer (or callable that returns one) which is the cost of a hit
* **override_defaults**: whether to override the default limits (default: True)
"""
return self.__limit_decorator(
@@ -755,6 +853,7 @@ class Limiter:
scope,
error_message=error_message,
exempt_when=exempt_when,
cost=cost,
override_defaults=override_defaults,
)

214
slowapi/middleware.py
View File

@@ -1,56 +1,206 @@
import inspect
from typing import Callable, Iterable, Optional, Tuple
from starlette.applications import Starlette
from starlette.datastructures import MutableHeaders
from starlette.middleware.base import (
BaseHTTPMiddleware,
RequestResponseEndpoint,
)
from starlette.requests import Request
from starlette.responses import Response
from starlette.routing import Match
from starlette.routing import BaseRoute, Match
from starlette.types import ASGIApp, Message, Scope, Receive, Send
from slowapi import Limiter, _rate_limit_exceeded_handler
def _find_route_handler(
routes: Iterable[BaseRoute], scope: Scope
) -> Optional[Callable]:
handler = None
for route in routes:
match, _ = route.matches(scope)
if match == Match.FULL and hasattr(route, "endpoint"):
handler = route.endpoint # type: ignore
return handler
def _get_route_name(handler: Callable):
return f"{handler.__module__}.{handler.__name__}"
def _check_limits(
limiter: Limiter, request: Request, handler: Optional[Callable], app: Starlette
) -> Tuple[Optional[Callable], bool, Optional[Exception]]:
"""
Utils to check (if needed) current requests limit.
It returns a tuple of size 3:
1. The exception handler to run, if needed
2. a bool, True if we need to inject some headers, False otherwise
3. the exception that happened, if any
"""
if limiter._auto_check and not getattr(
request.state, "_rate_limiting_complete", False
):
try:
limiter._check_request_limit(request, handler, True)
except Exception as e:
# handle the exception since the global exception handler won't pick it up if we call_next
exception_handler = app.exception_handlers.get(
type(e), _rate_limit_exceeded_handler
)
return exception_handler, False, e
return None, True, None
return None, False, None
def sync_check_limits(
limiter: Limiter, request: Request, handler: Optional[Callable], app: Starlette
) -> Tuple[Optional[Response], bool]:
"""
Returns a `Response` object if an error occurred, as well as a boolean to know
whether we should inject headers or not.
Used in our WSGI middleware, it only supports synchronous exception_handler.
This will fallback on _rate_limit_exceeded_handler otherwise.
"""
exception_handler, _bool, exc = _check_limits(limiter, request, handler, app)
if not exception_handler or not exc:
return None, _bool
# cannot execute asynchronous code in a synchronous middleware,
# -> fallback on default exception handler
if inspect.iscoroutinefunction(exception_handler):
exception_handler = _rate_limit_exceeded_handler
return exception_handler(request, exc), _bool # type: ignore
async def async_check_limits(
limiter: Limiter, request: Request, handler: Optional[Callable], app: Starlette
) -> Tuple[Optional[Response], bool]:
"""
Returns a `Response` object if an error occurred, as well as a boolean to know
whether we should inject headers or not.
Used in our ASGI middleware, this support both synchronous or asynchronous exception handlers.
"""
exception_handler, _bool, exc = _check_limits(limiter, request, handler, app)
if not exception_handler:
return None, _bool
if inspect.iscoroutinefunction(exception_handler):
return await exception_handler(request, exc), _bool
else:
return exception_handler(request, exc), _bool
def _should_exempt(limiter: Limiter, handler: Optional[Callable]) -> bool:
# if we can't find the route handler
if handler is None:
return True
name = _get_route_name(handler)
# if exempt no need to check
if name in limiter._exempt_routes:
return True
# there is a decorator for this route we let the decorator handle it
if name in limiter._route_limits:
return True
return False
class SlowAPIMiddleware(BaseHTTPMiddleware):
async def dispatch(
self, request: Request, call_next: RequestResponseEndpoint
) -> Response:
app: Starlette = request.app
limiter: Limiter = app.state.limiter
handler = None
if not limiter.enabled:
return await call_next(request)
for route in app.routes:
match, _ = route.matches(request.scope)
if match == Match.FULL and hasattr(route, "endpoint"):
handler = route.endpoint # type: ignore
# if we can't find the route handler
if handler is None:
handler = _find_route_handler(app.routes, request.scope)
if _should_exempt(limiter, handler):
return await call_next(request)
name = "%s.%s" % (handler.__module__, handler.__name__)
# if exempt no need to check
if name in limiter._exempt_routes:
return await call_next(request)
error_response, should_inject_headers = sync_check_limits(
limiter, request, handler, app
)
if error_response is not None:
return error_response
# there is a decorator for this route we let the decorator handle it
if name in limiter._route_limits:
return await call_next(request)
# let the decorator handle if already in
if limiter._auto_check and not getattr(
request.state, "_rate_limiting_complete", False
):
try:
limiter._check_request_limit(request, handler, True)
except Exception as e:
# handle the exception since the global exception handler won't pick it up if we call_next
exception_handler = app.exception_handlers.get(
type(e), _rate_limit_exceeded_handler
)
return exception_handler(request, e)
# request.state._rate_limiting_complete = True
response = await call_next(request)
response = await call_next(request)
if should_inject_headers:
response = limiter._inject_headers(response, request.state.view_rate_limit)
return response
return await call_next(request)
return response
class SlowAPIASGIMiddleware:
def __init__(self, app: ASGIApp) -> None:
self.app = app
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] != "http":
return await self.app(scope, receive, send)
await _ASGIMiddlewareResponder(self.app)(scope, receive, send)
class _ASGIMiddlewareResponder:
def __init__(self, app: ASGIApp) -> None:
self.app = app
self.error_response: Optional[Response] = None
self.initial_message: Message = {}
self.inject_headers = False
async def send_wrapper(self, message: Message) -> None:
if message["type"] == "http.response.start":
# do not send the http.response.start message now, so that we can edit the headers
# before sending it, based on what happens in the http.response.body message.
self.initial_message = message
elif message["type"] == "http.response.body":
if self.error_response:
self.initial_message["status"] = self.error_response.status_code
if self.inject_headers:
headers = MutableHeaders(raw=self.initial_message["headers"])
headers = self.limiter._inject_asgi_headers(
headers, self.request.state.view_rate_limit
)
# send the http.response.start message just before the http.response.body one,
# now that the headers are updated
await self.send(self.initial_message)
await self.send(message)
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
self.send = send
_app: Starlette = scope["app"]
limiter: Limiter = _app.state.limiter
if not limiter.enabled:
return await self.app(scope, receive, self.send)
handler = _find_route_handler(_app.routes, scope)
request = Request(scope, receive=receive, send=self.send)
if _should_exempt(limiter, handler):
return await self.app(scope, receive, self.send)
error_response, should_inject_headers = await async_check_limits(
limiter, request, handler, _app
)
if error_response is not None:
return await error_response(scope, receive, self.send_wrapper)
if should_inject_headers:
self.inject_headers = True
self.limiter = limiter
self.request = request
return await self.app(scope, receive, self.send_wrapper)

10
slowapi/util.py
View File

@@ -11,11 +11,17 @@ def get_ipaddr(request: Request) -> str:
if "X_FORWARDED_FOR" in request.headers:
return request.headers["X_FORWARDED_FOR"]
else:
return request.client.host or "127.0.0.1"
if not request.client or not request.client.host:
return "127.0.0.1"
return request.client.host
def get_remote_address(request: Request) -> str:
"""
Returns the ip address for the current request (or 127.0.0.1 if none found)
"""
return request.client.host or "127.0.0.1"
if not request.client or not request.client.host:
return "127.0.0.1"
return request.client.host

5
slowapi/wrappers.py
View File

@@ -18,6 +18,7 @@ class Limit(object):
methods: Optional[List[str]],
error_message: Optional[Union[str, Callable[..., str]]],
exempt_when: Optional[Callable[..., bool]],
cost: Union[int, Callable[..., int]],
override_defaults: bool,
) -> None:
self.limit = limit
@@ -27,6 +28,7 @@ class Limit(object):
self.methods = methods
self.error_message = error_message
self.exempt_when = exempt_when
self.cost = cost
self.override_defaults = override_defaults
@property
@@ -65,6 +67,7 @@ class LimitGroup(object):
methods: Optional[List[str]],
error_message: Optional[Union[str, Callable[..., str]]],
exempt_when: Optional[Callable[..., bool]],
cost: Union[int, Callable[..., int]],
override_defaults: bool,
):
self.__limit_provider = limit_provider
@@ -74,6 +77,7 @@ class LimitGroup(object):
self.methods = methods and [m.lower() for m in methods] or methods
self.error_message = error_message
self.exempt_when = exempt_when
self.cost = cost
self.override_defaults = override_defaults
self.request = None
@@ -100,6 +104,7 @@ class LimitGroup(object):
self.methods,
self.error_message,
self.exempt_when,
self.cost,
self.override_defaults,
)

76
tests/__init__.py
View File

@@ -1,37 +1,69 @@
import asyncio
import logging
import pytest
from fastapi import FastAPI
from mock import mock # type: ignore
from starlette.applications import Starlette
from starlette.requests import Request
from slowapi.errors import RateLimitExceeded
from slowapi.extension import Limiter, _rate_limit_exceeded_handler
from slowapi.middleware import SlowAPIMiddleware
from slowapi.middleware import SlowAPIMiddleware, SlowAPIASGIMiddleware
from slowapi.util import get_remote_address
async def _async_rate_limit_exceeded_handler(request: Request, exc: RateLimitExceeded):
await asyncio.sleep(0)
return _rate_limit_exceeded_handler(request, exc)
class TestSlowapi:
def build_starlette_app(self, config={}, **limiter_args):
limiter_args.setdefault("key_func", get_remote_address)
limiter = Limiter(**limiter_args)
app = Starlette(debug=True)
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
app.add_middleware(SlowAPIMiddleware)
@pytest.fixture(
params=[
(SlowAPIMiddleware, _rate_limit_exceeded_handler),
(SlowAPIASGIMiddleware, _rate_limit_exceeded_handler),
(SlowAPIASGIMiddleware, _async_rate_limit_exceeded_handler),
]
)
def build_starlette_app(self, request):
def _factory(config={}, **limiter_args):
middleware, exception_handler = request.param
mock_handler = mock.Mock()
mock_handler.level = logging.INFO
limiter.logger.addHandler(mock_handler)
return app, limiter
limiter_args.setdefault("key_func", get_remote_address)
limiter = Limiter(**limiter_args)
app = Starlette(debug=True)
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, exception_handler)
app.add_middleware(middleware)
def build_fastapi_app(self, config={}, **limiter_args):
limiter_args.setdefault("key_func", get_remote_address)
limiter = Limiter(**limiter_args)
app = FastAPI()
app.state.limiter = limiter
app.add_middleware(SlowAPIMiddleware)
mock_handler = mock.Mock()
mock_handler.level = logging.INFO
limiter.logger.addHandler(mock_handler)
return app, limiter
mock_handler = mock.Mock()
mock_handler.level = logging.INFO
limiter.logger.addHandler(mock_handler)
return app, limiter
return _factory
@pytest.fixture(
params=[
(SlowAPIMiddleware, _rate_limit_exceeded_handler),
(SlowAPIASGIMiddleware, _rate_limit_exceeded_handler),
(SlowAPIASGIMiddleware, _async_rate_limit_exceeded_handler),
]
)
def build_fastapi_app(self, request):
def _factory(config={}, **limiter_args):
middleware, exception_handler = request.param
limiter_args.setdefault("key_func", get_remote_address)
limiter = Limiter(**limiter_args)
app = FastAPI()
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, exception_handler)
app.add_middleware(middleware)
mock_handler = mock.Mock()
mock_handler.level = logging.INFO
limiter.logger.addHandler(mock_handler)
return app, limiter
return _factory

138
tests/test_fastapi_extension.py
View File

@@ -9,8 +9,8 @@ from tests import TestSlowapi
class TestDecorators(TestSlowapi):
def test_single_decorator(self):
app, limiter = self.build_fastapi_app(key_func=get_ipaddr)
def test_single_decorator(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr)
@app.get("/t1")
@limiter.limit("5/minute")
@@ -22,8 +22,8 @@ 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)
def test_single_decorator_with_headers(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr, headers_enabled=True)
@app.get("/t1")
@limiter.limit("5/minute")
@@ -39,8 +39,8 @@ class TestDecorators(TestSlowapi):
)
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)
def test_single_decorator_not_response(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr)
@app.get("/t1")
@limiter.limit("5/minute")
@@ -52,8 +52,8 @@ class TestDecorators(TestSlowapi):
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)
def test_single_decorator_not_response_with_headers(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr, headers_enabled=True)
@app.get("/t1")
@limiter.limit("5/minute")
@@ -69,8 +69,8 @@ class TestDecorators(TestSlowapi):
)
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)
def test_multiple_decorators(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr)
@app.get("/t1")
@limiter.limit(
@@ -94,8 +94,8 @@ class TestDecorators(TestSlowapi):
== 429
)
def test_multiple_decorators_not_response(self):
app, limiter = self.build_fastapi_app(key_func=get_ipaddr)
def test_multiple_decorators_not_response(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr)
@app.get("/t1")
@limiter.limit(
@@ -119,8 +119,8 @@ class TestDecorators(TestSlowapi):
== 429
)
def test_multiple_decorators_not_response_with_headers(self):
app, limiter = self.build_fastapi_app(key_func=get_ipaddr, headers_enabled=True)
def test_multiple_decorators_not_response_with_headers(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr, headers_enabled=True)
@app.get("/t1")
@limiter.limit(
@@ -144,8 +144,8 @@ class TestDecorators(TestSlowapi):
== 429
)
def test_endpoint_missing_request_param(self):
app, limiter = self.build_fastapi_app(key_func=get_ipaddr)
def test_endpoint_missing_request_param(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr)
with pytest.raises(Exception) as exc_info:
@@ -158,8 +158,8 @@ class TestDecorators(TestSlowapi):
r"""^No "request" or "websocket" argument on function .*"""
)
def test_endpoint_missing_request_param_sync(self):
app, limiter = self.build_fastapi_app(key_func=get_ipaddr)
def test_endpoint_missing_request_param_sync(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr)
with pytest.raises(Exception) as exc_info:
@@ -172,8 +172,8 @@ class TestDecorators(TestSlowapi):
r"""^No "request" or "websocket" argument on function .*"""
)
def test_endpoint_request_param_invalid(self):
app, limiter = self.build_fastapi_app(key_func=get_ipaddr)
def test_endpoint_request_param_invalid(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr)
@app.get("/t4")
@limiter.limit("5/minute")
@@ -187,8 +187,8 @@ 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)
def test_endpoint_response_param_invalid(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr, headers_enabled=True)
@app.get("/t4")
@limiter.limit("5/minute")
@@ -202,8 +202,8 @@ class TestDecorators(TestSlowapi):
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)
def test_endpoint_request_param_invalid_sync(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr)
@app.get("/t5")
@limiter.limit("5/minute")
@@ -217,8 +217,8 @@ class TestDecorators(TestSlowapi):
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)
def test_endpoint_response_param_invalid_sync(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr, headers_enabled=True)
@app.get("/t5")
@limiter.limit("5/minute")
@@ -232,7 +232,7 @@ class TestDecorators(TestSlowapi):
r"""parameter `response` must be an instance of starlette.responses.Response"""
)
def test_dynamic_limit_provider_depending_on_key(self):
def test_dynamic_limit_provider_depending_on_key(self, build_fastapi_app):
def custom_key_func(request: Request):
if request.headers.get("TOKEN") == "secret":
return "admin"
@@ -243,7 +243,7 @@ class TestDecorators(TestSlowapi):
return "10/minute"
return "5/minute"
app, limiter = self.build_fastapi_app(key_func=custom_key_func)
app, limiter = build_fastapi_app(key_func=custom_key_func)
@app.get("/t1")
@limiter.limit(dynamic_limit_provider)
@@ -259,11 +259,11 @@ class TestDecorators(TestSlowapi):
response = client.get("/t1", headers={"TOKEN": "secret"})
assert response.status_code == 200 if i < 10 else 429
def test_disabled_limiter(self):
def test_disabled_limiter(self, build_fastapi_app):
"""
Check that the limiter does nothing if disabled (both sync and async)
"""
app, limiter = self.build_fastapi_app(key_func=get_ipaddr, enabled=False)
app, limiter = build_fastapi_app(key_func=get_ipaddr, enabled=False)
@app.get("/t1")
@limiter.limit("5/minute")
@@ -291,3 +291,81 @@ class TestDecorators(TestSlowapi):
for i in range(0, 10):
response = client.get("/t3")
assert response.status_code == 200
def test_cost(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr)
@app.get("/t1")
@limiter.limit("50/minute", cost=10)
async def t1(request: Request):
return PlainTextResponse("test")
@app.get("/t2")
@limiter.limit("50/minute", cost=15)
async def t2(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
response = client.get("/t2")
assert response.status_code == 200 if i < 3 else 429
def test_callable_cost(self, build_fastapi_app):
app, limiter = build_fastapi_app(key_func=get_ipaddr)
@app.get("/t1")
@limiter.limit("50/minute", cost=lambda request: int(request.headers["foo"]))
async def t1(request: Request):
return PlainTextResponse("test")
@app.get("/t2")
@limiter.limit(
"50/minute", cost=lambda request: int(request.headers["foo"]) * 1.5
)
async def t2(request: Request):
return PlainTextResponse("test")
client = TestClient(app)
for i in range(0, 10):
response = client.get("/t1", headers={"foo": "10"})
assert response.status_code == 200 if i < 5 else 429
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
)

142
tests/test_starlette_extension.py
View File

@@ -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
@@ -10,8 +11,8 @@ from tests import TestSlowapi
class TestDecorators(TestSlowapi):
def test_single_decorator_async(self):
app, limiter = self.build_starlette_app(key_func=get_ipaddr)
def test_single_decorator_async(self, build_starlette_app):
app, limiter = build_starlette_app(key_func=get_ipaddr)
@limiter.limit("5/minute")
async def t1(request: Request):
@@ -26,8 +27,8 @@ class TestDecorators(TestSlowapi):
if i < 5:
assert response.text == "test"
def test_single_decorator_sync(self):
app, limiter = self.build_starlette_app(key_func=get_ipaddr)
def test_single_decorator_sync(self, build_starlette_app):
app, limiter = build_starlette_app(key_func=get_ipaddr)
@limiter.limit("5/minute")
def t1(request: Request):
@@ -42,8 +43,8 @@ class TestDecorators(TestSlowapi):
if i < 5:
assert response.text == "test"
def test_shared_decorator(self):
app, limiter = self.build_starlette_app(key_func=get_ipaddr)
def test_shared_decorator(self, build_starlette_app):
app, limiter = build_starlette_app(key_func=get_ipaddr)
shared_lim = limiter.shared_limit("5/minute", "somescope")
@@ -65,8 +66,8 @@ class TestDecorators(TestSlowapi):
# the shared limit has already been hit via t1
assert client.get("/t2").status_code == 429
def test_multiple_decorators(self):
app, limiter = self.build_starlette_app(key_func=get_ipaddr)
def test_multiple_decorators(self, build_starlette_app):
app, limiter = build_starlette_app(key_func=get_ipaddr)
@limiter.limit("10 per minute", lambda: "test")
@limiter.limit("5/minute") # per ip as per default key_func
@@ -89,10 +90,8 @@ class TestDecorators(TestSlowapi):
== 429
)
def test_multiple_decorators_with_headers(self):
app, limiter = self.build_starlette_app(
key_func=get_ipaddr, headers_enabled=True
)
def test_multiple_decorators_with_headers(self, build_starlette_app):
app, limiter = build_starlette_app(key_func=get_ipaddr, headers_enabled=True)
@limiter.limit("10 per minute", lambda: "test")
@limiter.limit("5/minute") # per ip as per default key_func
@@ -116,8 +115,8 @@ class TestDecorators(TestSlowapi):
== 429
)
def test_headers_no_breach(self):
app, limiter = self.build_starlette_app(
def test_headers_no_breach(self, build_starlette_app):
app, limiter = build_starlette_app(
headers_enabled=True, key_func=get_remote_address
)
@@ -149,8 +148,8 @@ class TestDecorators(TestSlowapi):
assert resp.headers.get("Retry-After") == str(1)
def test_headers_breach(self):
app, limiter = self.build_starlette_app(
def test_headers_breach(self, build_starlette_app):
app, limiter = build_starlette_app(
headers_enabled=True, key_func=get_remote_address
)
@@ -172,10 +171,10 @@ class TestDecorators(TestSlowapi):
)
assert resp.headers.get("Retry-After") == str(int(50))
def test_retry_after(self):
def test_retry_after(self, build_starlette_app):
# FIXME: this test is not actually running!
app, limiter = self.build_starlette_app(
app, limiter = build_starlette_app(
headers_enabled=True, key_func=get_remote_address
)
@@ -193,8 +192,8 @@ class TestDecorators(TestSlowapi):
resp = cli.get("/t1")
assert resp.status_code == 200
def test_exempt_decorator(self):
app, limiter = self.build_starlette_app(
def test_exempt_decorator(self, build_starlette_app):
app, limiter = build_starlette_app(
headers_enabled=True,
key_func=get_remote_address,
default_limits=["1/minute"],
@@ -235,8 +234,8 @@ class TestDecorators(TestSlowapi):
assert resp2.status_code == 200
# todo: more tests - see https://github.com/alisaifee/flask-limiter/blob/55df08f14143a7e918fc033067a494248ab6b0c5/tests/test_decorators.py#L187
def test_default_and_decorator_limit_merging(self):
app, limiter = self.build_starlette_app(
def test_default_and_decorator_limit_merging(self, build_starlette_app):
app, limiter = build_starlette_app(
key_func=lambda: "test", default_limits=["10/minute"]
)
@@ -259,3 +258,102 @@ class TestDecorators(TestSlowapi):
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429
)
def test_cost(self, build_starlette_app):
app, limiter = build_starlette_app(key_func=get_ipaddr)
@limiter.limit("50/minute", cost=10)
async def t1(request: Request):
return PlainTextResponse("test")
app.add_route("/t1", t1)
@limiter.limit("50/minute", cost=15)
async def t2(request: Request):
return PlainTextResponse("test")
app.add_route("/t2", t2)
client = TestClient(app)
for i in range(0, 10):
response = client.get("/t1")
assert response.status_code == 200 if i < 5 else 429
if i < 5:
assert response.text == "test"
else:
assert "error" in response.json()
response = client.get("/t2")
assert response.status_code == 200 if i < 3 else 429
if i < 3:
assert response.text == "test"
else:
assert "error" in response.json()
def test_callable_cost(self, build_starlette_app):
app, limiter = build_starlette_app(key_func=get_ipaddr)
@limiter.limit("50/minute", cost=lambda request: int(request.headers["foo"]))
async def t1(request: Request):
return PlainTextResponse("test")
app.add_route("/t1", t1)
@limiter.limit(
"50/minute", cost=lambda request: int(request.headers["foo"]) * 1.5
)
async def t2(request: Request):
return PlainTextResponse("test")
app.add_route("/t2", t2)
client = TestClient(app)
for i in range(0, 10):
response = client.get("/t1", headers={"foo": "10"})
assert response.status_code == 200 if i < 5 else 429
if i < 5:
assert response.text == "test"
else:
assert "error" in response.json()
response = client.get("/t2", headers={"foo": "5"})
assert response.status_code == 200 if i < 6 else 429
if i < 6:
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
)