Websocket

WebSocket Security Skill

File Organization

• SKILL.md: Core principles, patterns, essential security (this file)
• references/security-examples.md: CSWSH examples and authentication patterns
• references/advanced-patterns.md: Connection management, scaling patterns
• references/threat-model.md: Attack scenarios including CSWSH

Validation Gates

Gate 0.2: PASSED (5+ vulnerabilities documented) - CVE-2024-23898, CVE-2024-26135, CVE-2023-0957

---

1. Overview

Risk Level: HIGH

Justification: WebSocket connections bypass Same-Origin Policy protections, making them vulnerable to Cross-Site WebSocket Hijacking (CSWSH). Persistent connections require careful authentication, session management, and input validation.

You are an expert in WebSocket security, understanding the unique vulnerabilities of persistent bidirectional connections.

Core Expertise Areas

• CSWSH (Cross-Site WebSocket Hijacking) prevention
• Origin header validation and token-based authentication
• Message validation and per-message authorization
• Rate limiting and connection lifecycle security

---

2. Core Responsibilities

Fundamental Principles

1. TDD First: Write tests before implementation - test security boundaries, connection lifecycle 2. Performance Aware: Optimize for low latency (<50ms), connection pooling, backpressure 3. Validate Origin: Always check Origin header against explicit allowlist 4. Authenticate First: Verify identity before accepting messages 5. Authorize Each Action: Don't assume connection equals unlimited access 6. Validate All Messages: Treat WebSocket messages as untrusted input 7. Limit Resources: Rate limit messages, timeout idle connections

Security Decision Framework

---

3. Technical Foundation

Version Recommendations

Security Configuration

WEBSOCKET_CONFIG = {
    "max_message_size": 1024 * 1024,  # 1MB
    "max_connections_per_ip": 10,
    "idle_timeout_seconds": 300,
    "messages_per_minute": 60,
}

# NEVER use "*" for origins
ALLOWED_ORIGINS = ["https://app.example.com", "https://admin.example.com"]

---

4. Implementation Workflow (TDD)

Step 1: Write Failing Test First

import pytest
from httpx import AsyncClient, ASGITransport
from fastapi.testclient import TestClient

# Test security boundaries first
@pytest.mark.asyncio
async def test_origin_validation_rejects_invalid():
    """CSWSH prevention - must reject invalid origins."""
    async with AsyncClient(
        transport=ASGITransport(app=app),
        base_url="http://test"
    ) as client:
        # This should fail until origin validation is implemented
        with pytest.raises(Exception):
            async with client.websocket_connect(
                "/ws?token=valid",
                headers={"Origin": "https://evil.com"}
            ):
                pass

@pytest.mark.asyncio
async def test_authentication_required():
    """Must reject connections without valid token."""
    with TestClient(app) as client:
        with pytest.raises(Exception):
            with client.websocket_connect("/ws") as ws:
                pass

@pytest.mark.asyncio
async def test_message_authorization():
    """Each message action must be authorized."""
    with TestClient(app) as client:
        with client.websocket_connect(
            "/ws?token=readonly_user",
            headers={"Origin": "https://app.example.com"}
        ) as ws:
            ws.send_json({"action": "delete", "id": "123"})
            response = ws.receive_json()
            assert response.get("error") == "Permission denied"

Step 2: Implement Minimum to Pass

# Implement only what's needed to pass the test
async def validate_origin(websocket: WebSocket) -> bool:
    origin = websocket.headers.get("origin")
    if not origin or origin not in ALLOWED_ORIGINS:
        await websocket.close(code=4003, reason="Invalid origin")
        return False
    return True

Step 3: Refactor and Verify

# Run all WebSocket tests
pytest tests/websocket/ -v --asyncio-mode=auto

# Check for security issues
bandit -r src/websocket/

# Verify no regressions
pytest tests/ -v

---

5. Performance Patterns

Pattern 1: Connection Pooling

# BAD - Create new connection for each request
ws = await create_connection(user_id)  # Expensive!

# GOOD - Reuse connections from pool
class ConnectionPool:
    def __init__(self, max_size: int = 100):
        self.connections: dict[str, WebSocket] = {}

    async def get_or_create(self, user_id: str) -> WebSocket:
        if user_id not in self.connections:
            self.connections[user_id] = await create_connection(user_id)
        return self.connections[user_id]

Pattern 2: Message Batching

# BAD - Send messages one at a time
for item in items:
    await websocket.send_json({"type": "item", "data": item})

# GOOD - Batch messages to reduce overhead
await websocket.send_json({"type": "batch", "data": items[:50]})

Pattern 3: Binary Protocols

# BAD - JSON for high-frequency data (~80 bytes)
await websocket.send_json({"x": 123.456, "y": 789.012, "z": 456.789})

# GOOD - Binary format (20 bytes)
import struct
await websocket.send_bytes(struct.pack('!3f', 123.456, 789.012, 456.789))

Pattern 4: Heartbeat Optimization

# BAD - Fixed frequent heartbeats
HEARTBEAT_INTERVAL = 5  # Every 5 seconds

# GOOD - Adaptive heartbeats based on activity
interval = 60 if (time() - last_activity) < 60 else 30

Pattern 5: Backpressure Handling

# BAD - Blocks on slow clients
await ws.send_json(message)

# GOOD - Timeout and bounded queue
from collections import deque
queue = deque(maxlen=100)  # Drop oldest when full
try:
    await asyncio.wait_for(ws.send_json(message), timeout=1.0)
except asyncio.TimeoutError:
    pass  # Client too slow

---

6. Implementation Patterns

Pattern 1: Origin Validation (Critical for CSWSH Prevention)

from fastapi import WebSocket

async def validate_origin(websocket: WebSocket) -> bool:
    """Validate WebSocket origin against allowlist."""
    origin = websocket.headers.get("origin")
    if not origin or origin not in ALLOWED_ORIGINS:
        await websocket.close(code=4003, reason="Invalid origin")
        return False
    return True

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    if not await validate_origin(websocket):
        return
    await websocket.accept()

Pattern 2: Token-Based Authentication

from jose import jwt, JWTError

async def authenticate_websocket(websocket: WebSocket) -> User | None:
    """Authenticate via token (not cookies - vulnerable to CSWSH)."""
    token = websocket.query_params.get("token")
    if not token:
        await websocket.close(code=4001, reason="Authentication required")
        return None
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
        user = await user_service.get(payload.get("sub"))
        if not user:
            await websocket.close(code=4001, reason="User not found")
            return None
        return user
    except JWTError:
        await websocket.close(code=4001, reason="Invalid token")
        return None

Pattern 3: Per-Message Authorization

from pydantic import BaseModel, field_validator

class WebSocketMessage(BaseModel):
    action: str
    data: dict

    @field_validator('action')
    @classmethod
    def validate_action(cls, v):
        if v not in {'subscribe', 'unsubscribe', 'send', 'query'}:
            raise ValueError(f'Invalid action: {v}')
        return v

async def handle_message(websocket: WebSocket, user: User, raw_data: dict):
    try:
        message = WebSocketMessage(**raw_data)
    except ValueError:
        await websocket.send_json({"error": "Invalid message format"})
        return

    if not user.has_permission(f"ws:{message.action}"):
        await websocket.send_json({"error": "Permission denied"})
        return

    result = await handlers[message.action](user, message.data)
    await websocket.send_json(result)

Pattern 4: Connection Manager with Rate Limiting

from collections import defaultdict
from time import time

class SecureConnectionManager:
    def __init__(self):
        self.connections: dict[str, WebSocket] = {}
        self.message_counts: dict[str, list[float]] = defaultdict(list)
        self.connections_per_ip: dict[str, int] = defaultdict(int)

    async def connect(self, websocket: WebSocket, user_id: str, ip: str) -> bool:
        if self.connections_per_ip[ip] >= WEBSOCKET_CONFIG["max_connections_per_ip"]:
            await websocket.close(code=4029, reason="Too many connections")
            return False
        await websocket.accept()
        self.connections[user_id] = websocket
        self.connections_per_ip[ip] += 1
        return True

    def check_rate_limit(self, user_id: str) -> bool:
        now = time()
        self.message_counts[user_id] = [
            ts for ts in self.message_counts[user_id] if ts > now - 60
        ]
        if len(self.message_counts[user_id]) >= WEBSOCKET_CONFIG["messages_per_minute"]:
            return False
        self.message_counts[user_id].append(now)
        return True

Pattern 5: Complete Secure Handler

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    if not await validate_origin(websocket):
        return
    user = await authenticate_websocket(websocket)
    if not user:
        return

    ip = websocket.client.host
    if not await manager.connect(websocket, user.id, ip):
        return

    try:
        while True:
            raw = await asyncio.wait_for(
                websocket.receive_json(),
                timeout=WEBSOCKET_CONFIG["idle_timeout_seconds"]
            )
            if not manager.check_rate_limit(user.id):
                await websocket.send_json({"error": "Rate limited"})
                continue
            await handle_message(websocket, user, raw)
    except (WebSocketDisconnect, asyncio.TimeoutError):
        pass
    finally:
        manager.disconnect(user.id, ip)

---

7. Security Standards

Domain Vulnerability Landscape

OWASP Top 10 Mapping

CSWSH Prevention Summary

async def secure_websocket_handler(websocket: WebSocket):
    # 1. VALIDATE ORIGIN (Critical)
    if websocket.headers.get("origin") not in ALLOWED_ORIGINS:
        await websocket.close(code=4003)
        return
    # 2. AUTHENTICATE with token (not cookies)
    user = await validate_token(websocket.query_params.get("token"))
    if not user:
        await websocket.close(code=4001)
        return
    # 3. Accept only after validation
    await websocket.accept()
    # 4. AUTHORIZE each message, 5. RATE LIMIT, 6. TIMEOUT idle

---

8. Common Mistakes & Anti-Patterns

No Origin Validation

# NEVER - vulnerable to CSWSH
@app.websocket("/ws")
async def vulnerable(websocket: WebSocket):
    await websocket.accept()  # Accepts any origin!

# ALWAYS - validate origin first
if websocket.headers.get("origin") not in ALLOWED_ORIGINS:
    await websocket.close(code=4003)
    return

Cookie-Only Authentication

# NEVER - cookies sent automatically in CSWSH attacks
session = websocket.cookies.get("session")

# ALWAYS - require explicit token parameter
token = websocket.query_params.get("token")

No Per-Message Authorization

# NEVER - assumes connection = full access
if data["action"] == "delete":
    await delete_resource(data["id"])

# ALWAYS - check permission for each action
if not user.has_permission("delete"):
    return {"error": "Permission denied"}

No Input Validation

# NEVER - trust WebSocket messages
await db.execute(f"SELECT * FROM {data['table']}")  # SQL injection!

# ALWAYS - validate with Pydantic
message = WebSocketMessage(**data)

---

9. Pre-Implementation Checklist

Phase 1: Before Writing Code

• [ ] Write failing tests for security boundaries (CSWSH, auth, authz)
• [ ] Write failing tests for connection lifecycle (connect, disconnect, timeout)
• [ ] Write failing tests for message validation and rate limiting
• [ ] Review threat model in references/threat-model.md
• [ ] Identify performance requirements (latency, throughput, connections)

Phase 2: During Implementation

• [ ] Origin validation against explicit allowlist
• [ ] Token-based authentication (not cookie-only)
• [ ] Per-message authorization checks
• [ ] Rate limiting and idle timeout implemented
• [ ] All messages validated with Pydantic
• [ ] Connection pooling for efficiency
• [ ] Backpressure handling for slow clients

Phase 3: Before Committing

• [ ] All security tests pass: pytest tests/websocket/ -v
• [ ] No security issues: bandit -r src/websocket/
• [ ] WSS (TLS) enforced in production config
• [ ] CSWSH test coverage verified
• [ ] Performance benchmarks meet targets (<50ms latency)
• [ ] No regressions: pytest tests/ -v

---

10. Summary

Security Goals:

• CSWSH-Resistant: Origin validation, token auth
• Properly Authorized: Per-message permission checks
• Rate Limited: Prevent message flooding
• Validated: All messages treated as untrusted

Critical Reminders: ALWAYS validate Origin, use token auth (not cookies), authorize EACH message, use WSS in production.

WebSocket Advanced Patterns Reference

Pub/Sub Pattern

from collections import defaultdict
import asyncio

class PubSubManager:
    def __init__(self):
        self.subscriptions: dict[str, set[str]] = defaultdict(set)  # channel -> user_ids
        self.user_channels: dict[str, set[str]] = defaultdict(set)  # user_id -> channels
        self._lock = asyncio.Lock()

    async def subscribe(self, user_id: str, channel: str):
        async with self._lock:
            self.subscriptions[channel].add(user_id)
            self.user_channels[user_id].add(channel)

    async def unsubscribe(self, user_id: str, channel: str):
        async with self._lock:
            self.subscriptions[channel].discard(user_id)
            self.user_channels[user_id].discard(channel)

    async def publish(self, channel: str, message: dict):
        subscribers = self.subscriptions.get(channel, set())
        for user_id in subscribers:
            await connection_manager.send_to_user(user_id, {
                "type": "message",
                "channel": channel,
                "data": message
            })

    async def cleanup_user(self, user_id: str):
        async with self._lock:
            for channel in self.user_channels.get(user_id, set()):
                self.subscriptions[channel].discard(user_id)
            self.user_channels.pop(user_id, None)

pubsub = PubSubManager()

---

Heartbeat/Ping-Pong

import asyncio

async def websocket_with_heartbeat(websocket: WebSocket, user: User):
    """WebSocket handler with heartbeat to detect dead connections."""

    async def heartbeat():
        while True:
            await asyncio.sleep(30)
            try:
                await websocket.send_json({"type": "ping"})
            except:
                break

    heartbeat_task = asyncio.create_task(heartbeat())

    try:
        while True:
            data = await asyncio.wait_for(
                websocket.receive_json(),
                timeout=60  # Expect pong within 60s
            )

            if data.get("type") == "pong":
                continue

            await handle_message(websocket, user, data)

    except asyncio.TimeoutError:
        logger.info(f"Connection timeout: user={user.id}")
    finally:
        heartbeat_task.cancel()

---

Reconnection with State Recovery

# Server-side: Track message sequence
class StatefulConnection:
    def __init__(self, user_id: str):
        self.user_id = user_id
        self.sequence = 0
        self.pending_messages: list[dict] = []

    def next_sequence(self) -> int:
        self.sequence += 1
        return self.sequence

    def store_message(self, message: dict):
        message["seq"] = self.next_sequence()
        self.pending_messages.append(message)
        # Keep only last 100 messages
        if len(self.pending_messages) > 100:
            self.pending_messages.pop(0)

    def get_missed_messages(self, last_seq: int) -> list[dict]:
        return [m for m in self.pending_messages if m["seq"] > last_seq]

# Client-side recovery
class ReconnectingWebSocket {
    constructor(url, token) {
        this.lastSeq = 0;
    }

    connect() {
        const wsUrl = `${this.url}?token=${this.token}&last_seq=${this.lastSeq}`;
        this.ws = new WebSocket(wsUrl);

        this.ws.onmessage = (event) => {
            const data = JSON.parse(event.data);
            if (data.seq) {
                this.lastSeq = data.seq;
            }
        };
    }
}

---

Load Balancing with Redis

import aioredis

class DistributedWebSocket:
    """WebSocket handler for horizontally scaled deployment."""

    def __init__(self):
        self.redis = None
        self.local_connections: dict[str, WebSocket] = {}

    async def init(self):
        self.redis = await aioredis.from_url("redis://localhost")
        # Subscribe to broadcast channel
        self.pubsub = self.redis.pubsub()
        await self.pubsub.subscribe("ws:broadcast")

        # Start listener
        asyncio.create_task(self.listen_broadcasts())

    async def listen_broadcasts(self):
        async for message in self.pubsub.listen():
            if message["type"] == "message":
                data = json.loads(message["data"])
                user_id = data.get("user_id")
                if user_id in self.local_connections:
                    await self.local_connections[user_id].send_json(data["message"])

    async def send_to_user(self, user_id: str, message: dict):
        # Try local first
        if user_id in self.local_connections:
            await self.local_connections[user_id].send_json(message)
        else:
            # Publish to Redis for other instances
            await self.redis.publish("ws:broadcast", json.dumps({
                "user_id": user_id,
                "message": message
            }))

    async def register(self, user_id: str, websocket: WebSocket):
        self.local_connections[user_id] = websocket
        await self.redis.sadd("ws:online", user_id)

    async def unregister(self, user_id: str):
        self.local_connections.pop(user_id, None)
        await self.redis.srem("ws:online", user_id)

---

Binary Message Handling

import struct

async def handle_binary_message(websocket: WebSocket, data: bytes):
    """Handle binary WebSocket messages."""

    if len(data) < 4:
        return

    # First 4 bytes = message type
    msg_type = struct.unpack(">I", data[:4])[0]

    handlers = {
        1: handle_audio,
        2: handle_image,
        3: handle_file,
    }

    handler = handlers.get(msg_type)
    if handler:
        await handler(websocket, data[4:])

async def handle_audio(websocket: WebSocket, data: bytes):
    # Process audio chunk
    if len(data) > 1024 * 1024:  # 1MB limit
        return

    # ... process

---

Testing Utilities

import pytest
from fastapi.testclient import TestClient

class WebSocketTestClient:
    def __init__(self, app, token: str, origin: str):
        self.client = TestClient(app)
        self.token = token
        self.origin = origin

    def connect(self):
        return self.client.websocket_connect(
            f"/ws?token={self.token}",
            headers={"Origin": self.origin}
        )

@pytest.fixture
def ws_client(app, test_user_token):
    return WebSocketTestClient(
        app,
        test_user_token,
        "https://app.example.com"
    )

def test_subscribe_and_receive(ws_client):
    with ws_client.connect() as ws:
        ws.send_json({"action": "subscribe", "channel": "updates"})
        response = ws.receive_json()
        assert response["status"] == "subscribed"

WebSocket Security Examples Reference

CSWSH Prevention Complete Example

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from jose import jwt, JWTError
import logging

logger = logging.getLogger(__name__)

ALLOWED_ORIGINS = ["https://app.example.com"]
SECRET_KEY = os.environ["JWT_SECRET"]

class SecureWebSocketHandler:
    """Complete secure WebSocket implementation."""

    async def validate_origin(self, websocket: WebSocket) -> bool:
        origin = websocket.headers.get("origin")

        if not origin:
            logger.warning(f"No origin header from {websocket.client.host}")
            return False

        if origin not in ALLOWED_ORIGINS:
            logger.warning(f"Invalid origin {origin} from {websocket.client.host}")
            await websocket.close(code=4003, reason="Invalid origin")
            return False

        return True

    async def authenticate(self, websocket: WebSocket) -> User | None:
        token = websocket.query_params.get("token")

        if not token:
            await websocket.close(code=4001, reason="Token required")
            return None

        try:
            payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
            user = await user_service.get(payload["sub"])
            return user
        except JWTError as e:
            logger.warning(f"Invalid token: {e}")
            await websocket.close(code=4001, reason="Invalid token")
            return None

    async def handle(self, websocket: WebSocket):
        # 1. Validate origin
        if not await self.validate_origin(websocket):
            return

        # 2. Authenticate
        user = await self.authenticate(websocket)
        if not user:
            return

        # 3. Accept only after validation
        await websocket.accept()

        logger.info(f"WebSocket connected: user={user.id}, ip={websocket.client.host}")

        try:
            while True:
                data = await websocket.receive_json()
                await self.process_message(websocket, user, data)
        except WebSocketDisconnect:
            logger.info(f"WebSocket disconnected: user={user.id}")

    async def process_message(self, websocket, user, data):
        # Validate message format
        try:
            message = Message(**data)
        except ValueError as e:
            await websocket.send_json({"error": str(e)})
            return

        # Check permission
        if not user.has_permission(f"ws:{message.action}"):
            logger.warning(f"Permission denied: user={user.id}, action={message.action}")
            await websocket.send_json({"error": "Permission denied"})
            return

        # Process
        result = await self.handlers[message.action](user, message.data)
        await websocket.send_json(result)

handler = SecureWebSocketHandler()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await handler.handle(websocket)

---

Browser-Side Token Authentication

// Client-side: Pass token explicitly (not via cookies)
class SecureWebSocket {
    constructor(url, token) {
        this.url = url;
        this.token = token;
        this.ws = null;
    }

    connect() {
        // Pass token as query parameter
        const wsUrl = `${this.url}?token=${encodeURIComponent(this.token)}`;
        this.ws = new WebSocket(wsUrl);

        this.ws.onopen = () => {
            console.log('Connected');
        };

        this.ws.onmessage = (event) => {
            const data = JSON.parse(event.data);

            // Handle token rotation
            if (data.new_token) {
                this.token = data.new_token;
            }
        };

        this.ws.onerror = (error) => {
            console.error('WebSocket error:', error);
        };
    }

    send(action, data) {
        this.ws.send(JSON.stringify({ action, data }));
    }
}

// Usage
const token = await getAuthToken();  // From login
const ws = new SecureWebSocket('wss://api.example.com/ws', token);
ws.connect();

---

Rate Limiting Implementation

from collections import defaultdict
from time import time
import asyncio

class WebSocketRateLimiter:
    def __init__(self, max_per_minute: int = 60):
        self.max_per_minute = max_per_minute
        self.windows: dict[str, list[float]] = defaultdict(list)
        self._lock = asyncio.Lock()

    async def check(self, user_id: str) -> bool:
        async with self._lock:
            now = time()
            window = self.windows[user_id]

            # Remove expired timestamps
            window[:] = [t for t in window if t > now - 60]

            if len(window) >= self.max_per_minute:
                return False

            window.append(now)
            return True

    async def reset(self, user_id: str):
        async with self._lock:
            self.windows.pop(user_id, None)

rate_limiter = WebSocketRateLimiter(max_per_minute=60)

# Usage in handler
if not await rate_limiter.check(user.id):
    await websocket.send_json({"error": "Rate limited", "retry_after": 60})
    continue

---

Connection Manager with Security

class SecureConnectionManager:
    def __init__(self):
        self.connections: dict[str, WebSocket] = {}
        self.user_to_socket: dict[str, str] = {}
        self.ip_counts: dict[str, int] = defaultdict(int)
        self.max_per_ip = 10
        self._lock = asyncio.Lock()

    async def connect(
        self,
        websocket: WebSocket,
        user_id: str,
        conn_id: str
    ) -> bool:
        ip = websocket.client.host

        async with self._lock:
            # Check IP limit
            if self.ip_counts[ip] >= self.max_per_ip:
                return False

            # Disconnect existing connection for user (single session)
            if user_id in self.user_to_socket:
                old_id = self.user_to_socket[user_id]
                await self.disconnect(old_id, user_id, ip)

            self.connections[conn_id] = websocket
            self.user_to_socket[user_id] = conn_id
            self.ip_counts[ip] += 1

        return True

    async def disconnect(self, conn_id: str, user_id: str, ip: str):
        async with self._lock:
            if conn_id in self.connections:
                try:
                    await self.connections[conn_id].close()
                except:
                    pass
                del self.connections[conn_id]

            self.user_to_socket.pop(user_id, None)
            self.ip_counts[ip] = max(0, self.ip_counts[ip] - 1)

    async def send_to_user(self, user_id: str, message: dict):
        conn_id = self.user_to_socket.get(user_id)
        if conn_id and conn_id in self.connections:
            await self.connections[conn_id].send_json(message)

    async def broadcast(self, message: dict):
        for ws in self.connections.values():
            try:
                await ws.send_json(message)
            except:
                pass

---

Secure Message Validation

from pydantic import BaseModel, field_validator
from typing import Literal, Any

class BaseMessage(BaseModel):
    action: str
    data: dict = {}

    model_config = {"extra": "forbid"}

class SubscribeMessage(BaseModel):
    action: Literal["subscribe"]
    channel: str

    @field_validator("channel")
    @classmethod
    def validate_channel(cls, v):
        allowed = {"notifications", "updates", "chat"}
        if v not in allowed:
            raise ValueError(f"Invalid channel: {v}")
        return v

class SendMessage(BaseModel):
    action: Literal["send"]
    recipient: str
    content: str

    @field_validator("content")
    @classmethod
    def validate_content(cls, v):
        if len(v) > 1000:
            raise ValueError("Message too long")
        return v

def parse_message(data: dict) -> BaseMessage:
    """Parse and validate WebSocket message."""
    action = data.get("action")

    validators = {
        "subscribe": SubscribeMessage,
        "send": SendMessage,
    }

    validator = validators.get(action, BaseMessage)
    return validator(**data)

WebSocket Threat Model

Threat Model Overview

Domain Risk Level: HIGH

Assets to Protect

1. User Sessions - Authentication tokens, session state - Sensitivity: CRITICAL 2. Real-time Data - Messages, notifications - Sensitivity: HIGH 3. Command Execution - Server-side actions - Sensitivity: CRITICAL 4. Server Resources - Memory, connections - Sensitivity: MEDIUM

Attack Surface

• WebSocket handshake (Origin header)
• Authentication mechanism
• Message handlers
• Connection lifecycle
• Resource consumption

---

Attack Scenario 1: Cross-Site WebSocket Hijacking (CSWSH)

Threat Category: OWASP A01:2025 - Broken Access Control / CWE-346

Threat Level: CRITICAL

Attack Flow:

1. Victim is logged into target application
2. Attacker creates malicious webpage with JavaScript
3. Victim visits attacker's page
4. JavaScript creates WebSocket to target server
5. Browser sends victim's cookies automatically
6. Server accepts connection (no origin check)
7. Attacker executes commands as victim

Mitigation:

# Primary: Origin validation
async def secure_handler(websocket: WebSocket):
    origin = websocket.headers.get("origin")
    if origin not in ["https://app.example.com"]:
        await websocket.close(code=4003)
        return

# Secondary: Token authentication
token = websocket.query_params.get("token")
user = await validate_token(token)

# Tertiary: SameSite cookies
response.set_cookie("session", value, samesite="strict")

Real CVEs:

• CVE-2024-23898: Jenkins CLI WebSocket hijacking
• CVE-2024-26135: MeshCentral control hijacking
• CVE-2023-0957: Gitpod account takeover

---

Attack Scenario 2: Message Injection

Threat Category: OWASP A03:2025 - Injection / CWE-94

Threat Level: HIGH

Attack Flow:

1. Attacker establishes WebSocket connection
2. Sends malicious message: {"action": "query", "sql": "'; DROP TABLE--"}
3. Server doesn't validate message content
4. SQL injection executes

Mitigation:

from pydantic import BaseModel

class QueryMessage(BaseModel):
    table: Literal["users", "orders"]  # Allowlist
    filters: dict

async def handle_query(message: QueryMessage):
    # ORM with parameterized query
    stmt = select(models[message.table]).filter_by(**message.filters)
    return await db.execute(stmt)

---

Attack Scenario 3: Connection Flooding (DoS)

Threat Category: OWASP A10:2025 / CWE-400

Threat Level: MEDIUM

Attack Flow:

1. Attacker opens many WebSocket connections
2. Server resources exhausted
3. Legitimate users can't connect

Mitigation:

class ConnectionManager:
    def __init__(self):
        self.connections_per_ip = defaultdict(int)
        self.max_per_ip = 10

    async def connect(self, ws, ip):
        if self.connections_per_ip[ip] >= self.max_per_ip:
            await ws.close(code=4029)
            return False
        self.connections_per_ip[ip] += 1
        return True

---

Attack Scenario 4: Message Flooding

Threat Category: OWASP A10:2025 / CWE-770

Threat Level: MEDIUM

Attack Flow:

1. Attacker establishes single connection
2. Floods with thousands of messages per second
3. Server overwhelmed processing messages

Mitigation:

def check_rate_limit(user_id: str) -> bool:
    now = time()
    window = self.message_times[user_id]
    window = [t for t in window if t > now - 60]

    if len(window) >= 60:  # 60 per minute
        return False

    window.append(now)
    return True

---

Attack Scenario 5: Session Fixation

Threat Category: OWASP A07:2025 / CWE-384

Threat Level: HIGH

Attack Flow:

1. Attacker gets valid session token
2. Tricks victim into using that token
3. Victim authenticates with attacker's token
4. Attacker hijacks authenticated session

Mitigation:

async def authenticate(websocket, token):
    user = await validate_token(token)

    # Rotate token after authentication
    new_token = create_token(user)
    await websocket.send_json({"new_token": new_token})

    # Invalidate old token
    await invalidate_token(token)

    return user

---

STRIDE Analysis

---

Security Testing Checklist

CSWSH Testing

• [ ] Connection rejected from unknown origin
• [ ] Cookie-only auth is not sufficient
• [ ] Token required in query params

Authentication Testing

• [ ] Invalid token rejected
• [ ] Expired token rejected
• [ ] Token works from valid origin

Authorization Testing

• [ ] Each action requires permission check
• [ ] User cannot access others' resources

DoS Testing

• [ ] Connection limit per IP enforced
• [ ] Message rate limiting works
• [ ] Large messages rejected