10.6 WebSocket 认证与安全

FastAPI WebSocket 认证与安全入门教程 | Python新手友好

FastAPI 教程

这个教程帮助有Python基础的初学者学习如何使用FastAPI实现WebSocket的认证和安全措施，通过简单示例和日常生活类比，快速上手并建立信心。

FastAPI WebSocket 认证与安全入门教程

欢迎！如果你有基础的Python语法知识，但没有Web开发经验，这个教程正适合你。我们将一起探索FastAPI中的WebSocket功能，并学习如何添加认证和安全措施。别担心，我会用简单的生活化类比解释概念，让你快速上手并体验“成功时刻”。

什么是WebSocket？一个简单类比

想象WebSocket就像你和朋友的一个持续电话通话——不需要每次说话都重新拨号，连接一直保持，可以实时交换信息。在Web开发中，WebSocket允许服务器和客户端（比如浏览器）建立持久连接，实现实时通信，比如聊天应用或游戏。

为什么需要认证和安全？

认证（Authentication）：就像进入俱乐部前，保安要检查你的ID，确保你是受邀客人。在WebSocket中，认证验证用户的身份，防止未授权访问。
安全（Security）：类似于给通话加密，防止别人偷听。WebSocket需要措施保护数据，避免被黑客攻击。

我们将一步步构建示例，让你在几分钟内看到效果。

准备工作

确保你安装了Python（建议版本3.7+）。打开终端或命令行，运行以下命令安装FastAPI和其他必要库：

pip install fastapi uvicorn websockets

Uvicorn是一个ASGI服务器，用于运行FastAPI应用；websockets库帮助我们处理WebSocket连接。

第一步：创建第一个WebSocket连接

让我们从一个超简单的WebSocket服务器开始。创建一个文件 main.py，并输入以下代码：

from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse

app = FastAPI()

# 一个简单的HTML页面作为客户端
html = """
<!DOCTYPE html>
<html>
    <head>
        <title>WebSocket Demo</title>
    </head>
    <body>
        <h1>WebSocket 连接测试</h1>
        <button onclick="connect()">连接</button>
        <button onclick="sendMessage()">发送消息</button>
        <p id="messages"></p>
        <script>
            let ws;
            function connect() {
                ws = new WebSocket('ws://localhost:8000/ws');
                ws.onmessage = function(event) {
                    document.getElementById('messages').innerHTML += event.data + '<br>';
                };
                alert('连接成功！'); // 成功时刻：看到弹出窗口
            }
            function sendMessage() {
                if (ws) {
                    ws.send('Hello from client!');
                }
            }
        </script>
    </body>
</html>
"""

@app.get("/")
async def get():
    return HTMLResponse(html)

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()  # 接受WebSocket连接
    while True:
        data = await websocket.receive_text()  # 接收客户端消息
        await websocket.send_text(f"服务器收到: {data}")  # 发送回复

运行服务器：

uvicorn main:app --reload

打开浏览器访问 http://localhost:8000，点击“连接”按钮。你会看到弹出窗口“连接成功！”，然后点击“发送消息”，页面会显示服务器回复。恭喜！你刚刚建立了第一个WebSocket连接——这是一个重要的成功时刻！

第二步：添加基本认证

现在，让我们添加一个简单的认证机制。假设我们使用令牌（token）来验证用户。我们将修改WebSocket端点，要求在连接时提供令牌。

更新 main.py 中的WebSocket部分：

from fastapi import WebSocket, WebSocketDisconnect

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket, token: str = None):
    # 检查令牌
    if token != "mysecrettoken":  # 简单硬编码令牌，实际中应从数据库或环境变量获取
        await websocket.close(code=1008)  # 关闭连接，代码1008表示策略违规
        return
    
    await websocket.accept()
    await websocket.send_text("认证成功！欢迎加入聊天。")  # 成功时刻：收到欢迎消息
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"你说: {data}")
    except WebSocketDisconnect:
        print("客户端断开连接")

更新HTML中的JavaScript，在连接时发送令牌：

<script>
    let ws;
    function connect() {
        ws = new WebSocket('ws://localhost:8000/ws?token=mysecrettoken'); // 添加查询参数
        ws.onmessage = function(event) {
            document.getElementById('messages').innerHTML += event.data + '<br>';
        };
        alert('连接成功！');
    }
    function sendMessage() {
        if (ws) {
            ws.send('Hello from client!');
        }
    }
</script>

重启服务器，刷新页面，点击“连接”。如果令牌正确，你会看到“认证成功！欢迎加入聊天。”的消息。哇，你刚刚实现了WebSocket认证——另一个成功时刻！ 尝试更改令牌为错误值，连接会失败，体验安全机制的作用。

第三步：增强安全性

WebSocket连接可以来自任何来源，我们需要防止恶意访问。类比：只允许来自信任朋友家的电话。我们将添加来源（origin）检查。

在 main.py 中，修改WebSocket端点：

from fastapi import WebSocket, WebSocketDisconnect

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket, token: str = None):
    # 检查来源：只允许本地开发环境
    origin = websocket.headers.get("origin")
    if origin not in ["http://localhost:8000", "http://127.0.0.1:8000"]:
        await websocket.close(code=1008)
        return
    
    if token != "mysecrettoken":
        await websocket.close(code=1008)
        return
    
    await websocket.accept()
    await websocket.send_text("认证和安全检查通过！开始聊天吧。")
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"你说: {data}")
    except WebSocketDisconnect:
        print("客户端断开连接")

这个例子中，我们检查请求头中的origin字段，确保只接受来自指定域的连接。在实际部署时，你可能需要更严格的检查，比如使用HTTPS和CORS配置。

小提示：为了生产环境安全，总是使用HTTPS（wss://）而不是HTTP（ws://），并避免硬编码敏感信息如令牌。

总结

通过这个教程，你学会了：

使用FastAPI创建基本WebSocket连接。
通过令牌实现简单认证。
添加来源检查增强安全性。

每个步骤都有可运行的代码示例，让你立即看到结果，建立学习信心。记住，认证就像验证身份，安全就像保护隐私——在Web开发中至关重要。

下一步

尝试扩展这个示例：

使用数据库存储用户和令牌。
实现更复杂的认证方式，如JWT（JSON Web Tokens）。
探索FastAPI的依赖注入来管理认证逻辑。

保持实践，你会很快掌握WebSocket的更多高级功能！

如果有问题，参考FastAPI官方文档或在线社区。祝你学习愉快！

上一章 10.5 消息广播与房间概念

下一章 10.7 与前端 WebSocket 客户端通信