Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 19 additions & 12 deletions slowapi/middleware.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,7 @@ def __init__(self, app: ASGIApp) -> None:
self.error_response: Optional[Response] = None
self.initial_message: Message = {}
self.inject_headers = False
self.started = False

async def send_wrapper(self, message: Message) -> None:
if message["type"] == "http.response.start":
Expand All @@ -164,18 +165,24 @@ async def send_wrapper(self, message: Message) -> None:
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)
if not self.started:
# Only modify and send the initial message once, on the first body chunk.
# Streaming responses (e.g. FileResponse) send multiple body chunks;
# re-sending http.response.start on each would cause
# "RuntimeError: ASGI flow error: Response already started".
self.started = True

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
)

await self.send(self.initial_message)

await self.send(message)

async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
Expand Down