📚 原创系列： “Go语言学习系列”
🔄 转载说明： 本文最初发布于"Gopher部落"微信公众号，经原作者授权转载。
🔗 关注原创： 欢迎扫描文末二维码，关注"Gopher部落"微信公众号获取第一手Go技术文章。
📑 Go语言学习系列导航
🚀 第三阶段：进阶篇本文是【Go语言学习系列】的第40篇，当前位于第三阶段（进阶篇）
- 并发编程（一）：goroutine基础
- 并发编程（二）：channel基础
- 并发编程（三）：select语句
- 并发编程（四）：sync包
- 并发编程（五）：并发模式
- 并发编程（六）：原子操作与内存模型
- 数据库编程（一）：SQL接口
- 数据库编程（二）：ORM技术
- Web开发（一）：路由与中间件
- Web开发（二）：模板与静态资源
- Web开发（三）：API开发
- Web开发（四）：认证与授权
- Web开发（五）：WebSocket 👈 当前位置
- 微服务（一）：基础概念
- 微服务（二）：gRPC入门
- 日志与监控
- 第三阶段项目实战：微服务聊天应用
📖 文章导读
在本文中，您将了解：
- WebSocket协议的基本原理与HTTP的区别
- 如何在Go中建立和管理WebSocket连接
- 实现WebSocket消息的发送、接收和处理
- 构建广播系统实现多客户端实时通信
- 心跳机制的实现确保连接稳定性
- 完整实现一个基于WebSocket的实时聊天应用
Web开发（五）：WebSocket
在前面的Web开发系列中，我们已经讨论了路由、中间件、模板渲染、API开发以及认证授权等技术。然而，这些技术主要用于构建传统的请求-响应模式的Web应用。现代Web应用通常需要更加实时的交互体验，例如聊天应用、协作工具或在线游戏等，这些场景都需要服务器能够主动推送数据给客户端。WebSocket正是为此类需求设计的协议，它实现了浏览器和服务器之间的全双工通信。
本文将详细介绍在Go语言中如何使用WebSocket，从基本概念到实际应用，帮助你掌握实时Web应用开发的核心技能。
1. WebSocket基础概念
1.1 WebSocket协议简介
WebSocket是一种在单个TCP连接上进行全双工通信的协议。WebSocket使得客户端和服务器之间的数据交换变得更加简单，允许服务端主动向客户端推送数据。在WebSocket API中，浏览器和服务器只需要完成一次握手，两者之间就可以创建持久性的连接，并进行双向数据传输。
WebSocket在2011年被IETF标准化为RFC 6455，并由W3C定义了对应的JavaScript API。
1.2 WebSocket vs HTTP
传统的HTTP协议有以下局限性：
- 单向通信：客户端发起请求，服务器响应，服务器不能主动推送数据
- 连接不持久：每次请求都需要建立新的TCP连接（虽然HTTP/1.1引入了Keep-Alive，但本质上仍是请求-响应模式）
- 头部开销大：HTTP请求和响应都包含大量头部信息
相比之下，WebSocket提供以下优势：
- 双向通信：服务器可以主动向客户端推送数据
- 持久连接：一次握手后保持TCP连接，减少连接建立的开销
- 较低的延迟：没有HTTP请求的额外开销，适合实时数据传输
- 更高效的数据传输：相比HTTP轮询，大大减少了数据传输量
1.3 WebSocket工作原理
WebSocket连接的建立过程如下：
-
握手阶段：客户端发起HTTP请求，请求升级到WebSocket协议
GET /chat HTTP/1.1 Host: server.example.com Upgrade: websocket Connection: Upgrade Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ== Sec-WebSocket-Version: 13 -
服务器响应：如果服务器支持WebSocket，则返回升级确认
HTTP/1.1 101 Switching Protocols Upgrade: websocket Connection: Upgrade Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo= -
建立连接：握手成功后，HTTP连接升级为WebSocket连接，此后的通信遵循WebSocket协议
-
数据传输：使用帧（frames）进行数据传输，帧可以是文本帧或二进制帧
-
关闭连接：任何一方都可以发送关闭帧来关闭连接
1.4 WebSocket应用场景
WebSocket特别适合以下场景：
- 聊天应用：用户间的实时消息传递
- 协作工具：多用户同时编辑文档、白板等
- 实时游戏：玩家间的实时交互
- 实时监控：数据仪表盘实时更新
- 股票行情：价格实时更新
- 体育赛事：比分实时显示
- 位置服务：实时位置追踪
2. Go语言中的WebSocket实现
在Go语言中，有多个WebSocket库可供选择，最常用的是Gorilla WebSocket库。接下来，我们将使用这个库实现WebSocket功能。
2.1 安装Gorilla WebSocket
首先，安装Gorilla WebSocket库：
go get github.com/gorilla/websocket
2.2 创建WebSocket服务器
下面是一个简单的WebSocket服务器示例：
package main
import (
"log"
"net/http"
"github.com/gorilla/websocket"
)
// 配置WebSocket升级器
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
// 允许所有CORS请求，生产环境应当配置更严格的检查
CheckOrigin: func(r *http.Request) bool {
return true
},
}
// WebSocket处理器
func handleWebSocket(w http.ResponseWriter, r *http.Request) {
// 将HTTP连接升级为WebSocket连接
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Println("升级连接失败:", err)
return
}
defer conn.Close()
log.Println("客户端已连接:", conn.RemoteAddr())
// 无限循环，持续读取客户端消息
for {
// 读取消息
messageType, message, err := conn.ReadMessage()
if err != nil {
log.Println("读取消息错误:", err)
break
}
// 打印收到的消息
log.Printf("收到消息: %s", message)
// 简单地将消息回送给客户端
err = conn.WriteMessage(messageType, message)
if err != nil {
log.Println("写入消息错误:", err)
break
}
}
}
func main() {
http.HandleFunc("/ws", handleWebSocket)
// 提供一个简单的客户端HTML页面
http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
http.ServeFile(w, r, "websocket.html")
})
log.Println("服务器启动在 :8080 端口")
log.Fatal(http.ListenAndServe(":8080", nil))
}
2.3 创建WebSocket客户端页面
创建一个简单的HTML客户端（websocket.html）：
<!DOCTYPE html>
<html>
<head>
<title>WebSocket测试</title>
<style>
body {
font-family: Arial, sans-serif;
max-width: 800px;
margin: 0 auto;
padding: 20px;
}
#messages {
height: 300px;
border: 1px solid #ccc;
margin-bottom: 10px;
padding: 10px;
overflow-y: auto;
}
input[type="text"] {
width: 70%;
padding: 8px;
}
button {
padding: 8px 15px;
background-color: #4CAF50;
color: white;
border: none;
cursor: pointer;
}
</style>
</head>
<body>
<h1>WebSocket测试</h1>
<div id="messages"></div>
<input type="text" id="messageInput" placeholder="输入消息">
<button onclick="sendMessage()">发送</button>
<script>
var socket;
var messagesDiv = document.getElementById('messages');
var messageInput = document.getElementById('messageInput');
// 建立WebSocket连接
function connect() {
// 创建一个WebSocket连接
socket = new WebSocket('ws://' + window.location.host + '/ws');
// 连接打开时的处理
socket.onopen = function(event) {
addMessage('系统', '连接已建立');
};
// 收到消息时的处理
socket.onmessage = function(event) {
addMessage('服务器', event.data);
};
// 连接关闭时的处理
socket.onclose = function(event) {
addMessage('系统', '连接已关闭');
// 尝试重连
setTimeout(function() {
connect();
}, 2000);
};
// 发生错误时的处理
socket.onerror = function(error) {
addMessage('错误', '发生错误');
console.error('WebSocket错误:', error);
};
}
// 发送消息
function sendMessage() {
if (!socket || socket.readyState !== WebSocket.OPEN) {
addMessage('错误', '连接未建立');
return;
}
var message = messageInput.value;
if (message) {
socket.send(message);
addMessage('我', message);
messageInput.value = '';
}
}
// 在消息区域添加一条消息
function addMessage(sender, message) {
var messageElement = document.createElement('div');
messageElement.textContent = sender + ': ' + message;
messagesDiv.appendChild(messageElement);
messagesDiv.scrollTop = messagesDiv.scrollHeight;
}
// 回车键发送消息
messageInput.addEventListener('keyup', function(event) {
if (event.key === 'Enter') {
sendMessage();
}
});
// 页面加载完成后建立连接
window.onload = connect;
</script>
</body>
</html>
这个简单的例子实现了客户端和服务器之间的实时消息交换。接下来，我们将深入探讨更高级的WebSocket功能和模式。
3. WebSocket消息广播系统
在许多WebSocket应用中，我们不仅需要处理单个客户端的连接，还需要向多个客户端广播消息。例如，在聊天室应用中，当一个用户发送消息时，所有在线用户都应该收到这条消息。接下来，我们将实现一个简单的广播系统。
3.1 设计广播系统
我们将设计一个广播系统，包含以下组件：
- 客户端结构体：用于存储每个WebSocket连接和相关信息
- 消息队列：用于存储要广播的消息
- 客户端管理器：用于管理所有已连接的客户端
下面是一个简单的实现：
package main
import (
"encoding/json"
"log"
"net/http"
"sync"
"time"
"github.com/gorilla/websocket"
)
// 配置WebSocket升级器
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
CheckOrigin: func(r *http.Request) bool {
return true
},
}
// 定义消息结构
type Message struct {
Type string `json:"type"`
Content string `json:"content"`
Sender string `json:"sender"`
Time time.Time `json:"time"`
}
// 客户端结构体
type Client struct {
ID string
Conn *websocket.Conn
Send chan []byte
Hub *Hub
Username string
}
// Hub管理所有已连接的客户端
type Hub struct {
clients map[*Client]bool
broadcast chan []byte
register chan *Client
unregister chan *Client
mutex sync.Mutex
}
// 创建新的Hub
func newHub() *Hub {
return &Hub{
broadcast: make(chan []byte),
register: make(chan *Client),
unregister: make(chan *Client),
clients: make(map[*Client]bool),
}
}
// Hub运行方法，管理客户端连接和消息广播
func (h *Hub) run() {
for {
select {
case client := <-h.register:
h.mutex.Lock()
h.clients[client] = true
h.mutex.Unlock()
// 通知所有客户端有新用户加入
joinMsg := Message{
Type: "system",
Content: client.Username + " 加入了聊天",
Time: time.Now(),
}
msgBytes, _ := json.Marshal(joinMsg)
h.broadcast <- msgBytes
case client := <-h.unregister:
h.mutex.Lock()
if _, ok := h.clients[client]; ok {
delete(h.clients, client)
close(client.Send)
// 通知所有客户端有用户离开
leaveMsg := Message{
Type: "system",
Content: client.Username + " 离开了聊天",
Time: time.Now(),
}
msgBytes, _ := json.Marshal(leaveMsg)
h.broadcast <- msgBytes
}
h.mutex.Unlock()
case message := <-h.broadcast:
h.mutex.Lock()
for client := range h.clients {
select {
case client.Send <- message:
// 消息发送成功
default:
// 客户端已无法接收消息
close(client.Send)
delete(h.clients, client)
}
}
h.mutex.Unlock()
}
}
}
// 客户端读取消息的goroutine
func (c *Client) readPump() {
defer func() {
c.Hub.unregister <- c
c.Conn.Close()
}()
// 设置最大消息大小
c.Conn.SetReadLimit(512)
// 设置读取超时
c.Conn.SetReadDeadline(time.Now().Add(60 * time.Second))
c.Conn.SetPongHandler(func(string) error {
c.Conn.SetReadDeadline(time.Now().Add(60 * time.Second))
return nil
})
// 持续读取客户端消息
for {
_, rawMessage, err := c.Conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
log.Printf("error: %v", err)
}
break
}
// 创建新消息并广播
message := Message{
Type: "message",
Content: string(rawMessage),
Sender: c.Username,
Time: time.Now(),
}
messageBytes, _ := json.Marshal(message)
c.Hub.broadcast <- messageBytes
}
}
// 客户端写入消息的goroutine
func (c *Client) writePump() {
ticker := time.NewTicker(54 * time.Second)
defer func() {
ticker.Stop()
c.Conn.Close()
}()
for {
select {
case message, ok := <-c.Send:
// 设置写入超时
c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
if !ok {
// 发送通道已关闭
c.Conn.WriteMessage(websocket.CloseMessage, []byte{})
return
}
// 获取写入器
writer, err := c.Conn.NextWriter(websocket.TextMessage)
if err != nil {
return
}
// 写入消息
writer.Write(message)
// 将Send通道中的所有消息一次性发送
n := len(c.Send)
for i := 0; i < n; i++ {
writer.Write([]byte{'\n'})
writer.Write(<-c.Send)
}
if err := writer.Close(); err != nil {
return
}
case <-ticker.C:
// 发送ping消息保持连接活跃
c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}
// 处理WebSocket连接
func serveWebSocket(hub *Hub, w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Println(err)
return
}
// 从查询参数中获取用户名，实际应用中可能从认证系统中获取
username := r.URL.Query().Get("username")
if username == "" {
username = "匿名用户"
}
client := &Client{
ID: conn.RemoteAddr().String(),
Hub: hub,
Conn: conn,
Send: make(chan []byte, 256),
Username: username,
}
client.Hub.register <- client
// 启动goroutine处理读写
go client.writePump()
go client.readPump()
}
func main() {
hub := newHub()
go hub.run()
// 提供WebSocket端点
http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
serveWebSocket(hub, w, r)
})
// 提供静态文件
http.Handle("/", http.FileServer(http.Dir("./static")))
log.Println("服务器启动在 :8080 端口")
log.Fatal(http.ListenAndServe(":8080", nil))
}
3.2 创建聊天室前端
接下来，创建一个更完善的聊天室前端界面（static/index.html）：
<!DOCTYPE html>
<html>
<head>
<title>Go WebSocket聊天室</title>
<style>
body {
font-family: Arial, sans-serif;
max-width: 800px;
margin: 0 auto;
padding: 20px;
}
#chatBox {
height: 400px;
border: 1px solid #ccc;
margin-bottom: 10px;
padding: 10px;
overflow-y: auto;
}
.message {
margin-bottom: 8px;
padding: 8px;
border-radius: 4px;
}
.message .time {
color: #999;
font-size: 0.8em;
}
.message .sender {
font-weight: bold;
margin-right: 5px;
}
.system {
background-color: #f8f8f8;
font-style: italic;
color: #666;
}
.user-message {
background-color: #e6f7ff;
}
.my-message {
background-color: #d4edda;
text-align: right;
}
.input-container {
display: flex;
}
#messageInput {
flex-grow: 1;
padding: 8px;
margin-right: 10px;
}
button {
padding: 8px 15px;
background-color: #4CAF50;
color: white;
border: none;
cursor: pointer;
}
button:hover {
background-color: #45a049;
}
#usernameContainer {
margin-bottom: 20px;
}
</style>
</head>
<body>
<h1>Go WebSocket聊天室</h1>
<div id="messages"></div>
<input type="text" id="messageInput" placeholder="输入消息">
<button onclick="sendMessage()">发送</button>
<script>
let socket;
let username = '';
const chatBox = document.getElementById('messages');
const messageInput = document.getElementById('messageInput');
// 建立WebSocket连接
function connect() {
// 创建一个WebSocket连接
socket = new WebSocket('ws://' + window.location.host + '/ws');
// 连接打开时的处理
socket.onopen = function(event) {
addMessage('系统', '连接已建立');
};
// 收到消息时的处理
socket.onmessage = function(event) {
addMessage('服务器', event.data);
};
// 连接关闭时的处理
socket.onclose = function(event) {
addMessage('系统', '连接已关闭');
// 尝试重连
setTimeout(function() {
connect();
}, 2000);
};
// 发生错误时的处理
socket.onerror = function(error) {
addMessage('错误', '发生错误');
console.error('WebSocket错误:', error);
};
}
// 发送消息
function sendMessage() {
if (!socket || socket.readyState !== WebSocket.OPEN) {
addMessage('错误', '连接未建立');
return;
}
var message = messageInput.value;
if (message) {
socket.send(message);
addMessage('我', message);
messageInput.value = '';
}
}
// 在消息区域添加一条消息
function addMessage(sender, message) {
var messageElement = document.createElement('div');
messageElement.textContent = sender + ': ' + message;
messagesDiv.appendChild(messageElement);
messagesDiv.scrollTop = messagesDiv.scrollHeight;
}
// 回车键发送消息
messageInput.addEventListener('keyup', function(event) {
if (event.key === 'Enter') {
sendMessage();
}
});
// 页面加载完成后建立连接
window.onload = connect;
</script>
</body>
</html>
这个简单的例子实现了客户端和服务器之间的实时消息交换。接下来，我们将深入探讨更高级的WebSocket功能和模式。
4. WebSocket连接管理与心跳机制
在实际应用中，WebSocket连接可能因为各种原因中断，如网络波动、防火墙超时等。为了保持连接的稳定性，我们需要实现心跳机制（Heartbeat）。
4.1 心跳机制的原理
心跳机制的基本原理是客户端和服务器定期交换短消息（通常称为"ping"和"pong"）以确认连接仍然有效。如果在预期的时间内没有收到响应，则认为连接已断开，需要重新连接。
WebSocket协议已经内置了ping/pong帧，我们可以使用它们实现心跳机制。
4.2 在Go中实现WebSocket心跳
Gorilla WebSocket库支持设置读取和写入的截止时间，以及处理ping/pong消息。下面是在之前示例基础上加入完整心跳机制的代码：
// 在Client.readPump方法中添加心跳检测
func (c *Client) readPump() {
defer func() {
c.Hub.unregister <- c
c.Conn.Close()
}()
// 设置最大消息大小
c.Conn.SetReadLimit(1024)
// 设置读取超时
c.Conn.SetReadDeadline(time.Now().Add(pongWait))
// 设置pong处理函数
c.Conn.SetPongHandler(func(string) error {
// 收到pong后延长读取截止时间
c.Conn.SetReadDeadline(time.Now().Add(pongWait))
return nil
})
// 持续读取消息
for {
_, message, err := c.Conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
log.Printf("error: %v", err)
}
break
}
// 处理消息（略）
}
}
// 在Client.writePump方法中添加ping发送
func (c *Client) writePump() {
// 定期发送ping消息的ticker
ticker := time.NewTicker(pingPeriod)
defer func() {
ticker.Stop()
c.Conn.Close()
}()
for {
select {
case message, ok := <-c.Send:
// 处理消息（略）
case <-ticker.C:
// 发送ping消息
c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}
在上面的代码中：
pongWait是等待客户端pong响应的超时时间，通常设为60秒pingPeriod是发送ping消息的间隔，通常设为pongWait * 9/10，例如54秒writeWait是写入超时，通常设为10秒
4.3 客户端心跳实现
浏览器中的WebSocket API会自动处理ping/pong消息，因此客户端通常不需要显式实现。但在某些特殊情况下，我们可能需要自定义心跳机制：
function setupHeartbeat(socket) {
const pingInterval = 30000; // 30秒
let pingTimer;
function sendPing() {
if (socket.readyState === WebSocket.OPEN) {
// 发送自定义ping消息
socket.send(JSON.stringify({type: 'ping'}));
}
}
// 开始定期发送ping
function startHeartbeat() {
pingTimer = setInterval(sendPing, pingInterval);
}
// 停止发送ping
function stopHeartbeat() {
clearInterval(pingTimer);
}
// 当连接打开时，开始心跳
socket.addEventListener('open', startHeartbeat);
// 当连接关闭时，停止心跳
socket.addEventListener('close', stopHeartbeat);
// 处理服务器响应
socket.addEventListener('message', function(event) {
const data = JSON.parse(event.data);
if (data.type === 'pong') {
console.log('Received pong from server');
}
});
}
5. 实现完整的WebSocket聊天应用
综合前面的知识，我们来实现一个完整的聊天应用，包含以下特性：
- 用户认证与管理
- 消息广播
- 私人消息
- 在线状态
- 消息历史记录
- 心跳机制
5.1 结构设计
首先，我们定义消息和用户的数据结构：
// message.go
package main
import "time"
type MessageType string
const (
TextMessage MessageType = "text"
SystemMessage MessageType = "system"
JoinMessage MessageType = "join"
LeaveMessage MessageType = "leave"
PrivateMessage MessageType = "private"
StatusMessage MessageType = "status"
)
type Message struct {
ID string `json:"id"`
Type MessageType `json:"type"`
Content string `json:"content"`
Sender string `json:"sender"`
Recipient string `json:"recipient,omitempty"` // 私信接收者
Time time.Time `json:"time"`
}
type User struct {
Username string `json:"username"`
Online bool `json:"online"`
}
5.2 客户端与连接管理
然后，我们实现客户端和消息处理：
// client.go
package main
import (
"encoding/json"
"log"
"time"
"github.com/google/uuid"
"github.com/gorilla/websocket"
)
const (
// 读取和写入超时
writeWait = 10 * time.Second
pongWait = 60 * time.Second
pingPeriod = (pongWait * 9) / 10
// 最大消息大小
maxMessageSize = 1024
)
// 客户端结构体
type Client struct {
hub *Hub
conn *websocket.Conn
send chan []byte
username string
}
// 读取来自客户端的消息
func (c *Client) readPump() {
defer func() {
c.hub.unregister <- c
c.conn.Close()
}()
c.conn.SetReadLimit(maxMessageSize)
c.conn.SetReadDeadline(time.Now().Add(pongWait))
c.conn.SetPongHandler(func(string) error {
c.conn.SetReadDeadline(time.Now().Add(pongWait))
return nil
})
for {
_, data, err := c.conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
log.Printf("error: %v", err)
}
break
}
// 解析消息
var msg struct {
Type string `json:"type"`
Content string `json:"content"`
To string `json:"to,omitempty"`
}
if err := json.Unmarshal(data, &msg); err != nil {
log.Println("Failed to unmarshal message:", err)
continue
}
// 创建消息对象
var message Message
switch msg.Type {
case "text":
// 普通文本消息
message = Message{
ID: uuid.New().String(),
Type: TextMessage,
Content: msg.Content,
Sender: c.username,
Time: time.Now(),
}
case "private":
// 私人消息
if msg.To == "" {
continue
}
message = Message{
ID: uuid.New().String(),
Type: PrivateMessage,
Content: msg.Content,
Sender: c.username,
Recipient: msg.To,
Time: time.Now(),
}
default:
continue
}
// 将消息发送到Hub进行处理
messageBytes, _ := json.Marshal(message)
c.hub.broadcast <- &BroadcastMessage{
message: messageBytes,
private: msg.Type == "private",
sender: c,
to: msg.To,
}
}
}
// 将消息写入WebSocket连接
func (c *Client) writePump() {
ticker := time.NewTicker(pingPeriod)
defer func() {
ticker.Stop()
c.conn.Close()
}()
for {
select {
case message, ok := <-c.send:
c.conn.SetWriteDeadline(time.Now().Add(writeWait))
if !ok {
// 通道已关闭
c.conn.WriteMessage(websocket.CloseMessage, []byte{})
return
}
w, err := c.conn.NextWriter(websocket.TextMessage)
if err != nil {
return
}
w.Write(message)
// 添加队列中的所有消息
n := len(c.send)
for i := 0; i < n; i++ {
w.Write([]byte{'\n'})
w.Write(<-c.send)
}
if err := w.Close(); err != nil {
return
}
case <-ticker.C:
c.conn.SetWriteDeadline(time.Now().Add(writeWait))
if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}
5.3 Hub与消息广播
接下来，我们实现消息管理和广播：
// hub.go
package main
import (
"encoding/json"
"log"
"sync"
"time"
)
// 广播消息结构
type BroadcastMessage struct {
message []byte
private bool
sender *Client
to string
}
// Hub管理所有客户端连接和消息广播
type Hub struct {
clients map[*Client]bool
usernames map[string]*Client
broadcast chan *BroadcastMessage
register chan *Client
unregister chan *Client
messages []Message
mutex sync.RWMutex
}
// 创建新Hub
func newHub() *Hub {
return &Hub{
broadcast: make(chan *BroadcastMessage),
register: make(chan *Client),
unregister: make(chan *Client),
clients: make(map[*Client]bool),
usernames: make(map[string]*Client),
messages: make([]Message, 0, 100),
}
}
// 运行Hub，处理消息
func (h *Hub) run() {
for {
select {
case client := <-h.register:
// 注册新客户端
h.mutex.Lock()
h.clients[client] = true
h.usernames[client.username] = client
h.mutex.Unlock()
// 通知所有用户有新用户加入
joinMsg := Message{
ID: "system-" + time.Now().String(),
Type: JoinMessage,
Content: client.username + " 加入了聊天",
Sender: "system",
Time: time.Now(),
}
msgBytes, _ := json.Marshal(joinMsg)
h.broadcastToAll(msgBytes)
// 发送用户列表给新用户
h.sendUserList(client)
// 发送最近的消息历史
h.sendMessageHistory(client)
case client := <-h.unregister:
h.mutex.Lock()
if _, ok := h.clients[client]; ok {
delete(h.clients, client)
delete(h.usernames, client.username)
close(client.send)
// 通知所有用户有用户离开
leaveMsg := Message{
ID: "system-" + time.Now().String(),
Type: LeaveMessage,
Content: client.username + " 离开了聊天",
Sender: "system",
Time: time.Now(),
}
msgBytes, _ := json.Marshal(leaveMsg)
h.mutex.Unlock()
h.broadcastToAll(msgBytes)
} else {
h.mutex.Unlock()
}
case broadcastMsg := <-h.broadcast:
// 解析消息
var msg Message
if err := json.Unmarshal(broadcastMsg.message, &msg); err != nil {
log.Println("Failed to unmarshal broadcast message:", err)
continue
}
// 保存消息历史（仅公共消息）
if !broadcastMsg.private && msg.Type == TextMessage {
h.mutex.Lock()
h.messages = append(h.messages, msg)
// 仅保留最近100条消息
if len(h.messages) > 100 {
h.messages = h.messages[len(h.messages)-100:]
}
h.mutex.Unlock()
}
// 发送消息
if broadcastMsg.private {
// 私人消息
h.mutex.RLock()
if recipient, ok := h.usernames[broadcastMsg.to]; ok {
select {
case recipient.send <- broadcastMsg.message:
// 消息发送成功
default:
// 接收方无法接收消息
close(recipient.send)
delete(h.clients, recipient)
delete(h.usernames, recipient.username)
}
}
// 也发送给发送方
if sender, ok := h.clients[broadcastMsg.sender]; ok {
select {
case sender.send <- broadcastMsg.message:
// 消息发送成功
default:
// 发送方无法接收消息
}
}
h.mutex.RUnlock()
} else {
// 公共消息，广播给所有人
h.broadcastToAll(broadcastMsg.message)
}
}
}
}
// 广播消息给所有客户端
func (h *Hub) broadcastToAll(message []byte) {
h.mutex.RLock()
defer h.mutex.RUnlock()
for client := range h.clients {
select {
case client.send <- message:
// 消息发送成功
default:
// 客户端无法接收消息
close(client.send)
delete(h.clients, client)
delete(h.usernames, client.username)
}
}
}
// 发送用户列表给指定客户端
func (h *Hub) sendUserList(client *Client) {
h.mutex.RLock()
userList := make([]User, 0, len(h.usernames))
for username := range h.usernames {
userList = append(userList, User{
Username: username,
Online: true,
})
}
h.mutex.RUnlock()
userListMsg := struct {
Type string `json:"type"`
Users []User `json:"users"`
}{
Type: "user_list",
Users: userList,
}
if data, err := json.Marshal(userListMsg); err == nil {
client.send <- data
}
}
// 发送消息历史给指定客户端
func (h *Hub) sendMessageHistory(client *Client) {
h.mutex.RLock()
messages := make([]Message, len(h.messages))
copy(messages, h.messages)
h.mutex.RUnlock()
historyMsg := struct {
Type string `json:"type"`
Messages []Message `json:"messages"`
}{
Type: "history",
Messages: messages,
}
if data, err := json.Marshal(historyMsg); err == nil {
client.send <- data
}
}
5.4 主函数与HTTP处理器
最后，我们实现主函数和HTTP处理器：
// main.go
package main
import (
"log"
"net/http"
"github.com/gorilla/websocket"
)
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
CheckOrigin: func(r *http.Request) bool {
return true
},
}
func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Println(err)
return
}
// 获取用户名
username := r.URL.Query().Get("username")
if username == "" {
conn.Close()
return
}
// 检查用户名是否已经存在
hub.mutex.RLock()
if _, exists := hub.usernames[username]; exists {
hub.mutex.RUnlock()
conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"error","content":"用户名已存在"}`))
conn.Close()
return
}
hub.mutex.RUnlock()
client := &Client{
hub: hub,
conn: conn,
send: make(chan []byte, 256),
username: username,
}
// 注册客户端
hub.register <- client
// 启动goroutine处理WebSocket
go client.writePump()
go client.readPump()
}
func main() {
// 创建Hub并运行
hub := newHub()
go hub.run()
// 静态文件服务
http.Handle("/", http.FileServer(http.Dir("./static")))
// WebSocket端点
http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
serveWs(hub, w, r)
})
// 启动服务器
addr := ":8080"
log.Printf("启动服务器在 %s", addr)
err := http.ListenAndServe(addr, nil)
if err != nil {
log.Fatal("ListenAndServe: ", err)
}
}
5.5 前端实现
为了使用我们的聊天服务器，需要创建一个更复杂的前端界面（static/index.html）：
<!DOCTYPE html>
<html>
<head>
<title>Go WebSocket实时聊天</title>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<style>
body {
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
max-width: 1000px;
margin: 0 auto;
padding: 20px;
background-color: #f7f9fc;
}
.container {
display: flex;
height: 80vh;
border-radius: 8px;
box-shadow: 0 4px 8px rgba(0,0,0,0.1);
overflow: hidden;
}
#sidebar {
width: 250px;
background-color: #2c3e50;
color: white;
padding: 15px;
overflow-y: auto;
}
#chatArea {
flex-grow: 1;
display: flex;
flex-direction: column;
background-color: white;
}
#chatMessages {
flex-grow: 1;
padding: 15px;
overflow-y: auto;
background-color: #f9f9f9;
}
.input-area {
padding: 15px;
background-color: #eee;
display: flex;
}
#messageInput {
flex-grow: 1;
padding: 10px;
border: 1px solid #ddd;
border-radius: 4px;
margin-right: 10px;
}
button {
padding: 10px 15px;
background-color: #3498db;
color: white;
border: none;
border-radius: 4px;
cursor: pointer;
}
button:hover {
background-color: #2980b9;
}
.login-container {
max-width: 400px;
margin: 100px auto;
padding: 30px;
background-color: white;
border-radius: 8px;
box-shadow: 0 4px 8px rgba(0,0,0,0.1);
text-align: center;
}
h1, h2, h3 {
color: #2c3e50;
}
.user {
padding: 10px;
margin: 5px 0;
border-radius: 4px;
cursor: pointer;
transition: background-color 0.3s;
}
.user:hover {
background-color: #34495e;
}
.user.active {
background-color: #3498db;
}
.message {
margin-bottom: 15px;
padding: 10px;
border-radius: 8px;
max-width: 70%;
}
.system-message {
background-color: #f1f1f1;
color: #777;
padding: 8px 15px;
margin: 10px auto;
text-align: center;
border-radius: 20px;
font-style: italic;
max-width: 80%;
}
.user-message {
background-color: #e8f4f8;
margin-right: auto;
margin-left: 5px;
border-top-left-radius: 0;
}
.my-message {
background-color: #dcf8c6;
margin-left: auto;
margin-right: 5px;
border-top-right-radius: 0;
}
.private {
background-color: #fff0c8;
border: 1px dashed #ffcc00;
}
.sender {
font-weight: bold;
margin-bottom: 5px;
}
.time {
font-size: 0.7em;
color: #888;
text-align: right;
margin-top: 5px;
}
.tabs {
display: flex;
border-bottom: 1px solid #ddd;
background-color: #f5f5f5;
}
.tab {
padding: 10px 20px;
cursor: pointer;
border-bottom: 2px solid transparent;
}
.tab.active {
border-bottom: 2px solid #3498db;
background-color: white;
}
#privateChat {
display: none;
}
</style>
</head>
<body>
<!-- 登录界面 -->
<div id="loginScreen" class="login-container">
<h1>Go WebSocket聊天</h1>
<p>请输入您的用户名开始聊天</p>
<div style="margin-top: 20px;">
<input type="text" id="usernameInput" placeholder="用户名" style="padding: 10px; width: 70%; margin-right: 10px;">
<button id="loginButton">登录</button>
</div>
<p id="loginError" style="color: red; display: none;"></p>
</div>
<!-- 聊天界面 -->
<div id="chatScreen" style="display: none;">
<h1>Go WebSocket聊天</h1>
<div class="container">
<!-- 侧边栏：用户列表 -->
<div id="sidebar">
<h3>在线用户</h3>
<div id="userList"></div>
</div>
<!-- 聊天区域 -->
<div id="chatArea">
<div class="tabs">
<div id="publicTab" class="tab active">公共聊天</div>
<div id="privateTab" class="tab">私人聊天</div>
</div>
<!-- 公共聊天 -->
<div id="publicChat">
<div id="chatMessages"></div>
<div class="input-area">
<input type="text" id="messageInput" placeholder="输入消息...">
<button id="sendButton">发送</button>
</div>
</div>
<!-- 私人聊天 -->
<div id="privateChat">
<div id="privateChatMessages"></div>
<div class="input-area">
<input type="text" id="privateMessageInput" placeholder="输入私人消息...">
<button id="sendPrivateButton">发送</button>
</div>
</div>
</div>
</div>
</div>
<script>
// DOM元素
const loginScreen = document.getElementById('loginScreen');
const chatScreen = document.getElementById('chatScreen');
const usernameInput = document.getElementById('usernameInput');
const loginButton = document.getElementById('loginButton');
const loginError = document.getElementById('loginError');
const userList = document.getElementById('userList');
const chatMessages = document.getElementById('chatMessages');
const privateChatMessages = document.getElementById('privateChatMessages');
const messageInput = document.getElementById('messageInput');
const privateMessageInput = document.getElementById('privateMessageInput');
const sendButton = document.getElementById('sendButton');
const sendPrivateButton = document.getElementById('sendPrivateButton');
const publicTab = document.getElementById('publicTab');
const privateTab = document.getElementById('privateTab');
const publicChat = document.getElementById('publicChat');
const privateChat = document.getElementById('privateChat');
// 全局变量
let socket;
let username = '';
let selectedUser = null;
// 标签切换
publicTab.addEventListener('click', () => {
publicTab.classList.add('active');
privateTab.classList.remove('active');
publicChat.style.display = 'block';
privateChat.style.display = 'none';
});
privateTab.addEventListener('click', () => {
if (selectedUser) {
publicTab.classList.remove('active');
privateTab.classList.add('active');
publicChat.style.display = 'none';
privateChat.style.display = 'block';
} else {
alert('请先选择一个用户进行私聊');
}
});
// 登录处理
loginButton.addEventListener('click', () => {
username = usernameInput.value.trim();
if (username) {
connectWebSocket();
} else {
showLoginError('请输入用户名');
}
});
usernameInput.addEventListener('keyup', (e) => {
if (e.key === 'Enter') {
loginButton.click();
}
});
// 显示登录错误
function showLoginError(message) {
loginError.textContent = message;
loginError.style.display = 'block';
}
// 连接WebSocket
function connectWebSocket() {
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
const wsUrl = `${protocol}//${window.location.host}/ws?username=${encodeURIComponent(username)}`;
socket = new WebSocket(wsUrl);
// 连接打开
socket.addEventListener('open', () => {
// 切换到聊天界面
loginScreen.style.display = 'none';
chatScreen.style.display = 'block';
// 添加系统消息
addSystemMessage('已连接到服务器');
});
// 收到消息
socket.addEventListener('message', (event) => {
const data = JSON.parse(event.data);
switch (data.type) {
case 'text':
case 'join':
case 'leave':
handleChatMessage(data);
break;
case 'private':
handlePrivateMessage(data);
break;
case 'user_list':
updateUserList(data.users);
break;
case 'history':
loadMessageHistory(data.messages);
break;
case 'error':
handleError(data);
break;
}
});
// 连接关闭
socket.addEventListener('close', () => {
addSystemMessage('与服务器的连接已关闭');
// 3秒后尝试重连
setTimeout(() => {
if (username) {
connectWebSocket();
}
}, 3000);
});
// 连接错误
socket.addEventListener('error', (error) => {
console.error('WebSocket错误:', error);
addSystemMessage('连接错误');
});
}
// 发送公共消息
sendButton.addEventListener('click', () => {
const message = messageInput.value.trim();
if (message && socket.readyState === WebSocket.OPEN) {
socket.send(JSON.stringify({
type: 'text',
content: message
}));
messageInput.value = '';
messageInput.focus();
}
});
messageInput.addEventListener('keyup', (e) => {
if (e.key === 'Enter') {
sendButton.click();
}
});
// 发送私人消息
sendPrivateButton.addEventListener('click', () => {
const message = privateMessageInput.value.trim();
if (message && selectedUser && socket.readyState === WebSocket.OPEN) {
socket.send(JSON.stringify({
type: 'private',
content: message,
to: selectedUser
}));
privateMessageInput.value = '';
privateMessageInput.focus();
}
});
privateMessageInput.addEventListener('keyup', (e) => {
if (e.key === 'Enter') {
sendPrivateButton.click();
}
});
// 处理普通聊天消息
function handleChatMessage(message) {
const messageElement = document.createElement('div');
if (message.type === 'join' || message.type === 'leave') {
messageElement.className = 'system-message';
messageElement.textContent = message.content;
} else {
messageElement.className = 'message';
if (message.sender === username) {
messageElement.className += ' my-message';
} else {
messageElement.className += ' user-message';
const senderElement = document.createElement('div');
senderElement.className = 'sender';
senderElement.textContent = message.sender;
messageElement.appendChild(senderElement);
}
const contentElement = document.createElement('div');
contentElement.className = 'content';
contentElement.textContent = message.content;
messageElement.appendChild(contentElement);
const timeElement = document.createElement('div');
timeElement.className = 'time';
timeElement.textContent = formatTime(new Date(message.time));
messageElement.appendChild(timeElement);
}
chatMessages.appendChild(messageElement);
chatMessages.scrollTop = chatMessages.scrollHeight;
}
// 处理私人消息
function handlePrivateMessage(message) {
const messageElement = document.createElement('div');
messageElement.className = 'message private';
if (message.sender === username) {
messageElement.className += ' my-message';
// 如果选中的不是消息接收者，就切换过去
if (selectedUser !== message.recipient) {
selectUser(message.recipient);
}
} else {
messageElement.className += ' user-message';
const senderElement = document.createElement('div');
senderElement.className = 'sender';
senderElement.textContent = message.sender;
messageElement.appendChild(senderElement);
// 如果选中的不是消息发送者，就切换过去
if (selectedUser !== message.sender) {
selectUser(message.sender);
}
}
const contentElement = document.createElement('div');
contentElement.className = 'content';
contentElement.textContent = message.content;
messageElement.appendChild(contentElement);
const timeElement = document.createElement('div');
timeElement.className = 'time';
timeElement.textContent = formatTime(new Date(message.time));
messageElement.appendChild(timeElement);
privateChatMessages.appendChild(messageElement);
privateChatMessages.scrollTop = privateChatMessages.scrollHeight;
// 自动切换到私聊标签
privateTab.click();
}
// 添加系统消息
function addSystemMessage(message) {
const messageElement = document.createElement('div');
messageElement.className = 'system-message';
messageElement.textContent = message;
chatMessages.appendChild(messageElement);
chatMessages.scrollTop = chatMessages.scrollHeight;
}
// 更新用户列表
function updateUserList(users) {
userList.innerHTML = '';
users.forEach(user => {
if (user.username !== username) {
const userElement = document.createElement('div');
userElement.className = 'user';
userElement.textContent = user.username;
userElement.addEventListener('click', () => {
selectUser(user.username);
});
if (user.username === selectedUser) {
userElement.classList.add('active');
}
userList.appendChild(userElement);
}
});
}
// 选择用户进行私聊
function selectUser(username) {
selectedUser = username;
// 更新UI
const userElements = userList.querySelectorAll('.user');
userElements.forEach(element => {
if (element.textContent === username) {
element.classList.add('active');
} else {
element.classList.remove('active');
}
});
privateTab.textContent = `与 ${username} 私聊`;
privateTab.click();
}
// 加载消息历史
function loadMessageHistory(messages) {
chatMessages.innerHTML = '';
messages.forEach(message => {
handleChatMessage(message);
});
chatMessages.scrollTop = chatMessages.scrollHeight;
}
// 处理错误消息
function handleError(data) {
if (loginScreen.style.display !== 'none') {
showLoginError(data.content);
} else {
addSystemMessage(`错误: ${data.content}`);
}
}
// 格式化时间
function formatTime(date) {
const hours = date.getHours().toString().padStart(2, '0');
const minutes = date.getMinutes().toString().padStart(2, '0');
return `${hours}:${minutes}`;
}
</script>
</body>
</html>
5.6 运行与测试
按照以下步骤运行和测试我们的聊天应用：
-
创建项目目录结构：
chat-app/ ├── main.go ├── client.go ├── hub.go ├── message.go └── static/ └── index.html -
安装依赖：
go get github.com/gorilla/websocket go get github.com/google/uuid -
运行应用：
go run *.go -
在多个浏览器窗口中打开 http://localhost:8080，使用不同的用户名登录
-
测试各种功能：
- 发送公共消息
- 选择用户发送私人消息
- 查看在线用户列表
- 测试断线重连
6. WebSocket安全与性能最佳实践
在生产环境中使用WebSocket时，以下是一些安全和性能最佳实践：
6.1 安全性考虑
-
输入验证：始终验证所有来自客户端的输入，防止注入攻击
-
使用HTTPS/WSS：在生产环境中使用WSS (WebSocket Secure)，而不是WS
-
认证与授权：
- 要求用户先通过HTTP认证，然后才能建立WebSocket连接
- 在查询参数或头部传递认证令牌
- 使用有状态或无状态的会话管理
-
来源检查：限制WebSocket连接的来源：
var upgrader = websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { // 检查Origin头，仅允许特定域 origin := r.Header.Get("Origin") return origin == "https://myapp.com" }, } -
连接限制：
- 限制每个IP或用户的连接数
- 设置连接超时
- 实现速率限制防止DDoS攻击
-
数据验证与清理：
- 验证消息格式和大小
- 清理HTML和JS代码防止XSS
6.2 性能优化
-
消息压缩：对大型消息使用压缩：
var upgrader = websocket.Upgrader{ EnableCompression: true, } -
批量操作：将多个小消息合并为一个大消息，减少传输开销
-
缓存：缓存不经常变化的数据，如用户列表和历史消息
-
水平扩展：
- 使用Redis等发布/订阅系统在多个服务器实例间广播消息
- 使用负载均衡在多个实例之间分发WebSocket连接
-
消息队列：使用消息队列（如RabbitMQ、NATS、Kafka）处理大量消息
-
定期清理：清除长时间不活动的连接
-
优化内存使用：
- 有效管理goroutine数量
- 限制每个连接的消息缓冲区大小
- 使用对象池减少内存分配
7. 总结
在本文中，我们深入探讨了WebSocket协议及其在Go语言中的实现，涵盖了以下关键内容：
-
WebSocket基础知识：了解了WebSocket协议的特性、适用场景和与HTTP的区别。
-
Go中的WebSocket库：学习了如何使用Gorilla WebSocket和标准库net/websocket包实现WebSocket通信。
-
WebSocket服务器实现：掌握了如何设置连接升级、处理连接、发送和接收消息等基本操作。
-
WebSocket客户端实现：探讨了前端JavaScript和Go客户端的实现方式。
-
实用WebSocket模式：学习了广播、房间和频道等常见的应用模式。
-
安全性和最佳实践：了解了WebSocket安全风险及其防范措施。
👨💻 关于作者与Gopher部落
"Gopher部落"专注于Go语言技术分享，提供从入门到精通的完整学习路线。
🌟 为什么关注我们？
- 系统化学习路径：本系列44篇文章循序渐进，带你完整掌握Go开发
- 实战驱动教学：理论结合实践，每篇文章都有可操作的代码示例
- 持续更新内容：定期分享最新Go生态技术动态与大厂实践经验
- 专业技术社区：加入我们的技术交流群，与众多Go开发者共同成长
📱 关注方式
- 微信公众号：搜索 “Gopher部落” 或 “GopherTribe”
- 优快云专栏：点击页面右上角"关注"按钮
💡 读者福利
关注公众号回复 “Go学习” 即可获取：
- 完整Go学习路线图
- Go面试题大全PDF
- Go项目实战源码
- 定制学习计划指导
期待与您在Go语言的学习旅程中共同成长！
7282
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
