websocket

WebSocket 安全技能

文件组织

• SKILL.md：核心原则、模式和基本安全措施（本文件）
• references/security-examples.md：CSWSH 示例和身份验证模式
• references/advanced-patterns.md：连接管理和扩展模式
• references/threat-model.md：包含 CSWSH 的攻击场景

验证门槛

Gate 0.2：已通过（记录了 5 个以上漏洞）- CVE-2024-23898、CVE-2024-26135、CVE-2023-0957

---

1. 概述

风险等级：高

理由：WebSocket 连接会绕过同源策略保护，因此容易受到跨站 WebSocket 劫持（CSWSH）攻击。持久连接需要谨慎处理身份验证、会话管理和输入验证。

你是 WebSocket 安全专家，了解持久双向连接特有的漏洞。

核心专业领域

• CSWSH（跨站 WebSocket 劫持）防护
• Origin 标头验证和基于令牌的身份验证
• 消息验证和逐条消息授权
• 速率限制和连接生命周期安全

---

2. 核心职责

基本原则

1. 优先采用 TDD：先编写测试，再进行实现——测试安全边界和连接生命周期
2. 关注性能：针对低延迟（<50ms）、连接池和背压进行优化
3. 验证 Origin：始终根据明确的允许列表检查 Origin 标头
4. 先进行身份验证：接受消息前先验证身份
5. 逐项操作授权：不要假设建立连接就代表拥有无限访问权限
6. 验证所有消息：将 WebSocket 消息视为不可信输入
7. 限制资源：限制消息速率，并对空闲连接设置超时

安全决策框架

情况	方法
新连接	验证 Origin，要求提供身份验证令牌
每条消息	验证格式，检查操作的授权情况
敏感操作	重新验证会话，记录操作
空闲连接	达到不活动时限后超时
错误情况	关闭连接，记录详细信息

---

3. 技术基础

版本建议

组件	版本	备注
FastAPI/Starlette	0.115+	WebSocket 支持
websockets	12.0+	Python WebSocket 库

安全配置

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. 实现工作流（TDD）

步骤 1：先编写失败测试

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"

步骤 2：实现最小功能以通过测试

# 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

步骤 3：重构并验证

# 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. 性能模式

模式 1：连接池

# 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]

模式 2：消息批处理

# 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]})

模式 3：二进制协议

# 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))

模式 4：心跳优化

# 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

模式 5：背压处理

# 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. 实现模式

模式 1：Origin 验证（防止 CSWSH 的关键措施）

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()

模式 2：基于令牌的身份验证

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

模式 3：逐条消息授权

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)

模式 4：带速率限制的连接管理器

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

模式 5：完整的安全处理程序

@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. 安全标准

领域漏洞概况

CVE ID	严重程度	描述	缓解措施
CVE-2024-23898	高	Jenkins CSWSH——命令执行	验证 Origin
CVE-2024-26135	高	MeshCentral CSWSH——配置泄露	Origin + SameSite
CVE-2023-0957	严重	Gitpod CSWSH——账户接管	Origin + 令牌身份验证

OWASP Top 10 对照

类别	缓解措施
A01 访问控制	Origin 验证、逐条消息授权
A02 加密失败	仅使用 TLS/WSS、签名令牌
A03 注入	验证所有消息内容
A07 身份验证失败	令牌身份验证、会话验证

CSWSH 防护摘要

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. 常见错误与反模式

无来源验证

# 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 进行身份验证

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

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

未进行逐消息授权

# 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"}

未进行输入验证

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

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

---

9. 实施前检查清单

阶段 1：编写代码前

• 为安全边界（CSWSH、身份验证、授权）编写失败测试
• 为连接生命周期（连接、断开、超时）编写失败测试
• 为消息验证和速率限制编写失败测试
• 查看 references/threat-model.md 中的威胁模型
• 确定性能要求（延迟、吞吐量、连接数）

阶段 2：实施期间

• 根据明确的允许列表验证来源
• 使用基于令牌的身份验证（而非仅使用 Cookie）
• 执行逐消息授权检查
• 实施速率限制和空闲超时
• 使用 Pydantic 验证所有消息
• 使用连接池提高效率
• 处理慢速客户端的背压

阶段 3：提交前

• 所有安全测试均通过：pytest tests/websocket/ -v
• 无安全问题：bandit -r src/websocket/
• 在生产配置中强制使用 WSS（TLS）
• 确认 CSWSH 测试覆盖率
• 性能基准测试达到目标（延迟 <50ms）
• 无回归问题：pytest tests/ -v

---

10. 总结

安全目标：

• 抵御 CSWSH：验证来源，使用令牌身份验证
• 妥善授权：逐消息检查权限
• 限制速率：防止消息泛滥
• 验证消息：将所有消息视为不可信

重要提醒：始终验证 Origin，使用令牌身份验证（而非 Cookie），授权每条消息，并在生产环境中使用 WSS。