How to gain access to headers / authenticate subscriptions with FastApi and graphql-transport-ws

[kristjanvalur]

So, working on a FastAPI integration, and using the graphql-transport-ws.
For subscriptions, info.context["request" just holds a starlette.websockets.WebSocket instance.

How can I get access to the headers that were provided when establishing the websocket connection?
Alternatively, can I get access to the payload member from the connection_init message?

Alternatively, is there a way to intercept and handle the websocket upgrade operation, and accept and/or reject it based on request headers? Optionally, attach some processed data (e.g. user permissions) to the resulting Websockets object?

Alternatively, is there a way to intercept and handler the connection_init message and either accept or reject the conneciton based on the request headers (from the original connection) or the optional payload and attach custom data to the resulting connection object?

Basically, I'm searching for the canonical way to do request authentication and authorization when working with requests that arrive via the graphql-transport-ws transport, be they subscriptions, mutations or queries.

Cheers!

[kristjanvalur]

Okay, actually opening stuff up in the debugger helped me reaslize that the starlette.websockets.WebSocket is a valid "request" and has headers member. So there is that.

Still, it would be very useful to be able to intercept the graphql connection request, perform some data processing, (authentication, for instance) and attach some data to the request as a result... Going to experiment a bit.

[kristjanvalur]

Okay, I've done some other experiments. It is possible to override the Router to intecept the connection_init message:

from strawberry.fastapi.handlers import GraphQLTransportWSHandler, GraphQLWSHandler
from strawberry.subscriptions.protocols.graphql_ws.types import (
    OperationMessage,
    OperationMessagePayload,
    StartPayload,
)
from strawberry.subscriptions.protocols.graphql_ws import (
    GQL_COMPLETE,
    GQL_CONNECTION_ACK,
    GQL_CONNECTION_INIT,
    GQL_CONNECTION_KEEP_ALIVE,
    GQL_CONNECTION_TERMINATE,
    GQL_DATA,
    GQL_ERROR,
    GQL_START,
    GQL_STOP,
)
class MyGraphQLWSHandler(GraphQLWSHandler):
    async def handle_connection_init(self, message: OperationMessage) -> None:
        if message['payload'].get("auth", 1) != 0:
            await self.close(4400, "this is the end")
            return
        payload = {"status":"ok"}
        self.success = True
        await self.finish_connection_init(payload)
        
    async def finish_connection_init(self, payload=None) -> None:
        data: OperationMessage = {"type": GQL_CONNECTION_ACK}
        if payload:
            data["payload"] = payload
        await self.send_json(data)

        if self.keep_alive:
            keep_alive_handler = self.handle_keep_alive()
            self.keep_alive_task = asyncio.create_task(keep_alive_handler)

class MyRouter(GraphQLRouter):
    graphql_ws_handler_class = MyGraphQLWSHandler
    pass

(this is using the older graphql-ws handler, since I'm working with GraphiQL when testing)

Two things:

1. It would be nice to have a method finish_connection_init() defined as part of the handler, taking an optional payload, as implemented here, and have it defined already as part of the base class, to minimize the amount of code duplication.
2. There really is no convenient place to store the results of any processing at this point. I can store it on the handler object (self.success = True) but during evaluation later, inside the graphql functions, I have no access to this state. Well, I suppose I could store it on the websockets object _ws but it is a bit dirty.

What would be the correct approach here?

Finally, I've also noticed that the requeyst has an .auth member, so authentication middleware could be used, of course.

[DoctorJohn]

Okay, I've done some other experiments. It is possible to override the Router to intecept the connection_init message

That's the intended way of doing it 💯️ we should probably document it.

Well, I suppose I could store it on the websockets object _ws but it is a bit dirty.

I faced a similar issue in the past. Maybe you could use the WebSocket objects state property. The state property is part of the HTTPConnection which WebSocket is a subclass of. (docs, source)

[kristjanvalur]

Well, I have a PR (#1652) implementing a change, making it simpler to do this subclassing, not having to replicate other logic.
But I was also wondering if this sort of thing could perhaps better be done via the Extensions mechanism? I'm unsure.