尝试 FastAPI WebSocket 写简单的聊天应用

听说 WebSocket 那么一个东西已经很久很久了,但始终未曾尝试,大概对它的了解就是它可以让服务与客户端之间保持一个通道,并实现 Server Push 的功能,这样客户端就无需使用长连接或轮询来与服务端交互,更重要是 WebSocket 能复用 Web 的 80 或 443 端口号,通常是防火墙友好的端口。最近在为 WebSocket 配置 Nginx 反射代理时意识到要加上 Upgrade, Connection 头,加之各种流行的 Web 框架都提供了 WebSocket 的支持,所以又不得不对它重视起来。


下面以 Python 的 FastAPI 框架为例,实现一个简单的聊天站点,主要是参考自 FastAPI 的 WebSockets 官方样例。 当然如果从一开始就询问 ChatGPT: Create a simple chat web app with Python FastAPI framework + WebSocket, 肯定能让你获得一个不错的实现。但本人仍然希望能自我深度思考,而不是让 AI 帮我深度思考,而不是深度思考如何向 ChatGPT 问问题。记得多年前看过一位台湾历史老师讲《史记》, 经常说当看到哪一段时,把课本扣过来,然后思考碰到书中的情形之下自己会怎么想,所以觉得学任何东西当如此。

记得在读当前明月的《明朝那些事》,作者常常这么做,特别是碰到皇帝问话时,如果是他本人应该如何应对,如何回答的不偏不倚,妥当得体,漏不透风,这也符合他的公务员的谨慎思维,所以最后他抑郁了。

要完成本文的简单聊天网站,在 Python 虚拟环境中需安装如下组件
pip install fastapi uvicorn websockets

然后是服务端代码 ws_demo.py
 1from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 2from fastapi.responses import HTMLResponse
 3import uvicorn
 4
 5app = FastAPI()
 6
 7websockets = {}
 8
 9@app.websocket("/ws")
10async def websocket_endpoint(websocket: WebSocket):
11    await websocket.accept()
12    print("WebSocket connection established")
13    try:
14        while True:
15            data = await websocket.receive_text()
16            if data.startswith("@setName"):
17                user_name = data.split(" ", 1)[1]
18                websockets[user_name] = websocket
19                continue
20
21            from_who = next((name for name, ws in websockets.items() if ws == websocket), "Unknown")
22            if data.startswith("@"):
23                to_who, message = data[1:].split(" ", 1)
24                out_msg = f"{from_who}: {message}"
25                await websockets[to_who].send_text(out_msg)
26                await websocket.send_text(out_msg)
27            else:
28                for ws in websockets.values():
29                    await ws.send_text(f"{from_who}: {data}")
30    except WebSocketDisconnect:
31        print("WebSocket closed by client")
32        websockets.pop(next((name for name, ws in websockets.items() if ws == websocket), None), None)
33    except KeyError:
34        pass
35
36@app.get("/")
37async def index():
38    with open("app.html", 'r') as f:
39        return HTMLResponse(content=f.read())
40
41if __name__ == '__main__':
42    uvicorn.run(app, host="0.0.0.0", port=8000)

实现了一个非常粗陋的 WebSocket 连接管理器,让用户名与相应的 WebSocket 关联起来。主要功能是

  1. 当客户端发送消息 @setName <username>, 把 <username> 与当前连接的 WebSocket 关联
  2. 客户端发送不带 @<username> 为首的消息将把消息广播到所有客户端
  3. 当客户端发送以 @<username> <message content> 格式的消息时,只把消息投送到 @<username> 相对应的客户端,同时给自己来一份

用命令 python ws_demo.py 起动服务,WebSocket 服务是 ws://localhost:8000/ws

客户端网页代码 app.html
 1<!DOCTYPE html>
 2<html lang="en">
 3<head>
 4    <title>WebSocket Chat</title>
 5</head>
 6<body>
 7<h3>WebSocket Chat</h3>
 8<form action="" onsubmit="sendMessage(event)">
 9    <label for="messageText" id="name"></label>
10    <input type="text" id="messageText" placeholder="input name" autocomplete="off"/>
11    <button>Send</button>
12</form>
13<ul id='messages'>
14</ul>
15<script>
16    let ws = new WebSocket("ws://localhost:8000/ws");
17    let userName = ""
18    ws.onmessage = function (event) {
19        let msgLis = document.querySelectorAll("#messages li")
20        if (msgLis.length === 5) {
21            msgLis[0].remove()
22        }
23        let messages = document.getElementById('messages')
24        let message = document.createElement('li')
25        let content = document.createTextNode(event.data)
26        message.appendChild(content)
27        messages.appendChild(message)
28    };
29
30    function sendMessage(event) {
31        event.preventDefault()
32        let input = document.getElementById("messageText")
33        let message = input.value
34        if(userName === "") {
35            userName = input.value
36            input.setAttribute("placeholder", "input message")
37            document.getElementById("name").innerHTML = input.value
38            message = "@setName " + userName
39        }
40        ws.send(message)
41        input.value = ""
42    }
43</script>
44</body>
45</html>

我们三个浏览器窗口中打开 http://localhost:8000/, 没有输入用户名的话看到的输入框是 input name

在第一个浏览器中输入 user1, 然后看到 user1 [input message] (Send) 的式样。

在另两个窗口中分别输入 user2 或 user3, 以此来标识各自的用户名分别是 user1, user2 和 user3.

下图是我们在不同窗口中按序操作的结果

  1. user3 窗口中输入 hello to all, 看到三个窗口中都显示 user3: hello to all
  2. user3 窗口中继续输入 @user1 send to user1, 在 user1 和 user3 窗口同时显示了 user3: send to user1
  3. user3 窗口接着输入 @user2 send to user2, 在 user2 和 user3 窗口同时显示了 user: send to user2
  4. 最后在 user1 窗口中输入 hello to all, 同样在三个窗口中都马上显示了 user1: hello to all

这样一个多用户聊天的 Web 站点略具雏形,我们可以润色更丰富的功能,如用户注册,则进入聊天室自动显示用户名。更强的连接管理,如用户加入与退出的通知,聊天消息的 LocalStorage 存储,更漂亮的聊天页面展示。由于 WebSocket 可传送二进制,浏览器能访问麦克风,实现语音聊天都不是难事,或者语音编码为 Base64 传送都行。

WebSocket 协议分析

下面我们简单分析下由 http 协议如何升级到 ws 协议的,使用 Wireshark 抓包工具,对 WebSocket 中的关键包进行分析。与 HTTP, WebSocket 相关的包

在 Wireshark 中能很好的识别出 WebSocket 协议, 下面分析几个关键报文

浏览器在执行
1let ws = new WebSocket("ws://localhost:8000/ws");

相应的 HTTP 协议文本为
 1GET /ws HTTP/1.1
 2Host: localhost:8000
 3Connection: Upgrade
 4Pragma: no-cache
 5Cache-Control: no-cache
 6User-Agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/137.0.0.0 Safari/537.36
 7Upgrade: websocket
 8Origin: http://localhost:8000
 9Sec-WebSocket-Version: 13
10Accept-Encoding: gzip, deflate, br, zstd
11Accept-Language: en-US,en;q=0.9,zh-CN;q=0.8,zh;q=0.7
12Sec-WebSocket-Key: A3+6KUIiYQSSwVf+9Ec+ag==
13Sec-WebSocket-Extensions: permessage-deflate; client_max_window_bits

请求协议升级到 websocket

从服务端回过来的 HTTP 报文是
1HTTP/1.1 101 Switching Protocols
2Upgrade: websocket
3Connection: Upgrade
4Sec-WebSocket-Accept: wsjTAriES7avXQbPbdm7cvo9C98=
5Sec-WebSocket-Extensions: permessage-deflate
6date: Wed, 11 Jun 2025 05:06:11 GMT
7server: uvicorn

通知协议成功切换到 websocket,并建立起了 WebSocket 连接, 这可认为是 WebSocket 的握手过程。后面的通信与 HTTP 协议没有什么关系了,应用层协议由 HTTP 切换到了 WebSocket 二进制协议了。

从这里的协议切换过程可以理解为什么在为 WebSocket 配置反射代理时需考虑到 Upgrade: websocket 和 Connection: Upgrade 这两个头。

后面就是服务端与客户端之间的 WebSocket 通信,当客户端执行
1ws.send("@setName user1")

时,WebSocket 报文是

注意,WebSocket 不像 HTTP, 它是一个二进制的协议

看下一个客户端向服务端发送 hello to all 的 WebSocket 报文

下面是服务端向客户端回送 hello to all 的 WebSocket 报文

最后是关闭 WebSocket 连接的包。

在浏览器中用 Inspect 功能可以看到由 HTTP 升级到 WebSocket 的请求响应包,之后的 WebSocket 通信就无从知晓了。

关于 WebSocket 协议本身的 RFC 参考文件请见 RFC 6455 - The WebSocket ProtocolRFC 7963 - Clarifying Registry Procedures for the WebSocket Suprotocol Name Registry.

最后

WebSocket 是与 HTTP 处于对等位置的应用程协议,由 HTTP 切换过来的,除了复用了相同的端口号就没有别的关系了. WebSocket 的 URL 格式为

  1. ws://host[:port]/path 或
  2. wss://host[:port]/path

wss 对应于 HTTPS 的安全协议, ws 默认端口为 80, wss  默认端口为 443

从 Wireshark 中观察到当从 HTTP 协议升级到 WebSocket 协议,它们所用的本地端口号都是一样的,即 WebSocket 复用了 TCP 连接

当我们刷新或关掉浏览器的时候,在服务端的 WebSocket 连接会自动关闭,即
1try:
2    while True:
3        data = await websocket.receive_text()
4except WebSocketDisconnect:
5    print("WebSocket closed by client")       

会触发 WebSocketDisconnect 中的代码执行。

WebSocket 是二进制协议,所以可用来传送任何类型的数据。

WebSocket 的数据是按帧传输的,它的 Opcode 标志是 4 位,有以下选择

%x0: 帧继续,%x1: 表示文本帧,%x2 表示二进制帧,%3-7 保留为非控制帧用途, %x8 表示关闭连接,%x9 心跳 ping, %xA 心跳 pong, %xB-F 保留为控制帧用途。 永久链接 https://yanbin.blog/experience-fastapi-websocket-simple-chat/, 来自 隔叶黄莺 Yanbin's Blog
[版权声明] 本文采用 署名-非商业性使用-相同方式共享 4.0 国际 (CC BY-NC-SA 4.0) 进行许可。