WebSocket fails via CSRFMiddleware

[johnpaulett]

Ideally, I'd like to use the CSRFMiddleware on a websocket route. But at present, the CSRFMiddleware makes hits an assertion error whenever the websocket route is accessed:

Error

.venv/lib/python3.11/site-packages/starlette/testclient.py:91: in __enter__
    message = self.receive()
.venv/lib/python3.11/site-packages/starlette/testclient.py:160: in receive
    raise message
.venv/lib/python3.11/site-packages/anyio/from_thread.py:219: in _call_func
    retval = await retval
.venv/lib/python3.11/site-packages/starlette/testclient.py:118: in _run
    await self.app(scope, receive, send)
.venv/lib/python3.11/site-packages/fastapi/applications.py:289: in __call__
    await super().__call__(scope, receive, send)
.venv/lib/python3.11/site-packages/starlette/applications.py:122: in __call__
    await self.middleware_stack(scope, receive, send)
.venv/lib/python3.11/site-packages/starlette/middleware/errors.py:149: in __call__
    await self.app(scope, receive, send)
.venv/lib/python3.11/site-packages/starlette/middleware/cors.py:75: in __call__
    await self.app(scope, receive, send)
.venv/lib/python3.11/site-packages/starlette/middleware/sessions.py:86: in __call__
    await self.app(scope, receive, send_wrapper)
.venv/lib/python3.11/site-packages/starlette_csrf/middleware.py:55: in __call__
    request = Request(scope)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <starlette.requests.Request object at 0x137109750>
scope = {'app': <fastapi.applications.FastAPI object at 0x117817e10>, 'client': ['testclient', 50000], 'headers': [(b'host', '...'), (b'connection', b'upgrade'), ...], 'path': '/ws, ...}
receive = <function empty_receive at 0x102887600>, send = <function empty_send at 0x105c49b20>

    def __init__(
        self, scope: Scope, receive: Receive = empty_receive, send: Send = empty_send
    ):
        super().__init__(scope)
>       assert scope["type"] == "http"
E       AssertionError

.venv/lib/python3.11/site-packages/starlette/requests.py:197: AssertionError

Attempted solutions

• Failed: Added my websocket route to exempt_urls -- think the error happens before this is every applied
• Failed: Added my websocket route to required_urls -- I hoped that maybe the initial HTTP connection that gets upgrade would pass thru
• Tried replacing request = Request(scope) with request = WebSocket(scope, receive=receive, send=send) if scope["type"] == "websocket" else Request(scope), but still occurred

Current workaround

I wrap CSRFMiddleware and only pass HTTP requests into it. This is suboptimal because I would like to enforce CSRF protection for my websocket route.

from starlette_csrf.middleware import CSRFMiddleware as _CSRFMiddleware
 
class CSRFMiddleware(_CSRFMiddleware):
    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        # type="websocket" will raise an exception at present
        if scope["type"] == "http":
            await super().__call__(scope, receive, send)
        else:
            await self.app(scope, receive, send)

Partial test case

I tried to document the error in a testcase, but using the httpx client does not expose the .websocket_connect() that starlette's testclient exposes. So this code does not yet fully work

def get_app(**middleware_kwargs) -> Starlette:
    async def get(request: Request):
        return JSONResponse({"hello": "world"})

    async def post(request: Request):
        json = await request.json()
        return JSONResponse(json)
    
    async def websocket(websocket: WebSocket):
        await websocket.accept()
        data = await websocket.receive_text()
        await websocket.send_text(data)
        await websocket.close()

    app = Starlette(
        debug=True,
        routes=[
            Route("/get", get, methods=["GET"]),
            Route("/post1", post, methods=["POST"]),
            Route("/post2", post, methods=["POST"]),
            WebSocketRoute("/ws", websocket),
        ],
        middleware=[Middleware(CSRFMiddleware, secret="SECRET", **middleware_kwargs)],
    )

    return app

@pytest.mark.asyncio
async def test_valid_websocket():
    async with get_test_client(get_app()) as client:
        response_get = await client.get("/get")
        csrf_cookie = response_get.cookies["csrftoken"]

        async with client.websocket_connect("/ws") as websocket:
            await websocket.send_text("hello world")
            data = await websocket.receive_text()
            assert data == "hello world"

Happy to try to help with some pointers.

Upvote & Fund

• We're using Polar.sh so you can upvote and help fund this issue.
• We receive the funding once the issue is completed & confirmed by you.
• Thank you in advance for helping prioritize & fund our backlog.

[frankie567]

Indeed, I've never considered the case of WebSocket, so it's probable the code is not adequate for such requests.

I'm not sure right now how this can be solved, but I'll investigate. Thank you for the report 🙏

[gantoine]

@johnpaulett Not sure if this'll work for ya but I "fixed" this by only applying CSRF to http requests

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

        await super().__call__(scope, receive, send)

[merlinz01]

Sharing my workaround FWIW. Because of whatwg/websockets#16 I am putting the CSRF token as a separate cookie instead of a header. I don't think it would be considered as secure but it's better than nothing.

from starlette.websockets import WebSocket

# Most of this logic is copied from the existing implementation.
class CSRFMiddleware(StarletteCSRFMiddleware):
    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] == "websocket":
            request = WebSocket(scope, receive, send)
            csrf_cookie = request.cookies.get(self.cookie_name)

            if self._url_is_required(request.url) or (
                not self._url_is_exempt(request.url)
                and self._has_sensitive_cookies(request.cookies)
            ):
                submitted_csrf_token = request.cookies.get("header_x_csrf_token")
                if (
                    not csrf_cookie
                    or not submitted_csrf_token
                    or not self._csrf_tokens_match(csrf_cookie, submitted_csrf_token)
                ):
                    response = self._get_error_response(request)  # type: ignore
                    await response(scope, receive, send)
                    return

            send = functools.partial(self.send, send=send, scope=scope)
            await self.app(scope, receive, send)
        else:
            await super().__call__(scope, receive, send)

    async def send(self, message: Message, send: Send, scope: Scope) -> None:
        request = (
            WebSocket(scope, None, send)  # type: ignore
            if scope["type"] == "websocket"
            else Request(scope)
        )
        csrf_cookie = request.cookies.get(self.cookie_name)

        if csrf_cookie is None:
            message.setdefault("headers", [])
            headers = MutableHeaders(scope=message)

            cookie: http.cookies.BaseCookie = http.cookies.SimpleCookie()
            cookie_name = self.cookie_name
            cookie[cookie_name] = self._generate_csrf_token()
            cookie[cookie_name]["path"] = self.cookie_path
            cookie[cookie_name]["secure"] = self.cookie_secure
            cookie[cookie_name]["httponly"] = self.cookie_httponly
            cookie[cookie_name]["samesite"] = self.cookie_samesite
            if self.cookie_domain is not None:
                cookie[cookie_name]["domain"] = self.cookie_domain  # pragma: no cover
            headers.append("set-cookie", cookie.output(header="").strip())

        await send(message)