FastAPI – WebSocket + JWT 驗證

簡介

在即時互動的 Web 應用程式（聊天、即時通知、多人協作編輯…）中，WebSocket 是最常見的雙向通訊技術。與傳統的 HTTP 請求不同，WebSocket 允許伺服器在任何時刻推送資料給客戶端，實現低延遲的即時體驗。

然而，開放的即時通道如果缺乏適當的身份驗證，將會成為安全漏洞的入口。JWT（JSON Web Token） 作為輕量級的自包含憑證，已成為 REST API、GraphQL 以及 WebSocket 認證的事實標準。將 JWT 與 FastAPI 的 WebSocket 結合，不僅能確保連線的合法性，還能在連線期間快速取得使用者資訊，方便權限控制與訊息過濾。

本篇文章將從概念說明、實作範例、常見陷阱到最佳實踐，完整展示 FastAPI + WebSocket + JWT 的開發流程，幫助你在專案中快速上手安全的即時功能。

---

核心概念

1. WebSocket 基礎

• 雙向通訊：客戶端與伺服器在同一個 TCP 連線上同時收發訊息。
• 升級協議：瀏覽器先發送 HTTP Upgrade 請求，伺服器回應 101 Switching Protocols 後完成升級。
• FastAPI 實作：使用 WebSocket 類別與 @app.websocket 裝飾器即可建立端點。

2. JWT 作用原理

• Header、Payload、Signature 三段結構，Payload 中可放入使用者 ID、角色等資訊。
• 簽名（Signature）確保 token 未被竄改，伺服器只要持有密鑰即可驗證。
• 無狀態：伺服器不必儲存 session，只要在每次請求或連線時驗證 token 即可。

3. 為什麼要在 WebSocket 握手階段驗證 JWT？

1. 防止未授權連線：惡意使用者若直接開啟 WebSocket，會造成資源濫用。
2. 取得使用者資訊：驗證成功後即可把使用者 ID、角色等寫入 WebSocket 物件的 state，供後續訊息處理使用。
3. 統一授權邏輯：與 REST API 共用同一套 JWT 驗證程式碼，降低維護成本。

---

程式碼範例

以下範例使用 Python 3.11、FastAPI 0.109、uvicorn，以及 PyJWT 來簽發與驗證 JWT。

1️⃣ 建立 JWT 工具函式

# jwt_utils.py
import time
from typing import Any, Dict
import jwt  # pip install PyJWT

SECRET_KEY = "your-very-secret-key"
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_SECONDS = 3600  # 1 小時

def create_access_token(data: Dict[str, Any]) -> str:
    """
    產生 JWT，payload 會自動帶入過期時間 (exp)。
    """
    to_encode = data.copy()
    expire = int(time.time()) + ACCESS_TOKEN_EXPIRE_SECONDS
    to_encode.update({"exp": expire})
    token = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return token

def verify_access_token(token: str) -> Dict[str, Any]:
    """
    驗證 JWT，成功回傳 payload，失敗拋出 jwt.exceptions.* 錯誤。
    """
    payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
    return payload

說明：create_access_token 接收任意字典（通常只放 sub、role），自動加入 exp。verify_access_token 直接返回 payload，若 token 無效或過期會拋例外。

---

2️⃣ 在 HTTP 登入路由中發行 JWT

# main.py (部分)
from fastapi import FastAPI, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordRequestForm
from jwt_utils import create_access_token

app = FastAPI()

# 假設有一個簡易的使用者資料庫
FAKE_USERS_DB = {
    "alice": {"username": "alice", "password": "wonderland", "role": "admin"},
    "bob": {"username": "bob", "password": "builder", "role": "user"},
}

@app.post("/token")
async def login(form_data: OAuth2PasswordRequestForm = Depends()):
    user = FAKE_USERS_DB.get(form_data.username)
    if not user or user["password"] != form_data.password:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="帳號或密碼錯誤",
        )
    access_token = create_access_token({"sub": user["username"], "role": user["role"]})
    return {"access_token": access_token, "token_type": "bearer"}

重點：這裡使用 OAuth2PasswordRequestForm 處理 application/x-www-form-urlencoded 表單，返回的 JSON 直接符合前端常見的 Authorization: Bearer <token> 使用方式。

---

3️⃣ 在 WebSocket 握手階段驗證 JWT

# websockets.py
from fastapi import WebSocket, WebSocketDisconnect, Depends, HTTPException, status
from jwt_utils import verify_access_token
from typing import Callable

async def get_current_user(websocket: WebSocket) -> dict:
    """
    從 query string 或 header 取得 token，驗證後回傳 payload。
    """
    token = websocket.query_params.get("token")
    # 也可以改成從 Sec-WebSocket-Protocol 或 Authorization Header 取得
    if not token:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        raise HTTPException(status_code=400, detail="Missing token")
    try:
        payload = verify_access_token(token)
    except Exception as e:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        raise HTTPException(status_code=401, detail="Invalid token")
    return payload

async def websocket_endpoint(websocket: WebSocket, user: dict = Depends(get_current_user)):
    """
    連線成功後，使用者資訊會被存入 websocket.state.user，之後可直接使用。
    """
    await websocket.accept()
    websocket.state.user = user  # 把使用者資訊掛在 state 上
    try:
        while True:
            data = await websocket.receive_text()
            # 這裡示範回傳使用者名稱給客戶端
            await websocket.send_text(f"[{user['sub']}] 你說: {data}")
    except WebSocketDisconnect:
        print(f"使用者 {user['sub']} 離線")

說明：

• Depends(get_current_user) 讓 FastAPI 在 WebSocket 握手前先執行 JWT 驗證。
• 若驗證失敗，我們主動關閉連線並回傳 WS_1008_POLICY_VIOLATION（政策違規）。
• 成功後把 payload 放入 websocket.state.user，後續的訊息處理就能直接取用。

---

4️⃣ 把 WebSocket 路由掛到 FastAPI 應用

# main.py (續)
from fastapi import APIRouter
from websockets import websocket_endpoint

router = APIRouter()
router.websocket("/ws/chat")(websocket_endpoint)

app.include_router(router)

小技巧：使用 APIRouter 可以將 WebSocket 路由與 REST API 分離，讓專案結構更清晰。

---

5️⃣ 前端範例（使用 JavaScript）

<!-- index.html -->
<!DOCTYPE html>
<html lang="zh-TW">
<head>
  <meta charset="UTF-8">
  <title>FastAPI WebSocket + JWT Demo</title>
</head>
<body>
  <input id="msg" placeholder="輸入訊息...">
  <button onclick="send()">送出</button>
  <pre id="log"></pre>

  <script>
    // 1. 先向 /token 取得 JWT（此處簡化為直接寫死 token）
    const token = "YOUR_JWT_FROM_login";   // 替換成實際取得的 token

    // 2. 建立 WebSocket 連線，將 token 放在 query string
    const ws = new WebSocket(`ws://localhost:8000/ws/chat?token=${token}`);

    ws.onopen = () => log('✅ 連線成功');
    ws.onmessage = (e) => log('📨 ' + e.data);
    ws.onclose = (e) => log(`❌ 連線關閉 (${e.code})`);
    ws.onerror = (e) => log('⚠️ 錯誤', e);

    function send() {
      const msg = document.getElementById('msg').value;
      ws.send(msg);
    }

    function log(...args) {
      document.getElementById('log').textContent += args.join(' ') + '\n';
    }
  </script>
</body>
</html>

重點：

• 前端把 JWT 直接放在 query string（實務上可改為 Sec-WebSocket-Protocol 或自訂 Header），簡易示範即可。
• 若 token 無效或過期，伺服器會立即關閉連線，前端可根據 close 事件做重新登入的流程。

---

常見陷阱與最佳實踐

陷阱	可能的後果	解決方案 / 最佳實踐
把 JWT 放在 URL query string	URL 可能被瀏覽器快取、日誌或第三方服務紀錄，造成憑證洩漏。	優先使用 Sec-WebSocket-Protocol 或自訂 Header（需要前端支援）。若必須使用 query，務必使用 HTTPS/WSS 加密。
在每條訊息都重新驗證 JWT	不必要的 CPU 開銷，且可能因頻繁拋例外導致連線不穩。	只在 握手階段 驗證一次，驗證成功後把使用者資訊存入 websocket.state。
未設定 token 失效時間	永久有效的 token 成為長期安全風險。	為 JWT 設定合理的 exp（如 1 小時），並在前端實作自動刷新機制。
忘記在 WebSocketDisconnect 處理資源釋放	記憶體或資料庫連線泄漏。	在 except WebSocketDisconnect 中 關閉 DB 連線、取消背景任務。
在多節點（Cluster）環境下只依賴 JWT	若需要即時推送給特定使用者，單靠 JWT 無法定位連線。	使用 Redis Pub/Sub 或 Message Queue 來同步不同節點的訊息，並在 payload 中帶上 user_id 作為分發依據。

其他最佳實踐

1. 使用 pydantic 定義 JWT Payload，提升型別安全與自動文件化。
2. 統一錯誤回應：在 verify_access_token 拋出自訂例外，讓前端只需要處理一次 401。
3. 限制同時連線數：透過 WebSocketLimiter（如 async-limiter）防止 DoS 攻擊。
4. 定期輪替密鑰：使用 kid（Key ID）機制支援多把金鑰，同時支援金鑰過期與更換。

---

實際應用場景

場景	為什麼需要 WebSocket + JWT	可能的實作細節
即時聊天系統	每位使用者必須先登入，且訊息僅能發送給已授權的聊天室成員。	JWT 中記錄 sub 與 room_id，WebSocket 連線時驗證後加入對應的 Room Group（可用 dict[room_id] = set[WebSocket]）
即時通知 (Push)	系統需要向特定使用者推送重要訊息，如訂單狀態變更。	後端在資料庫觸發事件時，根據 JWT sub 查找對應的 WebSocket，使用 await ws.send_json(...)
多人協作編輯	編輯文件時必須確保每位使用者都有編輯權限，且變更要即時同步。	JWT 包含 role 與 document_id，伺服器在收到編輯指令前檢查權限，再廣播給同文件的其他連線
遊戲即時對戰	玩家必須驗證身份，避免偽造玩家資料。	JWT 中加入 player_id、match_id，伺服器根據 match_id 建立對戰房間，僅允許同房間玩家互相通訊

---

總結

• WebSocket 為即時雙向通訊提供了高效的基礎，而 JWT 則是現代 Web 應用最常用的無狀態驗證機制。
• 在 FastAPI 中，我們只需要在 WebSocket 握手階段使用 Depends 呼叫 JWT 驗證函式，即可把使用者資訊安全地掛在 websocket.state，後續的訊息處理就能直接使用。
• 實作時要注意 憑證傳遞方式（優先 Header 或 Sec-WebSocket-Protocol）、Token 失效時間、以及 多節點同步 等安全與可擴展性議題。
• 透過本文提供的 完整範例，你可以快速在自己的專案中加入 安全的即時功能，無論是聊天、通知、或是協作編輯，都能在保護使用者身份的前提下提供流暢的使用體驗。

祝開發順利，玩得開心！ 🎉