Security on Websocket

[developer992]

Hello

I am trying to get this working: https://fastapi.tiangolo.com/advanced/websockets/#using-depends-and-others

Here is my basic example that fails:

from fastapi import APIRouter, WebSocket, Security, Depends
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from pydantic import BaseModel

class User(BaseModel):
    name: str
    email: str



async def login_required(
    websocket: WebSocket,
    bearer_key: HTTPAuthorizationCredentials = Security(
        HTTPBearer(auto_error=False, description="Add the token to your bearer authentication")
    ),
) -> User:
    # all good
    return User(name='Joe', email='joe@gmail.com')

@router.websocket('/ws')
async def chat(
    websocket: WebSocket,
    user: User = Security(login_required),
):
    logger.debug(f'CONNECTION REQUEST from user={user}')
    
    await websocket.accept()
    while True:
        raw_message: str = await websocket.receive_text()
        await websocket.send_text(f"Thank you for {raw_message}")

Regular REST API endpoint works with HTTPAuthorizationCredentials ... it takes the "Authorization Bearer XXX" header, but not when using websocket endpoint ...

What am i misunderstanding?

I dont see any errors in my logs though, not sure why... i just get "Error: Unexpected server response: 500"

Works for regular REST API endpoints though:

@router.get('/preview')
async def get_preview_url(
    item: str,
    zoom: int = 1,
    page: str = "1",
    user: User = Security(login_required),
):
    pass 

I also tried Depends() instead of Security() but it's the same result.

Many thanks!

[developer992]

Hmm i thought HTTPBearer is problematic, but this doesn't work either ...

async def get_token(
    websocket: WebSocket,
    token: Annotated[Union[str, None], Query()] = None,
):
    if token is None:
        raise WebSocketException(code=111)
    
    return token


@router.websocket('/ws')
async def chat_ws(
    *,
    websocket: WebSocket,
    token: Annotated[str, Depends(get_token)],
):

This should work via query param, right, like so:

localhost:8000/chat/ws?token=abc

Isn't this the same as in docs? i'm confused now ...

[tomasvotava]

To be honest, I am not very experienced with WebSockets, your example looks almost the same as the one in FastAPI docs. When you say "it doesn't work", what exactly happens? Could you try running it either in a debugger or adding some logging to make sure that your dependency get_token is actually called?

[developer992]

Hey, sorry, i solved this using sub-protocol header.

Websocket connection does not support headers, except a few specific ones, for instance "Sub-Protocol" which is often used to hack/hide the token with. Wanted to avoid query params as to not expose the token.

Basically, the header will look like this:

Sec-WebSocket-Protocol: token, eyabcblabla

Which can then be parsed on the backend side and validate the token.

Fastapi code:

WEBSOCKET_TOKEN_SUBPROTOCOL = "token"    # this is the name of the subprotocol and can be any string

async def login_required_ws(
    websocket: WebSocket,
) -> User:
    """ 
    Websocket endpoints do not support all headers, so we need to send the token in sub-protocol header
    Function will return User object not defined here but basically, token is decoded and user object is initialized and returned.
    """

    protocol, token = websocket.headers.get('sec-websocket-protocol', '').split(',')
    token = token.strip()

    if token and protocol == WEBSOCKET_TOKEN_SUBPROTOCOL:
        return await _login_required(token=token)
    
    raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION, reason="Authentication failed!")


@router.websocket('/ws')
async def chat_ws(
    *,
    websocket: WebSocket,
    user: User = Security(login_required_ws),
):

    await websocket.accept(subprotocol=WEBSOCKET_TOKEN_SUBPROTOCOL)   # this tells it which subprotocols are supported by our backend
    logger.debug(f'New connection established!')
    while True:
        ....
        # we can access logged in user here