This commit is contained in:
Trevor Currie
2021-08-24 08:56:29 -07:00
parent 9bd22a1f00
commit 9f164179bc
3 changed files with 27 additions and 18 deletions

View File

@@ -378,7 +378,10 @@ class Limiter:
existing_retry_after_header = response.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)
reset_in = max(
self._determine_retry_time(existing_retry_after_header),
reset_in,
)
response.headers[self._header_mapping[HEADERS.RETRY_AFTER]] = (
formatdate(reset_in)
@@ -449,7 +452,9 @@ class Limiter:
def _determine_retry_time(self, retry_header_value) -> int:
try:
retry_after_date: Optional[datetime] = parsedate_to_datetime(retry_header_value)
retry_after_date: Optional[datetime] = parsedate_to_datetime(
retry_header_value
)
except TypeError:
retry_after_date = None
@@ -459,14 +464,16 @@ class Limiter:
try:
retry_after_int: int = int(retry_header_value)
except TypeError:
raise ValueError("Retry-After Header does not meet RFC2616 - value is not of http-date or int type.")
raise ValueError(
"Retry-After Header does not meet RFC2616 - value is not of http-date or int type."
)
return int(time.time() + retry_after_int)
def _check_request_limit(
self,
request: Request,
endpoint_func: Callable[..., Any],
self,
request: Request,
endpoint_func: Callable[..., Any],
in_middleware: bool = True,
) -> None:
"""

View File

@@ -90,8 +90,8 @@ class TestDecorators(TestSlowapi):
assert cli.get("/t1").status_code == 429
assert (
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429
)
def test_multiple_decorators_not_response(self):
@@ -115,8 +115,8 @@ class TestDecorators(TestSlowapi):
assert cli.get("/t1").status_code == 429
assert (
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429
)
def test_multiple_decorators_not_response_with_headers(self):
@@ -140,14 +140,15 @@ class TestDecorators(TestSlowapi):
assert cli.get("/t1").status_code == 429
assert (
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429
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)
with pytest.raises(Exception) as exc_info:
@app.get("/t3")
@limiter.limit("5/minute")
async def t3():
@@ -161,6 +162,7 @@ class TestDecorators(TestSlowapi):
app, limiter = self.build_fastapi_app(key_func=get_ipaddr)
with pytest.raises(Exception) as exc_info:
@app.get("/t3_sync")
@limiter.limit("5/minute")
def t3():

View File

@@ -87,8 +87,8 @@ class TestDecorators(TestSlowapi):
assert cli.get("/t1").status_code == 429
assert (
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429
)
def test_multiple_decorators_with_headers(self):
@@ -114,8 +114,8 @@ class TestDecorators(TestSlowapi):
assert cli.get("/t1").status_code == 429
assert (
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429
)
def test_headers_no_breach(self):
@@ -245,6 +245,6 @@ class TestDecorators(TestSlowapi):
assert cli.get("/t1").status_code == 429
assert (
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429
)