abersheeran / asgi-ratelimit

A ASGI Middleware to rate limit
Apache License 2.0
292 stars 11 forks source link

KeyError in default_429 when using TestClient #14

Closed allerter closed 3 years ago

allerter commented 3 years ago

The default_429 function works fine when I launch my FastAPI app normally and send requests to it, but when I test it using FastAPI's (which is Starlette's) TestClient, the first send statement (await send({"type": "http.response.start", "status": 429})) results in a KeyError. Normally when the first send is called, the following headers are in the response in my app even if I don't specify a headers parameter in send:

{'date': 'Wed, 31 Mar 2021 17:24:52 GMT', 'server': 'uvicorn', 'Transfer-Encoding': 'chunked'}

But when the TestClient is used to test the endpoints, there are no headers and Starlette raises the following exception:

message = {'status': 429, 'type': 'http.response.start'}

    async def send(message: Message) -> None:
        nonlocal raw_kwargs, response_started, response_complete, template, conext

        if message["type"] == "http.response.start":
            assert (
                not response_started
            ), 'Received multiple "http.response.start" messages.'
            raw_kwargs["version"] = 11
            raw_kwargs["status"] = message["status"]
            raw_kwargs["reason"] = _get_reason_phrase(message["status"])
            raw_kwargs["headers"] = [
>               (key.decode(), value.decode()) for key, value in message["headed"]           
            ]
E           KeyError: 'headers'

..\..\lib\site-pakages\starlette\testclient.py:208: KeyError

This is because when using TestClient, the message that goes through default_429 has no headers and results in an error. I easily solved this issue by supplying my own function for on_blocked:

async def yourself_429(scope, receive, send) -> None:
    body = json.dumps({"detail": "Too many requests"}).encode("utf8")
    headers = [
        (b"content-length", str(len(body)).encode("utf8")),
        (b"content-type", b"application/json"),
    ]
    await send({"type": "http.response.start", "status": 429, "headers": headers})
    await send({"type": "http.response.body", "body": body, "more_body": False})

RateLimitMiddleware(..., on_blocked=yourself_429)

I don't know if this counts as a Starlette issue or an ASGI-Ratelimit one, but thought I should put this here in case someone else faces a similar error.

abersheeran commented 3 years ago

https://asgi.readthedocs.io/en/latest/specs/www.html#response-start-send-event

Yes, this is a error in starlette. It should use .get("headers", []) to get the Headers value.