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建立一个连接，打开握手详情，可以看到这些信息：

1. 首先，客户端发送一个 HTTP GET 请求，包含以下关键头信息：

   • Upgrade: websocket：表示希望升级到 WebSocket 协议。

   • Connection: Upgrade：表示要求升级连接。

   • Sec-WebSocket-Key：一个随机生成的 Base64 编码字符串，用于验证服务器。

   • Sec-WebSocket-Version：指定 WebSocket 协议版本（通常为 13）。

2. 服务器检查请求头，如果支持 WebSocket，返回 HTTP 101 状态码（Switching Protocols），并包含以下头信息：

   • Upgrade: websocket：确认升级到 WebSocket 协议。
   • Connection: Upgrade：确认连接升级。
   • Sec-WebSocket-Accept：基于客户端的 Sec-WebSocket-Key 计算的值，用于验证。

关于wss和ws

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

简单实现

from fastapi import FastAPI, APIRouter, WebSocket, WebSocketDisconnect

app = 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}")
            break

app.include_router(router)

实战-实现一个聊天室

import json

from fastapi import FastAPI, APIRouter, WebSocket, WebSocketDisconnect

app = 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] = websocket

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

这里使用了user_id作为标志，根据用户id向指定的对象发送消息，可以看到8888用户给6666用户发送了个hello，而6666用户也接收到了相应的消息