FastAPI (SlowAPI Dynamic Rate Limit using keyfunc)

Question

Currently, I am trying to find a way to use a custom _keyfunc with SlowAPI that will allow me to access user data. The customer function I wanted to pass through looks like this:

async def get_limiter_id(current_user: User = Depends(get_current_user)):
    print(current_user)
    return str(current_user.retailer_id)

limiter = Limiter(
    key_func= get_limiter_id,
    default_limits=["200 per day", "50 per hour", "1 per minute"]
)

app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

The problem is that it doesn't accept asynchronous functions, so if I try to pass in: current_user: User = Depends(get_current_user) as an argument, it won't work because either you make the function non async, in which case the current user won't get passed through properly, or you make it async but then it will it just pass the coroutine object through instead because you didn't await the result (which isn't possible with how SlowAPI is written atm).

This is the get_current_user function for clarity sake:

async def get_current_user(token: str = Depends(oauth2_scheme)):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
        print(username)
        token_data = TokenData(username=username)
    except JWTError:
        raise credentials_exception
    user = get_user(fake_user_db, username=token_data.username)
    if user is None:
        raise credentials_exception
    if user.enabled == 0:
        raise HTTPException(status_code=400, detail="Inactive user")
    return user

One solution I thought of was to grab to put the retailer_id in the header so that I could have a custom function like this:

def get_retailer_id(request: Request):
    token = request.headers.get('retailer_id')
    return retailer_id

But I don't know how I can get this user specific information written into the header of all paths at creation of the bearer token. Any thoughts on how I could solve this?

Kevin Longe · Accepted Answer · 2022-01-10 00:30:57Z

0

Solved the issue. Didn't realise that I could just store this custom info in the token payload (upon creation of the token) and then bring it down from the header. There is a bit of code duplication here, but it's a good enough solution:

def get_retailer_id(request: Request):

    token = request.headers.get('Authorization')
    token = token.replace("Bearer ", "")
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Token is invalid or payload is corrupt",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        retailer_id: str = payload.get("retailer_id")
        return str(retailer_id)

    except Exception:
        raise credentials_exception

edited Jan 10, 2022 at 0:30

answered Jan 9, 2022 at 18:51

Kevin Longe

211 silver badge3 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Amine Bk · Accepted Answer · 2023-03-12 15:53:48Z

0

If there is a direct mapping between retailer_id to access token you can simple do ratelimit without DB interaction for faster handling:

import hashlib
from fastapi import HTTPException, status

def get_limiter_id(request: Request):
    try:
        token = request.headers['Authorization']
    except KeyError:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED)
    hashed_token = hashlib.md5(token.encode())
    hashed_token = hashed.hexdigest()
    return hashed_token

answered Mar 12, 2023 at 15:53

Amine Bk

1211 silver badge2 bronze badges

Collectives™ on Stack Overflow

FastAPI (SlowAPI Dynamic Rate Limit using keyfunc)

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related