当前位置：首页 > news >正文

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

news 2025/7/11 17:16:29

概述

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