asgi_ws

Creating a standalone WebSocket Server with FastApi and JWT Authentication in Python

In this post, I will show you how to create a WebSocket server in Python that uses JWT tokens for authentication. The server is designed to be independent of the main process, making it easy to integrate into existing applications. The client-side JavaScript will handle reconnections incrementally.

The WebSocket server will be created using FastApi, the web framework built on top of Starlette. This is the entrypoint.

import logging

from fastapi import FastAPI

from asgi_ws import setup_app

logging.basicConfig(
    format='%(asctime)s [%(levelname)s] %(message)s',
    level='INFO',
    datefmt='%d/%m/%Y %X')

logger = logging.getLogger(__name__)
SECRET_KEY = "your_secret_key"

app = FastAPI()

app = setup_app(
    app=app,
    base_path='/ws',
    jwt_secret_key=SECRET_KEY,
)

The setup_app function is defined in the lib.websockets module. This function will set up the WebSocket server and the necessary routes.

def setup_app(app, jwt_secret_key: str, base_path='/ws', jwt_algorithm: str = "HS256"):
    ws_router = get_ws_router(
        jwt_secret_key=jwt_secret_key,
        jwt_algorithm=jwt_algorithm,
        base_path=base_path
    )
    current_dir = Path(__file__).parent
    static_dir = current_dir / "static/js"

    app.mount("/js", StaticFiles(directory=static_dir), name="js")
    app.include_router(ws_router)

    return app

The get_ws_router function is defined in the same module. This function will create the WebSocket router and the necessary routes.

def get_ws_router(jwt_secret_key: str, base_path='ws', jwt_algorithm: str = "HS256"):
    ws_router = APIRouter()

    current_dir = Path(__file__).parent
    static_dir = current_dir / "static/js"
    ws_router.mount(f"{base_path}/js", StaticFiles(directory=static_dir), name="js")

    manager = ConnectionManager(jwt_secret_key=jwt_secret_key, jwt_algorithm=jwt_algorithm)

    @ws_router.post(f"{base_path}/emmit")
    async def emmit_endpoint(request: Request):
        payload = await request.json()
        await manager.broadcast(payload["channel"], payload["payload"])
        return True

    @ws_router.websocket(f"{base_path}/")
    async def websocket_endpoint(websocket: WebSocket):
        token = websocket.query_params.get("token")
        if not token:
            await websocket.close(code=1008)
            raise HTTPException(status_code=401, detail="Token required")

        await manager.connect(websocket, token)
        try:
            while True:
                message: Message = await websocket.receive()
                if message["type"] == "websocket.disconnect":
                    manager.disconnect(websocket)
                    break
        except WebSocketDisconnect:
            manager.disconnect(websocket)

    return ws_router

Websockets are bidirectional communication channels that allow real-time data transfer between clients and servers, but I prefer to avoid the communication from the client to the server. When a client wants to send a message to the server, it will send an HTTP POST request to the /emit endpoint (via the main process). The server will then broadcast the message to all connected clients. The client will only receive messages from the server. Because of that we need a main wsgi process using FastApi or another web framework to handle the HTTP requests.

This an example with FastApi:

<!DOCTYPE html>
<html>
<head>
    <title>Chat</title>
</head>
<body>
<h1>WebSocket Chat</h1>

<input type="text" id="messageText" autocomplete="off"/>
<button onclick="sendMessage()">Send</button>

<ul id='messages'>
</ul>
<script src="//localhost:8000/js/websockets.js"></script>
<script>
    async function sendMessage() {
        const channel = 'chat';
        const url = `/emit/${channel}`;
        const input = document.getElementById("messageText");
        const message = input.value;
        input.value = '';
        const body = JSON.stringify({channel: 'chat1', payload: message});
        const headers = {'Content-Type': 'application/json'};

        try {
            const response = await fetch(url, {method: 'POST', headers: headers, body: body});
        } catch (error) {
            console.error('Error:', error);
        }
    }

    (async function () {
        const getToken = async () => {
            const response = await fetch('/token');
            const {token} = await response.json();
            return token;
        };

        const messageCallback = (event) => {
            const messages = document.getElementById('messages');
            const message = document.createElement('li');
            message.textContent = event.data;
            messages.appendChild(message);
        };

        const wsManager = new WebSocketManager('ws://localhost:8000/ws/', getToken, messageCallback);
        await wsManager.connect();
    })();

</script>
</body>
</html>

Metadata

Release files for asgi_ws 1.0.0

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for asgi_ws 1.0.0
File	Size	Uploaded
asgi_ws-1.0.0.tar.gz	3.7 kB	Dec 18, 2024	Details

Built distribution (wheel)

Table of built distributions (wheels) for asgi_ws 1.0.0
File	Interpreter	ABI	Platform	Reset
asgi_ws-1.0.0-py3-none-any.whl 4.8 kB Dec 18, 2024	Python 3	none	any	Details

Total release size: 8.5 kB

Release files / asgi_ws-1.0.0.tar.gz

Download URL	asgi_ws-1.0.0.tar.gz
Size	3.7 kB
Tags	Source
SHA-256 checksum How to use checksums	`8813b89b94454385e20992c3871d7c13f30336bee8e577280f21046ce11bffda`
BLAKE2b-256 checksum How to use checksums	`639c8f6edb27c95523fe32882b0e3898eb2571a1ef6c246ee40ad1e7ef8b3980`
Upload date	Dec 18, 2024
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`poetry/1.8.2 CPython/3.9.13 Darwin/24.1.0`

Release files / asgi_ws-1.0.0-py3-none-any.whl

Download URL	asgi_ws-1.0.0-py3-none-any.whl
Size	4.8 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`849a559eafb3dc72d3a8441a35ab380b4f37ea085e587b143d58ac70871636da`
BLAKE2b-256 checksum How to use checksums	`ac812e1cd029291733735d1e3aea72d250f584f9ffe8c9f56a7344247104b7fb`
Upload date	Dec 18, 2024
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`poetry/1.8.2 CPython/3.9.13 Darwin/24.1.0`