Safeguard
Application Security

Building an authenticated, TLS-secured WebSocket server in Python

WebSockets skip same-origin checks by default — CWE-1385 exists because of it. Here's how to build one in Python with origin checks, TLS, and rate limits.

Safeguard Research Team
Research
6 min read

A WebSocket handshake starts as an ordinary HTTP GET request with an Upgrade: websocket header, and that ordinary-looking detail is the root of a real vulnerability class: MITRE's CWE database catalogs it directly as CWE-1385, Missing Origin Validation in WebSockets. Unlike a same-origin-policy-protected fetch() call, a browser will happily let any page open a WebSocket connection to any host, sending along whatever cookies the browser already holds for that host. If your server never checks the Origin header sent during the handshake, a malicious page can silently ride the victim's authenticated session — a pattern known as Cross-Site WebSocket Hijacking (CSWSH). Add to that the fact that TLS termination, authentication, and rate limiting are not automatic for WebSockets the way they often are for REST endpoints behind a framework's default middleware, and it's easy to ship a WS server that's wide open in three different ways at once. This tutorial builds a Python WebSocket server using the websockets library that closes all three gaps: origin allowlisting, wss:// transport encryption, token-based authentication at handshake time, and a token-bucket rate limiter — with the specific library parameters and code patterns for each.

Why doesn't the browser's same-origin policy protect a WebSocket connection?

Same-origin policy restricts what JavaScript can read back from a cross-origin response, but it was never designed to block the connection from being opened in the first place — and WebSockets inherit that gap. A page on evil.example can execute new WebSocket("wss://bank.example/ws") and the browser will attach bank.example's cookies automatically, exactly as it would for an <img> tag. The server sees a valid, authenticated-looking handshake and, unless it inspects the Origin header itself, has no way to know the request didn't originate from its own frontend. This is functionally CSRF applied to a persistent, bidirectional channel — and because a hijacked WebSocket stays open, an attacker gets a live feed of every message the server pushes, not just a single forged action. CWE-1385 exists as a named entry precisely because this was common enough across production WebSocket deployments to warrant its own classification, distinct from CWE-352 (CSRF).

How do you validate the Origin header in Python's websockets library?

The websockets library (maintained by Aymeric Augustin, python-websockets/websockets on GitHub) exposes an origins parameter directly on serve() for exactly this purpose — pass it a list of acceptable values and the library rejects the handshake before your application code ever runs:

import websockets

async def handler(websocket):
    async for message in websocket:
        await websocket.send(f"echo: {message}")

async def main():
    async with websockets.serve(
        handler,
        "0.0.0.0",
        8443,
        origins=["https://app.example.com"],
    ):
        await asyncio.Future()  # run forever

Requests with a missing or mismatched Origin header get an HTTP 403 during the handshake, never reaching handler. This is a strict allowlist, not a substring match — wildcard or regex matching against Origin is a common way teams accidentally reintroduce the exact hole this parameter closes, since an attacker-controlled subdomain or a loosely written pattern can slip past a naive check.

How do you authenticate a connection when JavaScript can't set custom headers on it?

Browsers deliberately don't let new WebSocket() set arbitrary headers like Authorization, so the two widely documented patterns (per the websocket.org security guide and freeCodeCamp's WebSocket security writeups) are passing a short-lived token as a query parameter on the handshake URL, or relying on the session cookie the browser attaches automatically and validating it server-side. The query-parameter approach is more explicit and easier to scope:

import jwt

async def handler(websocket):
    token = websocket.request.path.split("token=")[-1]
    try:
        claims = jwt.decode(token, SECRET, algorithms=["HS256"])
    except jwt.InvalidTokenError:
        await websocket.close(code=4001, reason="invalid token")
        return
    # proceed as claims["sub"]

Validate the token before entering your message loop, and reject with a close frame rather than accepting and failing later. If you use the cookie pattern instead, treat it as a second CSRF-adjacent surface: cookie presence alone is not proof of legitimate origin, which is why origin validation and token validation are complementary controls, not substitutes for each other.

How do you get wss:// without hand-rolling TLS logic?

Because the handshake is HTTP, TLS termination for WebSockets works exactly like HTTPS: either terminate at a reverse proxy (nginx or Caddy sitting in front of your Python process) or hand websockets.serve() an ssl.SSLContext directly and let it terminate TLS in-process:

import ssl

ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
ssl_context.load_cert_chain("fullchain.pem", "privkey.pem")

async with websockets.serve(handler, "0.0.0.0", 8443, ssl=ssl_context):
    ...

A reverse proxy is generally the better production default — it centralizes certificate renewal and lets you terminate TLS once for both your HTTP API and your WebSocket upgrade path — but the in-process option matters for services that must terminate TLS themselves, such as on-prem or air-gapped deployments without a fronting proxy.

How do you rate-limit a protocol that has no built-in throttle?

Nothing in TCP or the WebSocket framing spec stops a client from sending thousands of messages per second down one open connection, so throttling has to be layered on in application code (or at a fronting proxy like nginx's limit_req, which only really governs the initial handshake request, not the ongoing message stream). A simple token-bucket counter per connection is enough for most services:

import time

async def handler(websocket):
    tokens, last = 20.0, time.monotonic()
    async for message in websocket:
        now = time.monotonic()
        tokens = min(20.0, tokens + (now - last) * 5.0)
        last = now
        if tokens < 1.0:
            await websocket.close(code=4008, reason="rate limit exceeded")
            return
        tokens -= 1.0
        await websocket.send(f"echo: {message}")

Separately, cap frame size with max_size on serve() (the library defaults to 1 MiB) and tune write_limit if you expect high-throughput streams — both guard against memory exhaustion from a single connection sending oversized or unbounded frames, a distinct failure mode from message-rate abuse. (The older read_limit argument only applies to the legacy asyncio implementation, which is deprecated; the current implementation has no library-level read buffer to size.)

How Safeguard helps

None of these controls — origin allowlisting, handshake-time token checks, TLS configuration, or a rate limiter — are things a dependency scanner will ever flag as missing, because they're application logic choices, not vulnerable package versions. Safeguard's SAST engine traces data flow through Python source, so a handler that reads a token query parameter without validating it before entering a message loop, or a websockets.serve() call with no origins argument, are the kind of source-level patterns reachability-aware analysis is built to surface with a CWE mapping attached rather than leaving them undiscovered until an incident. And for the deployed service itself, Safeguard's DAST capability — currently HTTP-request-based, with verified targets, strict scope, and mandatory rate limits and safety modes — is a reasonable first check on the HTTP surface your WebSocket server shares, even though protocol-level WebSocket testing sits outside what it documents today.

Never miss an update

Weekly insights on software supply chain security, delivered to your inbox.