Securing WebSocket connections with OAuth2

Securing WebSocket connections is paramount to safeguarding the privacy and security of user interactions in real-time applications. By implementing authentication and access control mechanisms, developers can mitigate risks associated with unauthorized access, eavesdropping, and data tampering. In this recipe, we will see how to create a secure WebSocket connection endpoint with OAuth2 token authorization in your FastAPI applications.

Getting ready

To follow the recipe, you should already know how to set up a basic WebSocket endpoint – explained in the Setting up WebSockets in FastAPI recipe in this chapter.

Furthermore, we are going to use OAuth2 with a password and a bearer token. We will apply the same strategy we used to secure HTTP endpoints in the Securing your API with OAuth2 recipe in Chapter 3, Building RESTful APIs with FastAPI. Feel free to have a look before starting the recipe.

Before starting the recipe, let’s create a simple WebSocket endpoint, /secured-ws, in the main.py module:

@app.websocket("/secured-ws")
async def secured_websocket(
    websocket: WebSocket,
    username: str
):
    await websocket.accept()
    await websocket.send_text(f"Welcome {username}!")
    async for data in websocket.iter_text():
        await websocket.send_text(
            f"You wrote: {data}"
        )

The endpoint will accept any connection with a parameter to specify the username. Then it will send a welcome message to the client and return each message received to the client.

The endpoint is insecure since it does not have any protection and can be easily reached. Let’s dive into the recipe to see how to protect it with OAuth2 authentication.

How to do it…

At the time of writing, there is no support for the OAuth2PasswordBearer class for WebSocket in FastAPI. This means that checking the bearer token in the headers for WebSocket is not as straightforward as it is for HTTP calls. However, we can create a WebSocket-specific class that is derived from the one used by HTTP to achieve the same functionality as follows.

1. Let’s do it in a dedicated module under the app folder called ws_password_bearer.py:

   from fastapi import (
       WebSocket,
       WebSocketException,
       status,
   )
   from fastapi.security import OAuth2PasswordBearer
   class OAuth2WebSocketPasswordBearer(
       OAuth2PasswordBearer
   ):
       async def __call__(
           self, websocket: WebSocket
       ) -> str:
           authorization: str = websocket.headers.get(
               "authorization"
           )
           if not authorization:
               raise WebSocketException(
                   code=status.HTTP_401_UNAUTHORIZED,
                   reason="Not authenticated",
               )
           scheme, param = authorization.split()
           if scheme.lower() != "bearer":
               raise WebSocketException(
                   code=status.HTTP_403_FORBIDDEN,
                   reason=(
                       "Invalid authentication "
                       "credentials"
                   ),
               )
           return param

   We will use it to create a get_username_from_token function to retrieve the username from the token. You can create the function in a dedicated module – security.py.

2. Let’s define the oauth2_scheme_for_ws object:

   from app.ws_password_bearer import (
       OAuth2WebSocketPasswordBearer,
   )
   oauth2_scheme_for_ws = OAuth2WebSocketPasswordBearer(
       tokenUrl="/token"
   )

3. The tokenUrl argument specifies the callback endpoint to call to retrieve the token. This endpoint should be built according to the token resolution you use. After that, we can create a function that retrieves the username from the token:

   def get_username_from_token(
       token: str = Depends(oauth2_scheme_for_ws),
   ) -> str:
       user = fake_token_resolver(token)
       if not user:
           raise WebSocketException(
               code=status.HTTP_401_UNAUTHORIZED,
               reason=(
                   "Invalid authentication credentials"
               )
           )
       return user.username

   The purpose of the fake_token_resolver function is to simulate the process of resolving a token. This function can be found in the security.py file in the GitHub repository of the chapter. Furthermore, the example contains only two users, johndoe and janedoe, who can be used later for testing. Also, the security.py module from the GitHub repository contains the POST /token endpoint to be used to retrieve the token.

   However, it is important to mention that this function does not provide any actual security and it is only used for example purposes. In a production environment, it is recommended to use a JWT Authorization token or an external provider for token resolution (see the Working with OAuth2 and JWT for authentication and Using third-party authentication recipes – both in Chapter 4, Authentication and Authorization).

4. Now let’s use it to secure our WebSocket endpoint, /secured-ws, in the main.py module:

   from import Annotated
   from fastapi import Depends
   from app.security import get_username_from_token
   @app.websocket("/secured-ws")
   async def secured_websocket(
       websocket: WebSocket,
       username: Annotated[
           get_username_from_token, Depends()
       ]
   ):
       # rest of the endpoint

This is all you need to build a secured WebSocket endpoint.

To test it, spin up the server from the terminal by running the following:

$ uvicorn app.main:app

When attempting to connect to the WebSocket endpoint using Postman or another tool to the address ws://localhost:8000/secured-ws, an authorization error will occur, and the connection will be rejected before the handshake.

To allow the connection, we need to retrieve the token and pass it through the headers of the WebSocket request in Postman. You can retrieve the token from the dedicated endpoint or, if you use the fake token generator from the GitHub repository, you simply append the tokenized string to the username. For example, for johndoe, the token would be tokenizedjohndoe.

Let’s pass it through the header. In Postman, you can pass the bearer token to the WebSocket request in the Headers tab by adding a new header. The header will have a key called Authorization and value that will be bearer tokenizedjohndoe.

Now, if you try to connect, it should connect and you will be able to exchange messages with the endpoint.

You have just secured a WebSocket endpoint in FastAPI. By implementing OAuth2 authorization, you can enhance the security posture of your FastAPI applications and safeguard WebSocket communication against potential threats and vulnerabilities.

Exercise

Try to build a secure chat functionality where users need to log in to participate in the chat.

Tips: The endpoint that returns the HTML page should check for the bearer token in the cookies. If the cookie is not found or the bearer token is not valid, it should redirect the client to a login page that puts the token in the browser’s cookies.

You can use the response.RedirectResponse class from the fastapi package to handle redirections. The usage is quite straightforward and you can have a look at the documentation page at the link:

https://fastapi.tiangolo.com/advanced/custom-response/#redirectresponse.

See also

Integrating OAuth2 into WebSockets in FastAPI with an OAuth2PasswordBearer-like class is a current topic of interest, and it is expected to evolve quickly over time. You can follow the ongoing discussion in the FastAPI GitHub repository:

• OAuth2PasswordBearer with WebSocket Discussion: https://github.com/tiangolo/fastapi/discussions/8983