abersheeran / asgi-ratelimit

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

EmptyInformation results in Internal Server Error #12

Closed allerter closed 3 years ago

allerter commented 3 years ago

I successfully set up the rate limiter for my app with JWT auth and RedisBacked, but when I test it with requests that miss the Authorization header, the error isn't picked up by FastAPI and instead it results in an Internal Server Error. I tried adding an exception handler for EmptyInformation to send a 401 response back to the user, but the exception won't get picked up by the exception handler and still results in an Internal Server Error. Am I missing something here? What should I do to handle the EmptyInformation exception?

@app.exception_handler(EmptyInformation)
async def unicorn_exception_handler(request: Request, exc: EmptyInformation):
    return JSONResponse(
        status_code=401,
        content={"detail": "Unauthorized access."},
    )

Code:

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from ratelimit import RateLimitMiddleware, Rule
from ratelimit.auths.jwt import create_jwt_auth
from ratelimit.auths import EmptyInformation
from ratelimit.backends.redis import RedisBackend

SECRET_KEY = "SOME_SECRET_KEY"
ALGORITHM = "HS256"

app = FastAPI()
app.add_middleware(
    RateLimitMiddleware,
    authenticate=create_jwt_auth(key=SECRET_KEY, algorithms=[ALGORITHM]),
    backend=RedisBackend(),
    config={
        r"^/$": [Rule(second=5, group="default"), Rule(group="unlimited")],
    },
)

# @app.exception_handler(EmptyInformation)
# async def emptyinformation_exception_handler(request: Request, exc: EmptyInformation):
#    return JSONResponse(
#        status_code=401,
#        content={"detail": "Unauthorized access."},
#    )

@app.get("/")
async def root():
    return {"message": "Hello World"}

Version info

To reproduce:

pip install asgi-ratelimiter[full]
uvicorn main:app
abersheeran commented 3 years ago

Unfortunately. fastAPI cannot catch exceptions thrown in user‘s ASGI middleware.

abersheeran commented 3 years ago

You can use code similar to the following to solve this problem.

from ratelimit import RateLimitMiddleware
from ratelimit.auths import EmptyInformation

class CustomRateLimitMiddleware(RateLimitMiddleware):
    async def __call__(self, scope, receive, send):
        try:
            await super().__call__(scope, receive, send)
        except EmptyInformation:
            response = JSONResponse(status_code=401, content={"detail": "Unauthorized access."})
            await response(scope, receive, send)
allerter commented 3 years ago

Since we're adding a middleware here, it makes sense that FastAPI wouldn't catch the exception, but nonetheless, this sounds like a good solution. Thanks! Why doesn't the package offer this by default like the on_blocked parameter? Something like this would be nice:


async def send_401(scope: Scope, receive: Receive, send: Send, exc: EmptyInformation) -> None:
    await send({"type": "http.response.start", "status": 401})
    await send({"type": "http.response.body", "body": b"Unauthorized", "more_body": False})

app.add_middleware(
    RateLimitMiddleware,
    authenticate=AUTH_FUNCTION,
    backend=RedisBackend(),
    on_error=send_401, # or on_empty, on_auth_error, ...
    config={
        r"^/towns": [Rule(second=1, group="default"), Rule(group="admin")],
        r"^/forests": [Rule(minute=1, group="default"), Rule(group="admin")],
    },
)
abersheeran commented 3 years ago

This is a good idea. If you are interested, you can create a PR about it.

allerter commented 3 years ago

I'll send a PR soon.