WebSocket 学习教程

6. Java WebSocket 整合

6.1 Java WebSocket API (JSR 356)

Java API for WebSocket (JSR 356) 是 Java EE 7 引入的标准规范，提供了在 Java 中创建 WebSocket 应用的官方 API。

6.1.1 创建 WebSocket 服务端

使用 Java WebSocket API 创建一个 WebSocket 端点非常简单，只需添加几个注解即可：

package com.example.websocket;

import java.io.IOException;
import java.util.Set;
import java.util.concurrent.CopyOnWriteArraySet;

import javax.websocket.OnClose;
import javax.websocket.OnError;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.ServerEndpoint;

@ServerEndpoint("/chat")
public class ChatEndpoint {
    
    private static Set<Session> sessions = new CopyOnWriteArraySet<>();
    
    @OnOpen
    public void onOpen(Session session) {
        sessions.add(session);
        System.out.println("新连接建立: " + session.getId());
        broadcastMessage("系统消息：用户 " + session.getId() + " 加入了聊天室");
    }
    
    @OnMessage
    public void onMessage(String message, Session session) {
        System.out.println("收到消息: " + message + " 来自 " + session.getId());
        broadcastMessage("用户 " + session.getId() + ": " + message);
    }
    
    @OnClose
    public void onClose(Session session) {
        sessions.remove(session);
        System.out.println("连接关闭: " + session.getId());
        broadcastMessage("系统消息：用户 " + session.getId() + " 离开了聊天室");
    }
    
    @OnError
    public void onError(Throwable error, Session session) {
        System.err.println("错误发生在会话 " + session.getId() + ": " + error.getMessage());
    }
    
    private void broadcastMessage(String message) {
        for (Session session : sessions) {
            try {
                session.getBasicRemote().sendText(message);
            } catch (IOException e) {
                System.err.println("发送消息失败: " + e.getMessage());
            }
        }
    }
}

6.1.2 服务端配置

在 Java EE 环境中，WebSocket 端点会被自动扫描并注册，无需额外配置。但在 Servlet 容器（如 Tomcat）中，需要确保在 web.xml 或通过注解配置 WebSocket 支持：

// 使用编程方式注册端点
import javax.websocket.server.ServerContainer;
import javax.websocket.server.ServerEndpointConfig;
import javax.servlet.ServletContextListener;
import javax.servlet.ServletContextEvent;
import javax.servlet.annotation.WebListener;

@WebListener
public class WebSocketInitializer implements ServletContextListener {
    
    @Override
    public void contextInitialized(ServletContextEvent sce) {
        ServerContainer serverContainer = (ServerContainer) sce.getServletContext()
                .getAttribute("javax.websocket.server.ServerContainer");
        try {
            serverContainer.addEndpoint(ChatEndpoint.class);
            // 或使用配置类
            serverContainer.addEndpoint(ServerEndpointConfig.Builder
                    .create(AdvancedChatEndpoint.class, "/advancedChat")
                    .configurator(new CustomConfigurator())
                    .build());
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
    
    @Override
    public void contextDestroyed(ServletContextEvent sce) {
        // 清理资源
    }
}

6.1.3 WebSocket 客户端 API

Java WebSocket API 也支持创建客户端应用：

package com.example.websocket.client;

import java.net.URI;
import javax.websocket.ClientEndpoint;
import javax.websocket.ContainerProvider;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.WebSocketContainer;

@ClientEndpoint
public class WebSocketClient {
    
    private Session session;
    
    public WebSocketClient(URI endpointURI) {
        try {
            WebSocketContainer container = ContainerProvider.getWebSocketContainer();
            container.connectToServer(this, endpointURI);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
    
    @OnOpen
    public void onOpen(Session session) {
        this.session = session;
        System.out.println("已连接到服务器");
    }
    
    @OnMessage
    public void onMessage(String message) {
        System.out.println("收到消息: " + message);
    }
    
    public void sendMessage(String message) {
        try {
            session.getBasicRemote().sendText(message);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
    
    // 使用示例
    public static void main(String[] args) {
        try {
            WebSocketClient client = new WebSocketClient(new URI("ws://localhost:8080/chat"));
            // 等待连接建立
            Thread.sleep(1000);
            // 发送消息
            client.sendMessage("Hello from Java client!");
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}

6.2 Spring WebSocket

Spring 框架提供了对 WebSocket 的强大支持，尤其是与 STOMP 子协议的集成，使得构建复杂的 WebSocket 应用更加简单。

6.2.1 Spring WebSocket 配置

在 Spring Boot 应用中配置 WebSocket：

package com.example.websocket.config;

import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {

    @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {
        // 启用简单的内存消息代理
        registry.enableSimpleBroker("/topic", "/queue");
        // 设置应用前缀
        registry.setApplicationDestinationPrefixes("/app");
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        // 注册 STOMP 端点
        registry.addEndpoint("/ws")
                .setAllowedOrigins("*")
                .withSockJS(); // 启用 SockJS 回退
    }
}

6.2.2 控制器和消息处理

使用 Spring 的消息映射机制处理 WebSocket 消息：

package com.example.websocket.controller;

import com.example.websocket.model.ChatMessage;
import com.example.websocket.model.OutputMessage;

import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.messaging.handler.annotation.SendTo;
import org.springframework.stereotype.Controller;

import java.text.SimpleDateFormat;
import java.util.Date;

@Controller
public class ChatController {

    @MessageMapping("/chat")
    @SendTo("/topic/messages")
    public OutputMessage send(ChatMessage message) {
        String time = new SimpleDateFormat("HH:mm:ss").format(new Date());
        return new OutputMessage(message.getFrom(), message.getText(), time);
    }
}

// 消息模型
package com.example.websocket.model;

public class ChatMessage {
    private String from;
    private String text;

    // getters and setters
    public String getFrom() {
        return from;
    }

    public void setFrom(String from) {
        this.from = from;
    }

    public String getText() {
        return text;
    }

    public void setText(String text) {
        this.text = text;
    }
}

public class OutputMessage {
    private String from;
    private String text;
    private String time;

    public OutputMessage(String from, String text, String time) {
        this.from = from;
        this.text = text;
        this.time = time;
    }

    // getters and setters
    public String getFrom() {
        return from;
    }

    public void setFrom(String from) {
        this.from = from;
    }

    public String getText() {
        return text;
    }

    public void setText(String text) {
        this.text = text;
    }

    public String getTime() {
        return time;
    }

    public void setTime(String time) {
        this.time = time;
    }
}

6.2.3 用户认证与消息拦截

Spring Security 与 WebSocket 集成，实现认证和消息拦截：

package com.example.websocket.config;

import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.simp.SimpMessageType;
import org.springframework.security.config.annotation.web.messaging.MessageSecurityMetadataSourceRegistry;
import org.springframework.security.config.annotation.web.socket.AbstractSecurityWebSocketMessageBrokerConfigurer;

@Configuration
public class WebSocketSecurityConfig extends AbstractSecurityWebSocketMessageBrokerConfigurer {

    @Override
    protected void configureInbound(MessageSecurityMetadataSourceRegistry messages) {
        messages
            .simpTypeMatchers(SimpMessageType.CONNECT, SimpMessageType.HEARTBEAT, SimpMessageType.DISCONNECT).permitAll()
            .simpDestMatchers("/app/**").hasRole("USER")
            .simpSubscribeDestMatchers("/topic/**", "/queue/**").hasRole("USER")
            .anyMessage().denyAll();
    }

    @Override
    protected boolean sameOriginDisabled() {
        // 允许跨域
        return true;
    }
}

// 自定义通道拦截器
package com.example.websocket.interceptor;

import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.stereotype.Component;

@Component
public class WebSocketInterceptor implements ChannelInterceptor {

    @Override
    public Message preSend(Message message, MessageChannel channel) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
        
        if (StompCommand.CONNECT.equals(accessor.getCommand())) {
            // 获取用户认证信息或其他自定义逻辑
            String authToken = accessor.getFirstNativeHeader("X-Auth-Token");
            if (authToken != null) {
                // 验证令牌并设置用户信息
                // accessor.setUser(...)
            }
        }
        
        return message;
    }
}

6.3 Spring WebFlux WebSocket

Spring WebFlux 提供了响应式 WebSocket 支持，适合构建高并发的 WebSocket 应用：

package com.example.websocket.reactive;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.reactive.HandlerMapping;
import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.server.support.WebSocketHandlerAdapter;

import java.util.HashMap;
import java.util.Map;

import reactor.core.publisher.Flux;

@Configuration
public class ReactiveWebSocketConfig {

    @Bean
    public HandlerMapping webSocketHandlerMapping() {
        Map map = new HashMap<>();
        map.put("/reactive-socket", new ReactiveWebSocketHandler());

        SimpleUrlHandlerMapping handlerMapping = new SimpleUrlHandlerMapping();
        handlerMapping.setOrder(1);
        handlerMapping.setUrlMap(map);
        return handlerMapping;
    }

    @Bean
    public WebSocketHandlerAdapter handlerAdapter() {
        return new WebSocketHandlerAdapter();
    }
}

class ReactiveWebSocketHandler implements WebSocketHandler {

    @Override
    public reactor.core.publisher.Mono handle(org.springframework.web.reactive.socket.WebSocketSession session) {
        // 处理传入消息
        Flux output = session.receive()
                .map(message -> message.getPayloadAsText())
                .map(text -> "响应: " + text)
                .map(session::textMessage);

        return session.send(output);
    }
}

6.4 实际应用示例：实时股票交易平台

下面是一个使用 Spring Boot 和 WebSocket 构建的股票交易平台后端示例：

6.4.1 项目结构

src/main/java/com/example/stocktrading/
├── StockTradingApplication.java
├── config/
│   ├── WebSocketConfig.java
│   └── SchedulerConfig.java
├── controller/
│   └── StockController.java
├── model/
│   ├── StockQuote.java
│   └── TradeRequest.java
├── service/
│   ├── StockService.java
│   └── StockQuoteGenerator.java
└── exception/
    └── TradingException.java

6.4.2 WebSocket 配置

package com.example.stocktrading.config;

import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {

    @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {
        registry.enableSimpleBroker("/topic");
        registry.setApplicationDestinationPrefixes("/app");
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/stock-trading")
               .setAllowedOrigins("*")
               .withSockJS();
    }
}

6.4.3 定时推送股票行情

package com.example.stocktrading.config;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.simp.SimpMessagingTemplate;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.scheduling.annotation.Scheduled;

import com.example.stocktrading.model.StockQuote;
import com.example.stocktrading.service.StockQuoteGenerator;

import java.util.List;

@Configuration
@EnableScheduling
public class SchedulerConfig {

    @Autowired
    private SimpMessagingTemplate messagingTemplate;
    
    @Autowired
    private StockQuoteGenerator quoteGenerator;
    
    @Scheduled(fixedRate = 2000) // 每2秒更新一次
    public void sendStockQuotes() {
        List quotes = quoteGenerator.generateRandomQuotes();
        messagingTemplate.convertAndSend("/topic/stock-quotes", quotes);
    }
}

6.4.4 控制器处理交易请求

package com.example.stocktrading.controller;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.messaging.handler.annotation.SendTo;
import org.springframework.messaging.simp.SimpMessagingTemplate;
import org.springframework.stereotype.Controller;

import com.example.stocktrading.exception.TradingException;
import com.example.stocktrading.model.TradeRequest;
import com.example.stocktrading.model.TradeResponse;
import com.example.stocktrading.service.StockService;

@Controller
public class StockController {

    @Autowired
    private StockService stockService;
    
    @Autowired
    private SimpMessagingTemplate messagingTemplate;
    
    @MessageMapping("/trade")
    public void executeTrade(TradeRequest request) {
        try {
            TradeResponse response = stockService.processTrade(request);
            // 向特定用户发送交易确认
            messagingTemplate.convertAndSendToUser(
                request.getUserId(), 
                "/queue/trade-confirmation", 
                response
            );
            
            // 广播交易更新（不含敏感信息）
            messagingTemplate.convertAndSend(
                "/topic/market-activity", 
                stockService.getMarketUpdate(request.getSymbol())
            );
        } catch (TradingException e) {
            // 发送错误响应
            messagingTemplate.convertAndSendToUser(
                request.getUserId(), 
                "/queue/errors", 
                e.getMessage()
            );
        }
    }
    
    @MessageMapping("/subscribe-symbol")
    @SendTo("/topic/symbols/")
    public void subscribeToSymbol(String symbol) {
        // 用户订阅特定股票的更新
        // 可以在这里记录用户的订阅以便定向推送
        return;
    }
}

6.4.5 数据模型

package com.example.stocktrading.model;

import java.math.BigDecimal;
import java.time.LocalDateTime;

public class StockQuote {
    private String symbol;
    private BigDecimal price;
    private BigDecimal change;
    private BigDecimal changePercent;
    private long volume;
    private LocalDateTime timestamp;
    
    // 构造函数、getter和setter
}

public class TradeRequest {
    private String userId;
    private String symbol;
    private int quantity;
    private String orderType; // "BUY" or "SELL"
    private BigDecimal limitPrice; // 可选的限价
    
    // 构造函数、getter和setter
}

public class TradeResponse {
    private String tradeId;
    private String status;
    private BigDecimal executedPrice;
    private LocalDateTime executionTime;
    private String message;
    
    // 构造函数、getter和setter
}

7. Python中的WebSocket实现

Python拥有多个强大的WebSocket库，可以轻松实现客户端和服务器端WebSocket应用。接下来我们将介绍几个常用的Python WebSocket库及其使用方法。

7.1 websockets库

websockets是一个流行的Python库，提供了异步WebSocket客户端和服务器端实现，基于Python的asyncio库。

7.1.1 安装

pip install websockets

7.1.2 服务器端实现

以下是一个简单的WebSocket服务器实现：

#!/usr/bin/env python

import asyncio
import websockets
import json
from datetime import datetime

# 存储所有连接的客户端
connected_clients = set()

async def chat(websocket, path):
    # 注册客户端
    connected_clients.add(websocket)
    
    try:
        # 发送欢迎消息
        await websocket.send(json.dumps({
            "type": "system",
            "message": "欢迎加入聊天室！",
            "timestamp": datetime.now().strftime("%H:%M:%S")
        }))
        
        # 通知其他用户有新用户加入
        if len(connected_clients) > 1:
            await notify_others(websocket, "system", "新用户加入了聊天室")
        
        # 接收消息并广播
        async for message in websocket:
            try:
                data = json.loads(message)
                # 添加时间戳
                data["timestamp"] = datetime.now().strftime("%H:%M:%S")
                
                # 广播给所有客户端
                await broadcast(json.dumps(data))
            except json.JSONDecodeError:
                # 处理非JSON格式的消息
                await websocket.send(json.dumps({
                    "type": "error",
                    "message": "消息格式错误，请发送JSON格式",
                    "timestamp": datetime.now().strftime("%H:%M:%S")
                }))
                
    except websockets.exceptions.ConnectionClosed:
        print("客户端连接断开")
    finally:
        # 客户端断开连接时，从集合中移除
        connected_clients.remove(websocket)
        # 通知其他用户有用户离开
        if connected_clients:
            await notify_others(websocket, "system", "有用户离开了聊天室")

async def broadcast(message):
    """向所有连接的客户端广播消息"""
    if connected_clients:
        await asyncio.gather(
            *[client.send(message) for client in connected_clients]
        )

async def notify_others(sender, msg_type, message):
    """通知除发送者外的所有客户端"""
    others = connected_clients - {sender}
    if others:
        notification = json.dumps({
            "type": msg_type,
            "message": message,
            "timestamp": datetime.now().strftime("%H:%M:%S")
        })
        await asyncio.gather(*[client.send(notification) for client in others])

# 启动WebSocket服务器
start_server = websockets.serve(chat, "localhost", 8765)

print("WebSocket服务器启动在 ws://localhost:8765")

# 启动事件循环
asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()

7.1.3 客户端实现

以下是一个基于asyncio的WebSocket客户端实现：

#!/usr/bin/env python

import asyncio
import websockets
import json
import aioconsole

async def receive_messages(websocket):
    """接收并显示服务器发送的消息"""
    try:
        async for message in websocket:
            data = json.loads(message)
            
            if data["type"] == "system":
                print(f"\n[系统消息 {data['timestamp']}] {data['message']}")
            elif data["type"] == "chat":
                print(f"\n[{data['username']} {data['timestamp']}] {data['message']}")
            elif data["type"] == "error":
                print(f"\n[错误 {data['timestamp']}] {data['message']}")
    except websockets.exceptions.ConnectionClosed:
        print("\n连接已关闭")

async def send_messages(websocket, username):
    """从控制台读取消息并发送到服务器"""
    try:
        while True:
            message = await aioconsole.ainput("") 
            
            if message.lower() == "/quit":
                break
                
            await websocket.send(json.dumps({
                "type": "chat",
                "username": username,
                "message": message
            }))
    except websockets.exceptions.ConnectionClosed:
        pass

async def chat_client():
    # 获取用户名
    username = input("请输入您的用户名: ")
    
    # 连接到WebSocket服务器
    uri = "ws://localhost:8765"
    print(f"正在连接到 {uri}...")
    
    try:
        async with websockets.connect(uri) as websocket:
            print("连接成功！输入消息开始聊天，输入 /quit 退出。")
            
            # 创建两个任务：一个接收消息，一个发送消息
            receive_task = asyncio.create_task(receive_messages(websocket))
            send_task = asyncio.create_task(send_messages(websocket, username))
            
            # 等待任一任务完成
            done, pending = await asyncio.wait(
                [receive_task, send_task],
                return_when=asyncio.FIRST_COMPLETED
            )
            
            # 取消未完成的任务
            for task in pending:
                task.cancel()
                
    except (websockets.exceptions.InvalidStatusCode, ConnectionRefusedError) as e:
        print(f"连接错误: {e}")
    except Exception as e:
        print(f"发生未知错误: {e}")

# 运行客户端
asyncio.run(chat_client())

7.2 FastAPI与WebSockets

FastAPI是一个现代、高性能的Python Web框架，它也提供了内置的WebSocket支持。

7.2.1 安装

pip install fastapi uvicorn

7.2.2 实现聊天服务器

下面是一个基于FastAPI的WebSocket聊天服务器实现：

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from fastapi.middleware.cors import CORSMiddleware
from typing import List, Dict, Optional
import json
from datetime import datetime, timedelta
from pydantic import BaseModel
import jwt
from passlib.context import CryptContext

# 创建FastAPI应用
app = FastAPI(title="FastAPI WebSocket聊天")

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 在生产环境中应该限制来源
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 用于储存用户数据（实际应用中应使用数据库）
fake_users_db = {
    "admin": {
        "username": "admin",
        "hashed_password": "$2b$12$EixZaYVK1fsbw1ZfbX3OXePaWxn96p36WQoeG6Lruj3vjPGga31lW",  # "password"
        "disabled": False,
    }
}

# 用于JWT认证的密钥
SECRET_KEY = "09d25e094faa6ca2556c818166b7a9563b93f7099f6f0f4caa6cf63b88e8d3e7"
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30

# 密码加密上下文
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
# OAuth2 密码流程
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

# 用户模型
class User(BaseModel):
    username: str
    disabled: Optional[bool] = None

class UserInDB(User):
    hashed_password: str

class Token(BaseModel):
    access_token: str
    token_type: str

class TokenData(BaseModel):
    username: Optional[str] = None

# WebSocket连接管理器
class ConnectionManager:
    def __init__(self):
        # 存储活跃的WebSocket连接
        self.active_connections: Dict[str, WebSocket] = {}

    async def connect(self, websocket: WebSocket, user_id: str):
        await websocket.accept()
        self.active_connections[user_id] = websocket
        
        # 系统消息 - 用户加入
        await self.broadcast(
            json.dumps({
                "type": "system",
                "content": f"{user_id} 加入了聊天",
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            })
        )

    def disconnect(self, user_id: str):
        if user_id in self.active_connections:
            del self.active_connections[user_id]

    async def send_personal_message(self, message: str, user_id: str):
        if user_id in self.active_connections:
            await self.active_connections[user_id].send_text(message)

    async def broadcast(self, message: str):
        # 向所有连接的客户端广播消息
        for connection in self.active_connections.values():
            await connection.send_text(message)

# 创建连接管理器实例
manager = ConnectionManager()

# 辅助函数
def verify_password(plain_password, hashed_password):
    return pwd_context.verify(plain_password, hashed_password)

def get_password_hash(password):
    return pwd_context.hash(password)

def get_user(db, username: str):
    if username in db:
        user_dict = db[username]
        return UserInDB(**user_dict)

def authenticate_user(fake_db, username: str, password: str):
    user = get_user(fake_db, username)
    if not user:
        return False
    if not verify_password(password, user.hashed_password):
        return False
    return user

def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=15)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt

async def get_current_user(token: str = Depends(oauth2_scheme)):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
        token_data = TokenData(username=username)
    except jwt.PyJWTError:
        raise credentials_exception
    user = get_user(fake_users_db, username=token_data.username)
    if user is None:
        raise credentials_exception
    return user

# 路由
@app.post("/token", response_model=Token)
async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):
    # 认证用户
    user = authenticate_user(fake_users_db, form_data.username, form_data.password)
    if not user:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username or password",
            headers={"WWW-Authenticate": "Bearer"},
        )
    
    # 创建访问令牌
    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": user.username}, expires_delta=access_token_expires
    )
    return {"access_token": access_token, "token_type": "bearer"}

@app.websocket("/ws/{token}")
async def websocket_endpoint(websocket: WebSocket, token: str):
    try:
        # 验证令牌并获取用户
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username = payload.get("sub")
        if username is None or username not in fake_users_db:
            await websocket.close(code=1008, reason="Invalid token")
            return
            
        # 接受WebSocket连接
        await manager.connect(websocket, username)
        
        try:
            # 监听WebSocket消息
            while True:
                data = await websocket.receive_text()
                
                # 解析消息数据
                try:
                    message_data = json.loads(data)
                    # 添加发送者和时间戳
                    message_data["sender"] = username
                    message_data["timestamp"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    
                    # 广播消息给所有用户
                    await manager.broadcast(json.dumps(message_data))
                    
                except json.JSONDecodeError:
                    # 发送错误消息给用户
                    await manager.send_personal_message(
                        json.dumps({
                            "type": "error",
                            "content": "Invalid JSON format",
                            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                        }),
                        username
                    )
                    
        except WebSocketDisconnect:
            # 用户断开连接
            manager.disconnect(username)
            # 通知其他用户
            await manager.broadcast(
                json.dumps({
                    "type": "system",
                    "content": f"{username} 离开了聊天",
                    "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                })
            )
    except jwt.PyJWTError:
        await websocket.close(code=1008, reason="Invalid token")

# 前端页面路由（实际应用中应该使用单独的前端）
@app.get("/")
async def get():
    return {"message": "WebSocket聊天服务器正在运行"}

# 运行应用（使用命令：uvicorn app:app --reload）
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

7.2.3 HTML客户端页面

下面是一个配合FastAPI WebSocket服务器使用的简单HTML客户端：

<!DOCTYPE html>
<html lang="zh">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>WebSocket聊天室</title>
    <style>
        body {
            font-family: Arial, sans-serif;
            max-width: 800px;
            margin: 0 auto;
            padding: 20px;
        }
        #login-form, #chat-container {
            border: 1px solid #ccc;
            padding: 20px;
            border-radius: 5px;
            margin-bottom: 20px;
        }
        #messages {
            height: 300px;
            overflow-y: auto;
            border: 1px solid #eee;
            padding: 10px;
            margin-bottom: 10px;
            border-radius: 5px;
        }
        .message {
            margin-bottom: 10px;
            padding: 8px;
            border-radius: 5px;
        }
        .system {
            background-color: #f0f0f0;
            color: #555;
            font-style: italic;
        }
        .user {
            background-color: #e1f5fe;
        }
        .self {
            background-color: #e8f5e9;
            text-align: right;
        }
        .error {
            background-color: #ffebee;
            color: #c62828;
        }
        .hidden {
            display: none;
        }
        input[type="text"], input[type="password"], button {
            padding: 8px;
            margin-bottom: 10px;
        }
        #message-input {
            width: 80%;
            margin-right: 10px;
        }
        .timestamp {
            font-size: 0.8em;
            color: #999;
            display: block;
        }
        .sender {
            font-weight: bold;
        }
    </style>
</head>
<body>
    <h1>WebSocket聊天室</h1>
    
    <div id="login-form">
        <h2>登录</h2>
        <div>
            <label for="username">用户名：</label>
            <input type="text" id="username" placeholder="输入用户名" required>
        </div>
        <div>
            <label for="password">密码：</label>
            <input type="password" id="password" placeholder="输入密码" required>
        </div>
        <button id="login-btn">登录</button>
        <p id="login-error" style="color: red; display: none;"></p>
    </div>
    
    <div id="chat-container" class="hidden">
        <h2>聊天室</h2>
        <div id="messages"></div>
        <div id="input-container">
            <input type="text" id="message-input" placeholder="输入消息...">
            <button id="send-btn">发送</button>
        </div>
    </div>
    
    <script>
        // DOM元素
        const loginForm = document.getElementById('login-form');
        const loginError = document.getElementById('login-error');
        const chatContainer = document.getElementById('chat-container');
        const messagesDiv = document.getElementById('messages');
        const messageInput = document.getElementById('message-input');
        const sendBtn = document.getElementById('send-btn');
        const loginBtn = document.getElementById('login-btn');
        
        let socket = null;
        let token = null;
        let username = '';
        
        // 登录处理
        loginBtn.addEventListener('click', async () => {
            username = document.getElementById('username').value;
            const password = document.getElementById('password').value;
            
            if (!username || !password) {
                showLoginError('用户名和密码不能为空');
                return;
            }
            
            try {
                // 发送登录请求
                const response = await fetch('/token', {
                    method: 'POST',
                    headers: {
                        'Content-Type': 'application/x-www-form-urlencoded',
                    },
                    body: `username=${encodeURIComponent(username)}&password=${encodeURIComponent(password)}`,
                });
                
                const data = await response.json();
                
                if (response.ok) {
                    token = data.access_token;
                    // 登录成功，切换到聊天界面
                    loginForm.classList.add('hidden');
                    chatContainer.classList.remove('hidden');
                    // 连接WebSocket
                    connectWebSocket();
                } else {
                    showLoginError(data.detail || '登录失败');
                }
            } catch (error) {
                showLoginError('连接服务器失败');
                console.error('Login error:', error);
            }
        });
        
        // 显示登录错误
        function showLoginError(message) {
            loginError.textContent = message;
            loginError.style.display = 'block';
        }
        
        // 连接WebSocket
        function connectWebSocket() {
            // 使用JWT token建立WebSocket连接
            socket = new WebSocket(`ws://${window.location.host}/ws/${token}`);
            
            // 连接打开
            socket.onopen = function(event) {
                addMessage('system', '已连接到聊天服务器');
            };
            
            // 收到消息
            socket.onmessage = function(event) {
                const message = JSON.parse(event.data);
                
                if (message.type === 'system') {
                    addMessage('system', message.content, message.timestamp);
                } else if (message.type === 'chat') {
                    const isSelf = message.sender === username;
                    addMessage(isSelf ? 'self' : 'user', message.content, message.timestamp, message.sender);
                } else if (message.type === 'error') {
                    addMessage('error', message.content, message.timestamp);
                }
            };
            
            // 连接关闭
            socket.onclose = function(event) {
                addMessage('system', '与服务器的连接已关闭');
                // 尝试重新连接
                setTimeout(() => {
                    if (token) {
                        connectWebSocket();
                    }
                }, 5000);
            };
            
            // 连接错误
            socket.onerror = function(error) {
                addMessage('error', '连接错误');
                console.error('WebSocket error:', error);
            };
        }
        
        // 发送消息
        sendBtn.addEventListener('click', sendMessage);
        messageInput.addEventListener('keypress', function(e) {
            if (e.key === 'Enter') {
                sendMessage();
            }
        });
        
        function sendMessage() {
            const message = messageInput.value.trim();
            if (message && socket && socket.readyState === WebSocket.OPEN) {
                // 发送聊天消息
                socket.send(JSON.stringify({
                    type: 'chat',
                    content: message
                }));
                
                // 清空输入框
                messageInput.value = '';
            }
        }
        
        // 向聊天窗口添加消息
        function addMessage(type, content, timestamp = null, sender = null) {
            const messageElem = document.createElement('div');
            messageElem.className = `message ${type}`;
            
            let messageContent = '';
            
            if (sender && type !== 'self') {
                messageContent += `<span class="sender">${sender}</span>: `;
            }
            
            messageContent += content;
            
            if (timestamp) {
                messageContent += `<span class="timestamp">${timestamp}</span>`;
            }
            
            messageElem.innerHTML = messageContent;
            messagesDiv.appendChild(messageElem);
            
            // 滚动到底部
            messagesDiv.scrollTop = messagesDiv.scrollHeight;
        }
    </script>
</body>
</html>

7.3 其他Python WebSocket库

除了上述库外，Python还有其他几个WebSocket库值得关注：

• aiohttp - 提供异步WebSocket客户端和服务器支持
• Tornado - 包含WebSocket支持的异步网络库
• Django Channels - 为Django添加WebSocket支持
• Flask-SocketIO - 将Socket.IO集成到Flask应用中
• pywebsocket - 适用于WSGI服务器的WebSocket实现

7.4 选择合适的Python WebSocket库的考虑因素

• 现有技术栈：如果您已经使用Django或Flask，选择与其集成的WebSocket库可能更容易
• 性能需求：对于高性能应用，基于asyncio的库通常是更好的选择
• 跨平台兼容性：如果需要支持旧浏览器，考虑提供回退机制的库
• 扩展性：评估是否需要自定义WebSocket协议或其他功能
• 学习曲线：一些库比其他库更容易上手