Using websockets with authorization filtering

[aranvir]

Hi, I've been playing around with litestar and would like to add a websocket handler to my app with which I can send messages to authenticated and authorized clients. Basically:

• Only accept connections from authenticated clients (JWT or session id)
• Send normal message -> send to all connected clients
• Send privileged message -> send only to clients with correct authorization

I already tested the Channels plugin which fulfills the first two cases (at least tested for JWT, struggling to get it to work with cookies). Here, I just create a thread to periodically publish notifications. But I can't really figure out how to do message filtering based on user role for example.

So, I tried to do an implementation based on WebsocketListener that is attached below. However, how would I access the instance of this handler so that I can call the notify() method? Be it for a constant loop or, for example, to implement chat room behavior (here, channels is probably better, but just for the sake of understanding the solution)?

EDIT: In the meantime, managed to get this listener to work with Sessions and attached the app code below. I used httpx for testing, so my solution might be tailored to how that library formats data but I hope it's somewhat universal. So, on_connect and on_disconnect are working fine (validating the session id, getting session data). But I still don't know how I could use this for pushing messages to the client.

Websocket handler

import asyncio
import json
from dataclasses import dataclass
from typing import Any

from litestar import WebSocket, Router
from litestar import exceptions
from litestar.handlers import WebsocketListener
from litestar.stores.base import Store
from litestar.logging.config import LoggingConfig


logger = LoggingConfig().configure()()


@dataclass
class Notification:
    content: str
    secret: bool


class NotificationsWS(WebsocketListener):
    path = "/ws/updates"

    def __init__(self, owner: Router):
        super().__init__(owner)
        self._connections = dict()
        self._privileged_connections = dict()

    async def _get_session_from_cookies(self, cookies: dict) -> dict:
        """Workaround to fetch the session data for the current session id.

        Since the websocket connection is not in the same session as the HTTP connection they don't share cookies.
        Additionally, cookies passed to websocket extra_headers are a bit mangled (e.g., are not interpreted as JSON)
        so we need to do some parsing here.

        In the end, we return the session dictionary that we would usually expect from `connection.session` or `request.session`.
        """
        # session.cookies is a dict with a single key (empty string) mapped to whatever we passed to the cookies
        cookie_string = cookies.get("", None)
        if cookie_string is None:
            raise exceptions.WebSocketException(detail="No cookies found.")
        # but the content is string (and somehow not json formatted)
        cookies = json.loads(cookie_string.replace("'", '"'))
        session_id = cookies.get("session", None)
        if session_id is None:
            raise exceptions.WebSocketException(detail="No session cookie found.")
        # And then, although we fetch an existing memory store here,
        # it returns the data as bytes instead of a dict so we decode that
        session_store: Store = self._owner.stores.get("sessions")
        encoded_session = await session_store.get(session_id)

        if encoded_session is None:
            raise exceptions.WebSocketException(detail="Invalid session id")
        session = json.loads(encoded_session)

        return session

    async def on_accept(self, socket: WebSocket) -> None:
        session = await self._get_session_from_cookies(socket.cookies)
        if session.get("clearance"):
            self._privileged_connections[session.get('user_id')] = socket
        else:
            self._connections[session.get('user_id')] = socket
        logger.info(f"User {session.get('user_id')} connected. Clearance: {session.get('clearance')}")

    async def on_disconnect(self, socket: WebSocket) -> None:
        session = await self._get_session_from_cookies(socket.cookies)
        if session.get('clearance'):
            self._privileged_connections.pop(session.get('user_id'))
        else:
            self._connections.pop(session.get('user_id'))
        logger.info("Connection closed")

    async def on_receive(self, data: Any) -> Any:
        return "This endpoint does not accept requests."

    async def notify(self, message: Notification):
        """Send a notification to all authorized users."""
        if message.secret:
            calls = [socket.send_text(message.content) for socket in self._privileged_connections.values()]
        else:
            calls = [
                socket.send_text(message.content) for socket in self._privileged_connections.values()
            ] + [
                socket.send_text(message.content) for socket in self._connections.values()
            ]
        await asyncio.gather(*calls)

app code

from litestar import Litestar, Request, get, post
from litestar.middleware.session.server_side import ServerSideSessionBackend, ServerSideSessionConfig
from litestar.security.session_auth import SessionAuth
from litestar.stores.memory import MemoryStore

from mydemo import db_config
from mydemo.model import User
from mydemo.controller import UserController, MovieController, retrieve_user_handler
from mydemo.notifications import NotificationsWS


session_config = ServerSideSessionConfig()
session_auth = SessionAuth[User, ServerSideSessionBackend](
    retrieve_user_handler=retrieve_user_handler,
    session_backend_config=session_config,
    # exclude any URLs that should not have authentication.
    exclude=["/login", "/schema", "/ws/updates"],
)


app = Litestar(
    route_handlers=[UserController, MovieController, NotificationsWS],
    on_startup=[db_config.on_startup],
    on_app_init=[session_auth.on_app_init],
    stores={"sessions": MemoryStore()},
    plugins=[db_config.sqlalchemy_plugin],
    debug=True
)

[provinzkraut]

But I can't really figure out how to do message filtering based on user role for example.

Well you've already to started to recreate some of the Channels logic in your app.
The general idea would be that you'd create two channels: One for all users, and one for authenticated users only. Once a user connects, you'd subscribe them to the public channel, and if they're authenticated, to the privileged channel as well. Basically something like:

@websocket("/ws")
async def handler(socket: WebSocket, channels: ChannelsPlugin) -> None:
    await socket.accept()

    channels_to_subscribe = ["public"]
    if is_authenticated(socket):
        channels_to_subscribe.append("privileged")

    async with channels.subscribe(channels_to_subscribe) as subscriber:
        ...

[provinzkraut]

Thanks! I will need to do some more reading to understand what's going on there.

I've written a brief article on this sort of usage here:
https://blog.litestar.dev/creating-a-websockets-chat-in-just-30-lines-with-litestar-2ca0f3767a47

If anything's unclear or there are areas where we could improve the docs, please let me know.

My motivation for the class-based approach was that I can update subscriptions throughout the session. E.g., if the clearance is revoked, I want to be able to adjust that in the session data and in the subscription.

The issue with your approach is that it only works for a single process. A lot of the heavy lifting done by the channels is with regards to inter-process communication. For example, if you have 2 processes, and you initiate the "revoke" action in one of them, you can't be sure that the client you're trying to revoke is connected within the same process, so you'd have to somehow relay that information.

Keeping track of the subscribers and channels across processes is done by Channels for you, that's why I wouldn't recommend building your own in-memory logic for that. At the very least, you'd have to use e.g. Redis for this, where you'd end up with some sort of queue, which, in the end, would most likely resemble a basic re-implementation of Channels again 😉

My suggestion would be to actually make use of Channels for this by implementing an additional control channel. Then you could send a message there (e.g. {"event": "unsubscribe", "data": ["privileged"]}) that would then cause the subscriber to unsubscribe from that channel.

[aranvir]

Thanks again for the extensive feedback! I'll need to experiment a bit more with this.

I think in general my problem is - I'm a beginner/intermediate when it comes to web development. And litestar is amazing for providing solutions out of the box for simple and for complex stuff - but a lot is hidden under the hood 😅 so then I often don't know what is happening or where it's happening or why it's happening. Simple example - I didn't know that websockets would not share a session with other http endpoints or are in general a bit different to dealing with cookies than regular http requests :D And that is obviously not the fault of litestar or a shortcoming of the documentation.

[aranvir]

I kept pondering about the documentation and if there is something to improve.

I think I had the hardest time wrapping my head around all the "implicit dependency injections". This is not related to my original question, but it was often not clear to me which dependencies I could use where and how they get added to the workflow (e.g., session, connection, request, response, db_connection, etc.)
I think a "primer" or "summary" on how litestar handles this could be useful. Like, which of these optional keyword arguments can always be added to an endpoint, which need a certain plugin active, etc.

[provinzkraut]

Like, which of these optional keyword arguments can always be added to an endpoint, which need a certain plugin active, etc.

This is generally documented here:

https://docs.litestar.dev/latest/usage/routing/handlers.html#reserved-keyword-arguments

I think it's not part of the tutorial though, it might be worth adding them there or to the dependency injection documentation.

I think I had the hardest time wrapping my head around all the "implicit dependency injections".

The dependency injection itself is not very implicit actually. Aside from those parameters reserved by Litestar itself, nothing will get injected unless explicitly declared.

Thanks again for the extensive feedback! I'll need to experiment a bit more with this.

Let me know if there's anything left unclear or can be improved in the documentation.

simple example - I didn't know that websockets would not share a session with other http endpoints or are in general a bit different to dealing with cookies than regular http requests :D

As a general rule of thumb, if you feel like you're having to write a huge amount of boilerplate for a seemingly simple task, it's a good idea to consult the documentation. I know it can be a bit vast but it also has to cover a lot of ground. The usage docs generally aren't meant to be read cover to cover, but going through at least the core parts about routing, requests and dependency injection is something I'd recommend as a next step after the tutorials.

The docs are in constant improvement, so feedback on how to make them easier to navigate and/or digest is highly welcome.

[aranvir]

Thanks for the references, I'll do some more reading ;)

Yea that was exactly my thoughts, once I started hacking away this solution - this can't be right :D

If I come up with suggestions for doc improvements I will provide them!