简单的 Go WebSocket 服务器

Linux

简单的 Go WebSocket 服务器

在现代 Web 开发中，WebSocket 是实现实时通信（Real-Time Communication）的标准解决方案。与传统的 HTTP 请求-响应模式不同，WebSocket 允许客户端和服务器之间建立全双工的长连接。

本博客通过编写一个简单的回显服务器，认识 Go 语言处理 WebSocket 的核心机制，特别是它如何从 HTTP 协议“升级”而来，以及底层的 Goroutine 模型是如何变化的。

服务端与客户端代码

使用 Go 社区最广泛使用的 WebSocket 库：gorilla/websocket。

这段代码建立了一个 HTTP 服务器，并将 /ws 路径的请求升级为 WebSocket 连接。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
package main

import (
	"log"
	"net/http"

	"github.com/gorilla/websocket"
)

// Upgrader 用于将普通的 HTTP 连接升级为 WebSocket 连接
var upgrader = websocket.Upgrader{
	// CheckOrigin 拦截器：控制允许连接的来源
	// 在开发阶段为了方便调试，我们直接返回 true，允许跨域访问
	// ⚠️ 警告：生产环境请务必进行严格的 Origin 校验以防御 CSRF 攻击
	CheckOrigin: func(r *http.Request) bool { return true },
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
	// 1) 握手阶段：HTTP -> WebSocket 升级
	// Upgrade 函数会劫持底层的 TCP 连接，并写入 101 Switching Protocols 响应
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("upgrade error:", err)
		return
	}
	// 确保连接在函数退出时关闭
	defer conn.Close()

	log.Println("client connected:", r.RemoteAddr)

	// 2) 数据传输阶段：在一个无限循环中读写消息
	for {
		// 读取消息（会阻塞，直到收到消息或出错）
		msgType, msg, err := conn.ReadMessage()
		if err != nil {
			log.Println("read error:", err)
			break // 任何读取错误都应视为连接断开的信号
		}
		log.Printf("recv: type=%d msg=%s\n", msgType, string(msg))

		// 回写消息（Echo）
		// 注意：msgType 需要保持一致（文本或二进制）
		if err := conn.WriteMessage(msgType, msg); err != nil {
			log.Println("write error:", err)
			break
		}
	}
}

func main() {
	// 注册 WebSocket 路由
	http.HandleFunc("/ws", wsHandler)

	// 提供静态文件服务（为了访问 index.html）
	http.Handle("/", http.FileServer(http.Dir(".")))

	addr := ":8080"
	log.Println("listening on", addr)
	// 启动 HTTP 服务
	log.Fatal(http.ListenAndServe(addr, nil))
}

一个原生的 HTML 页面，用于连接上面的 Go 服务器进行测试。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
<!doctype html>
<html>
  <body>
    <h3>WebSocket Test</h3>
    <input id="txt" placeholder="say something" />
    <button id="send">Send</button>
    <pre id="log" style="border: 1px solid #ccc; padding: 10px; min-height: 100px;"></pre>

    <script>
      // 辅助函数：向页面追加日志
      // 箭头函数 (s) => ... 是一种更简洁的写法
      const log = (s) => (document.getElementById("log").textContent += s + "\n");
      
      // 1. 创建 WebSocket 对象，自动触发握手请求
      const ws = new WebSocket("ws://localhost:8080/ws");

      // 2. 绑定生命周期回调函数
      ws.onopen = () => log("connected: connection established");
      ws.onmessage = (e) => log("echo: " + e.data);
      ws.onclose = () => log("closed: connection lost");
      ws.onerror = (e) => log("error: " + e);

      // 3. 发送消息
      document.getElementById("send").onclick = () => {
        const v = document.getElementById("txt").value;
        if (!v) return;
        ws.send(v);
        log("sent: " + v);
        document.getElementById("txt").value = ""; // 清空输入框
      };
    </script>
  </body>
</html>

从 HTTP 到 WebSocket 的“夺权”

如果了解 Go 的 HTTP 编程（参考之前的博客：最简单的 Go HTTP 服务器），会发现上面的代码结构非常相似。net/http 库本身并不理解 WebSocket 的帧格式，它只负责 HTTP 协议。

那么，WebSocket 是如何介入的？答案在于 Upgrade（协议升级）。

握手过程

WebSocket 的建立始于一个标准的 HTTP 请求。浏览器发送的请求头中包含特殊的字段：

1
2
3
4
5
6
GET /ws HTTP/1.1
Host: localhost:8080
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
Sec-WebSocket-Version: 13

这一步依然由 Go 的 standard library net/http 处理。当请求到达 wsHandler 时，我们调用了 upgrader.Upgrade(w, r, nil)。这个函数内部执行了关键的“夺权”操作：

  1. 校验：检查 HTTP 请求头中的 Upgrade、Connection 和 Sec-WebSocket-Key 是否合法。
  2. 响应：向客户端写入 HTTP 状态码 101 Switching Protocols，告知浏览器“我们要切换协议了”。
  3. 劫持 (Hijack)：这是最关键的一步。gorilla/websocket 会通过 http.ResponseWriter 提供的 Hijack 接口，将底层的 TCP 连接对象 (net.Conn) 直接提取出来。
  4. 接管：一旦 TCP 连接被提取，net/http 库就失去了对该连接的控制权。此后，这条 TCP 连接的所有读写操作完全由 WebSocket 库接管。

Goroutine 模型的质变

理解 WebSocket 在 Go 中运行机制的关键，在于理解 Goroutine 生命周期的变化。

普通 HTTP 模式（短连接）

net/http 的设计模型是“一请求一处理”：

  1. 接收请求。
  2. 启动一个 Goroutine 执行 Handler。
  3. Handler 写入响应并返回。
  4. Goroutine 结束，连接处于 Idle 状态或被关闭。
1
2
3
4
TCP 连接
   |
   |-- HTTP 请求 1 → Goroutine A (处理完即销毁)
   |-- HTTP 请求 2 → Goroutine B (处理完即销毁)

WebSocket 模式（长连接）

在调用 Upgrade 之后，Handler Goroutine 的角色发生了本质改变：它不再是一个处理完就跑的短工，而变成了维护这条长连接的“守护者”。

1
2
3
4
5
6
7
TCP 连接
   |
   |-- HTTP Upgrade → Goroutine A
   |                       |
   |                       |-- 劫持 TCP 连接
   |                       |-- 进入 for 循环 (Read/Write)
   |                       |-- 此 Goroutine 持续存活，直到连接断开

关键点：Upgrade 之后，Handler 函数是否返回，完全决定了连接的生死。

  • 如果 return —— 连接关闭。
  • 如果 for {} —— 连接保持。
  • 如果 go func() —— 可以将连接对象传递给其他 Goroutine 管理，主 Handler 退出也没关系（因为底层 TCP 连接已经被劫持，不再受 net/http 的超时控制）。

服务端：发送与接收循环

一旦进入 WebSocket 模式，我们就进入了纯粹的网络编程领域。

1
2
3
4
for {
    msgType, msg, err := conn.ReadMessage()
    // ...
}

ReadMessage 是一个阻塞调用。如果没有数据到达，当前 Goroutine 会暂停执行（挂起），并不消耗 CPU。当以下情况发生时，它会唤醒并返回：

  1. 收到完整的 WebSocket 消息：返回数据。
  2. 收到关闭帧 (Close Frame)：客户端主动断开。
  3. 底层 TCP 断开：网线拔了、WiFi 断了等。
  4. 读取超时：如果设置了 SetReadDeadline。

在网络编程中，错误 (Error) 是常态，而非异常。

1
2
3
4
if err := conn.WriteMessage(msgType, msg); err != nil {
    log.Println("write error:", err)
    break // 退出循环，触发 defer conn.Close()
}

任何读写错误通常都意味着连接不可用（发送缓冲区满、网络中断、对端关闭）。因此，标准的处理范式是：一旦发生 I/O 错误，立即跳出循环，关闭连接。这也是为什么我们在代码开头写了 defer conn.Close() 的原因。

浏览器客户端

为什么用箭头函数？

在 JavaScript 代码中，我们使用了 const log = (s) => ...。

1
2
3
4
5
// 传统写法
var log = function(s) { ... };

// ES6 箭头函数写法
const log = (s) => ...;

这里主要有两点考虑： 1. 简洁性：对于这种单行的辅助函数，箭头函数写起来更干净。 2. this 绑定机制： * function 定义的函数有自己的 this 上下文（动态作用域）。 * 箭头函数 => 不会创建自己的 this，它会“捕获”外层上下文的 this（词法作用域）。 * 虽然在这个简单的 log 函数中我们没有用到 this，但在编写 React/Vue 组件或类的方法时，箭头函数能避免很多 this 指向错误的坑。

WebSocket 的生命周期事件

我们在 JS 中定义了四个关键的回调函数，它们完整描述了 WebSocket 的一生：

  1. ws.onopen
    • 含义：连接握手成功，通道已建立。
    • 时机：服务器返回 101 Switching Protocols 之后。
    • 用途：通常在这里更新 UI（如显示“已连接”），或发送第一条初始化消息（如身份 Token）。
  2. ws.onmessage
    • 含义：收到了服务器发来的数据。
    • 参数：e.data 包含了实际的消息内容（字符串或 Blob/ArrayBuffer）。
    • 注意：这是异步触发的，服务器随时可能发消息过来。
  3. ws.onerror
    • 含义：发生了错误（如网络不可达、协议解析错误）。
    • 注意：onerror 触发后，通常紧接着会触发 onclose。
  4. ws.onclose
    • 含义：连接已彻底关闭。
    • 用途：清理资源、重置 UI、或者发起断线重连。
    • 注意：无论是服务器主动踢人、客户端主动关闭、还是网络异常断开，最终都会走到这里。

浏览器里的 URL

1
const ws = new WebSocket("ws://localhost:8080/ws");

浏览器看到 ws://（或加密的 wss://）协议头时，会自动构建一个 HTTP 请求，并附带必要的 Upgrade 头。

注意：不能直接在浏览器地址栏输入 ws://...，地址栏只支持 http/https。WebSocket 必须由 JavaScript 发起。

总结

通过这个简单的示例，我们揭示了 WebSocket 的核心：

  1. 它始于 HTTP，但在服务器端通过 Upgrade 劫持了底层 TCP 连接。
  2. Go 的处理模型从“请求响应式”转变为“长连接独占 Goroutine”模式。
  3. 需要自己编写循环来维护消息的读写，并时刻警惕网络错误的发生。

掌握了这些，就拥有了构建聊天室、实时游戏或即时通知系统的基石。