Go 实现（一）：gorilla/websocket 升级与收发

gorilla/websocket 做两件事：把 HTTP 连接升级（upgrade）成帧读写器，并把"一读一写"的并发契约交给使用者。协议握手与帧解析它都处理好了，业务只需在连接上收发。但它的默认值不是生产值——读限制、读写超时、跨域校验都得自己设，否则第一个恶意客户端就能把你打疼。安装只需要 go get github.com/gorilla/websocket（版本与 API 细节以官方 README 与 pkg.go.dev 为准，目前仍是 v1 系）。

选型：几个 Go WebSocket 库的差别

库	特点	适合场景
`github.com/gorilla/websocket`	使用最广、API 直白、回调式（pong/close handler）、示例丰富	自建 IM 网关的默认选择
`github.com/coder/websocket`（原 `nhooyr.io/websocket`）	API 全程带 `context`、默认并发安全、体积小、支持 wasm	新项目，希望超时与取消统一走 context
`github.com/gobwas/ws`	贴近帧层、零分配、控制力最强	极致性能、愿意自己维护状态机
`golang.org/x/net/websocket`	官方已不再演进	只建议维护遗留代码

切换到 coder/websocket 时概念完全相同，只是 Upgrade／Dial 都换成接收 context.Context 的形式。

Upgrader：升级前必须做的几件事

var upgrader = websocket.Upgrader{
	HandshakeTimeout: 5 * time.Second, // 慢握手会白占连接与内存
	ReadBufferSize:   4096,
	WriteBufferSize:  4096,
	// nil 时默认执行同源校验；一旦自己写，就要真的收紧
	CheckOrigin: func(r *http.Request) bool {
		origin := r.Header.Get("Origin")
		return origin == "" || allowedOrigin(origin) // allowedOrigin 是项目内的白名单校验
	},
}

字段	作用	生产建议
HandshakeTimeout	握手阶段的读写超时	默认 0（不限制），设 5s
ReadBufferSize / WriteBufferSize	底层 bufio 缓冲大小	4 KB 起步；设 0 会退回默认小缓冲，反而多系统调用
CheckOrigin	跨域校验	用白名单，别一律 `return true`
Subprotocols	子协议协商	需要协议版本协商时显式列出
EnableCompression	permessage-deflate	默认关闭，小消息场景通常没必要开
WriteBufferPool	复用写缓冲	连接数很大时可显著降低 GC 压力

一个能跑起来的服务端

package main

import (
	"errors"
	"log"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
)

const maxMessageSize = 32 << 10 // 32 KiB

func wsHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		// Upgrade 内部已写过响应并关闭连接，这里只记日志，不要再写 http 响应
		log.Printf("upgrade failed: %v", err)
		return
	}
	defer conn.Close()

	conn.SetReadLimit(maxMessageSize)

	for {
		// 读超时的正确做法要配合 ping/pong 重置，见第 06 篇心跳章节
		_ = conn.SetReadDeadline(time.Now().Add(60 * time.Second))

		mt, data, err := conn.ReadMessage()
		if err != nil {
			var ce *websocket.CloseError
			if errors.As(err, &ce) { // 对端带关闭码断开
				log.Printf("closed by peer: code=%d reason=%q", ce.Code, ce.Text)
			}
			return
		}

		// 示例语义：把收到的消息原样回显
		_ = conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
		if err := conn.WriteMessage(mt, data); err != nil {
			log.Printf("write failed: %v", err)
			return
		}
	}
}

func main() {
	http.HandleFunc("/ws", wsHandler)
	log.Println("listening on :8080")
	if err := http.ListenAndServe(":8080", nil); err != nil {
		log.Fatal(err)
	}
}

验证用 wscat -c ws://127.0.0.1:8080/ws（需要 Node.js），或用 websocket.DefaultDialer.Dial 写一个 Go 客户端：握手失败时返回的 *http.Response 里带着状态码，能直接区分 401 与 403。

收发 API：什么时候用哪个

API	语义	适用
`ReadMessage`	读一条完整消息，自动拼装分片	绝大多数业务读
`NextReader` / `NextWriter`	流式读写单条消息	大消息，避免整块进内存
`WriteMessage`	写一条完整消息	绝大多数业务写
`ReadJSON` / `WriteJSON`	在消息之上做 JSON 编解码	控制类小消息
`WriteControl`	写 close/ping/pong 控制帧，可与其他写并发	心跳、关闭握手
`WritePreparedMessage`	预编码消息，可反复广播	高频广播同一条内容

写超时的规律很简单：任何写操作前都先 SetWriteDeadline。没设超时时，对端不读、TCP 发送缓冲写满，WriteMessage 就会一直阻塞，慢客户端能把写协程钉死（第 10 篇背压）。

错误处理：在 `ReadMessage` 返回 err 之后

switch {
case websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway):
	log.Printf("对端正常关闭")
case websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway):
	log.Printf("异常关闭: %v", err)
default:
	// 非关闭帧错误：读超时（net.Error.Timeout() 为真）说明对方没心跳，
	// 连接重置说明对端进程已消失
	log.Printf("读错误: %v", err)
}

读错误要分三类对待：对端发来 close 帧（拿到关闭码，可统计被踢、超长消息）、本地读超时、连接被重置。三类混在一起，线上就只剩"连接数莫名其妙下降"。

常见坑与排查

Upgrade 失败还去写响应：库里已经写过，会打日志 superfluous response.WriteHeader call。
不设 SetReadLimit：默认上限很大，单个客户端发一条超大消息就能吃掉可观内存。
多个 goroutine 同时 WriteMessage：直接 panic concurrent write to websocket connection，必须收敛到唯一的写协程（第 04 篇）。
CheckOrigin 一律放行、忘记 defer conn.Close()：前者配合 Cookie 鉴权会被跨站 WebSocket 劫持（CSWSH），务必白名单；后者会导致连接与文件句柄泄漏，进程最终撑爆 ulimit -n（第 17 篇）。

小结：Upgrader 负责握手与跨域校验，ReadMessage／WriteMessage 负责业务收发，ReadLimit 与读写 Deadline 负责保护自己，错误里必须读出关闭码。只要记住"一个连接同时只有一个写者"，剩下的复杂度都会落到下一章的读写分离模型里。

笔记加载中…

标签：#WebSocket #长连接 #IM

返回知识库