Implementing a Payload Validation after the Response Body is generated with a token

[trashdid]

I'm relatively new to using Litestar so forgive me if I am missing something. I am working on a project where I need to validate the response data against user roles after the response has been made.

Essentially the flow is Request with Access token -> Grab data from a database -> Build Response body -> Validate the roles from the Access Token match the required roles from the data -> Return Response/Raise 403.

I know that this isn't normally how one would approach validating a user's ability to access data, but due to the legacy systems that is being used, this was the approach my team and I determined was best for our use case.

Based on what I understood about Litestar, I had initially tried to use the after_request hook to grab the response before it was returned, and validate it there. This led to any issue because the ASGI connection was not passed with the response body in the after_request.

I then tried to use a contextvar that would put the token into a custom class. I used a guard on my controller that would set the contextvar and tried to get it in the after_request function. I could not figure out why, but the contextvar would save with the guard, but it would not be returned when I called it from the after_request.

The approach I am looking at now is making a custom Middleware protocol that will grab the token from the app and perform the role verification there before sending the response. I think this approach will work, but I am not sure if this is the proper way to implement the validation I want to do.

If anyone has any insights into what I could have been doing wrong in my various approaches or if there is a better one, I would greatly appreciate some advice. I am still getting a handle on using litestar, but I really like the framework and if there was someway to implement this with the after_request, that would be extremely helpful as easy for our team to maintain.

Thanks😄

[provinzkraut]

You won't have much luck with a middleware, since the response body has already been serialised.

There are several ways you could do this. One would be to implement a custom response class:

class ValidatingResponse[T](Response[T]):
    async def to_asgi_response[UserT, AuthT](self, request: Request[UserT, AuthT, Any]) -> ASGIResponse:
        content = self.content  # the unserialised response content
        user = request.user  # the user set by your auth middleware
        auth = request.auth  # the token instance created by your auth middleware


@get("/", response_class=ValidatingResponse)
async def get_things() -> Things:
    things = await fetch_things_from_db()
    return things

[trashdid]

Thanks for the the swift response @provinzkraut. If Im understanding this right, I would wrap my response in the custom response class and have the validation occur within that custom class. For example:

class ValidatingResponse[T](Response[T]):
    async def to_asgi_response[UserT, AuthT](self, request: Request[UserT, AuthT, Any]) -> ASGIResponse:
        content = self.content  # the unserialised response content
        auth = request.auth  # the token instance created by your auth middleware

        if auth.role != content.required_role:
               raise ForbiddenException()


@get("/", response_class=ValidatingResponse)
async def get_things() -> Things:
    things = await fetch_things_from_db()
    return things

In this case, if the user role does not match the required role for the content, the exception would be sent back instead of the response payload?

[provinzkraut]

Correct, that's the idea!

[trashdid]

Awesome,

Thanks again for your help.