websocket学习手册及python实现简单的聊天室

概述

WebSocket 是一种网络通信协议，允许在单个 TCP 连接上进行全双工通信。它最核心的优势就在于实现了持久连接，实现了实时的数据传输。
HTTP 协议有一个很大的缺点，通信只能由客户端发起，服务器返回响应后连接就会关闭，如果我们想要知道任务连续的状态变化的话，就需要通过轮询来获取状态，这种方法的效率就会非常低。
WebSocket需要维护连接，这需要额外的开销，包括内存和CPU，没有必要需求无需使用
WebSocket协议是一种可靠的、高效的、双向的、持久的通信协议，支持文本和二进制数据。

如何建立连接

+---------+                         +---------+
| Client  |                         | Server  |
+---------+                         +---------+|                                   || 1. HTTP GET (Upgrade request)     ||---------------------------------->||                                   || 2. HTTP 101 Switching Protocols   ||<----------------------------------||                                   || 3. WebSocket Connection Established ||<---------------------------------->||                                   |

WebSocket 连接的建立过程基于 HTTP 协议，我使用postman建立一个连接，打开握手详情，可以看到这些信息：

在这里插入图片描述

首先，客户端发送一个 HTTP GET 请求，包含以下关键头信息：
- Upgrade: websocket：表示希望升级到 WebSocket 协议。
- Connection: Upgrade：表示要求升级连接。
- Sec-WebSocket-Key：一个随机生成的 Base64 编码字符串，用于验证服务器。
- Sec-WebSocket-Version：指定 WebSocket 协议版本（通常为 13）。
服务器检查请求头，如果支持 WebSocket，返回 HTTP 101 状态码（Switching Protocols），并包含以下头信息：
- Upgrade: websocket：确认升级到 WebSocket 协议。
- Connection: Upgrade：确认连接升级。
- Sec-WebSocket-Accept：基于客户端的 Sec-WebSocket-Key 计算的值，用于验证。

关于wss和ws

ws 和 wss 是 WebSocket 协议的两种不同形式，主要区别在于 安全性 和 数据传输方式。

	`ws`	`wss`
特性	非加密	加密
协议	普通 TCP	基于 TLS/SSL 的 TCP
安全性	数据明文传输，不安全	数据加密传输，安全
适用场景	内部网络、开发环境	生产环境、敏感数据传输
默认端口	80	443
URL 前缀	`ws://`	`wss://`

简单实现

from fastapi import FastAPI, APIRouter, WebSocket, WebSocketDisconnectapp = FastAPI()
router = APIRouter(prefix='/demo', tags=['demo'])# WebSocket 端点
@router.websocket("/ws/{user_id}")
async def websocket_endpoint(websocket: WebSocket, user_id: str):# 接受客户端连接await websocket.accept()while True:try:# 接收客户端发送的消息data = await websocket.receive_text()# 打印接收到的消息print(f"Received message: {data}")# 将消息原样发送回客户端await websocket.send_text(f"Echo: {data}")except Exception as e:print(f"Error: {e}")breakapp.include_router(router)

实战-实现一个聊天室

import jsonfrom fastapi import FastAPI, APIRouter, WebSocket, WebSocketDisconnectapp = FastAPI()
router = APIRouter(prefix='/demo', tags=['demo'])class ConnectionManager:def __init__(self):self.active_connections: dict[str, WebSocket] = {}async def connect(self, websocket: WebSocket, user_id: str):await websocket.accept()self.active_connections[user_id] = websocketdef disconnect(self, user_id: str):if user_id in self.active_connections:del self.active_connections[user_id]async def send_message(self, message: str, user_id: str, target_user: str):# 向目标用户发送消息if websocket := self.active_connections.get(target_user):data = json.dumps({"message": message, "from_user": user_id})await websocket.send_text(data)manager = ConnectionManager()# WebSocket 端点
@router.websocket("/ws/{user_id}")
async def websocket_endpoint(websocket: WebSocket, user_id: str):await manager.connect(websocket, user_id)try:while True:data = await websocket.receive_text()data = json.loads(data)# 接收到消息后，向目标用户发送消息await manager.send_message(data.get("message"), user_id, data.get("target_user"))except WebSocketDisconnect:manager.disconnect(user_id)print(f"用户 {user_id} 断开连接")app.include_router(router)