Inspecting incoming websocket data in middleware

[nickolas1]

In the Inspecting or modifying the request section of the docs, they state:

Likewise, WebSocket middleware may manipulate websocket.receive ASGI event messages to inspect or alter incoming WebSocket data.

I've spent the afternoon looking with no luck, are there any good examples of doing this? An initial attempt with something like this:

    message = await receive()
    # inspect, take action, etc.

    async def wrapped_receive() -> Message:
        return message
        
    await self.app(scope, wrapped_receive, send)    

fails while negotiating the connection because of a state change in the WebSocket class to self.client_state = WebSocketState.CONNECTED on the first receive.

If I hack around and bypass the receive by inspecting headers for connection: Upgrade the middleware is not actually called when the client sends data after the initial connection. That's a separate issue I'm also looking into, but for both questions I'd love to see some examples of this kind of thing if anyone has some.

[Kludex]

You should call the message = await receive() inside the wrapped_receive.

Example:

from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.routing import WebSocketRoute
from starlette.types import ASGIApp, Message, Receive, Scope, Send
from starlette.websockets import WebSocket


class WSMiddleware:
    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "websocket":
            return await self.app(scope, receive, send)

        async def wrapped_receive() -> Message:
            message = await receive()
            print(message)
            return message

        await self.app(scope, wrapped_receive, send)


async def ws_endpoint(websocket: WebSocket) -> None:
    await websocket.accept()
    await websocket.send_text("Hello, world!")
    await websocket.close()


app = Starlette(
    routes=[WebSocketRoute("/", ws_endpoint)],
    middleware=[Middleware(WSMiddleware)],
)