Go 标准库 + gorilla/websocket：120 行实现 IM Mock Server

基于 YunRong 项目 Mock Server 的实际开发过程。零框架依赖，仅 gorilla/websocket 一个外部库。

业务场景

IM 客户端开发需要一个可控的后端来验证网络层代码。真实后端开发周期长、依赖 PostgreSQL 等服务，不适合在客户端开发阶段使用。Mock Server 的目标是：

用最少代码模拟登录鉴权、消息收发
客户端开发时不依赖真实后端
服务器行为完全可控（延时、丢包、异常响应均可注入）

最终实现：120 行 Go 代码完成 HTTP 登录 + JWT 签发 + WebSocket 回声 + Hub 单聊路由。

完整代码：server/cmd/server/main.go

第 1 步：最小 HTTP 服务

Go 1.22 的 http.NewServeMux 支持方法路由，不再需要第三方路由器：

mux := http.NewServeMux()
mux.HandleFunc("GET /health", func(w http.ResponseWriter, r *http.Request) {
    w.Header().Set("Content-Type", "application/json")
    w.Write([]byte(`{"status":"ok"}`))
})

http.ListenAndServe(":8080", mux)

Go 1.22 的 HandleFunc 新语法："GET /health" 和 "POST /api/v1/auth/login" 在路由注册时就绑定了 HTTP 方法。这消除了以前在每个 handler 里 if r.Method != "POST" 的样板代码。

// Go 1.22+ 写法
mux.HandleFunc("POST /api/v1/auth/login", loginHandler)

// 旧写法（Go 1.21 及以前）
mux.HandleFunc("/api/v1/auth/login", func(w http.ResponseWriter, r *http.Request) {
    if r.Method != "POST" {
        http.Error(w, "method not allowed", 405)
        return
    }
    loginHandler(w, r)
})

第 2 步：JWT 鉴权（标准库手写）

真实项目会用 golang-jwt/jwt，但 Mock Server 只有 3 个用户、1 个 secret，引一个 2000 行的库不值得。JWT 的结构其实很简单：

base64(header).base64(payload).hmac_sha256(header.payload, secret)

三步实现：

var jwtSecret = []byte("yunrong-mock-secret")

func makeJWT(userID int64, username string) (string, error) {
    // 1. Header
    header := base64.RawURLEncoding.EncodeToString(
        []byte(`{"alg":"HS256","typ":"JWT"}`))

    // 2. Payload
    payloadBytes, _ := json.Marshal(map[string]any{
        "user_id":  userID,
        "username": username,
        "exp":      time.Now().Add(2 * time.Hour).Unix(),
        "iat":      time.Now().Unix(),
    })
    payload := base64.RawURLEncoding.EncodeToString(payloadBytes)

    // 3. Signature
    unsigned := header + "." + payload
    mac := hmac.New(sha256.New, jwtSecret)
    mac.Write([]byte(unsigned))
    sig := base64.RawURLEncoding.EncodeToString(mac.Sum(nil))

    return unsigned + "." + sig, nil
}

三个要点：

RawURLEncoding 而非 StdEncoding——JWT 要求 URL-safe Base64（用 - 和 _ 替代 + 和 /）
makeJWT 返回的 token 可以用 jwt.io 直接解码验证——这就是验收标准中”有效的 JWT”的含义
Mock 阶段的 parseTokenUserID 只解码不验签——因为 secret 只有服务端知道，验签需要和签发者共享 secret

func parseTokenUserID(token string) int64 {
    parts := strings.Split(token, ".")
    if len(parts) != 3 { return 0 }         // 格式不对

    payload, err := base64.RawURLEncoding.DecodeString(parts[1])
    if err != nil { return 0 }               // Base64 解码失败

    var claims map[string]any
    json.Unmarshal(payload, &claims)         // JSON 解码

    id, _ := claims["user_id"].(float64)     // JSON 数字默认 float64
    return int64(id)
}

第 3 步：WebSocket 回声

gorilla/websocket 的使用分两步：升级和收发。

3.1 HTTP → WebSocket 升级

var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool { return true },
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Printf("ws upgrade failed: %v", err)
        return
    }
    defer conn.Close()

    for {
        _, msg, err := conn.ReadMessage()
        if err != nil {
            log.Printf("ws disconnected: %v", err)
            return
        }
        // 处理 msg...
    }
}

CheckOrigin 陷阱：返回 true 意味着允许任意来源连接。Mock Server 可以这样，但生产环境必须校验 Origin 头防止 CSWSH 攻击。

3.2 回声逻辑

for {
    _, msg, err := conn.ReadMessage()
    if err != nil { return }

    var req map[string]any
    json.Unmarshal(msg, &req)

    msgType, _ := req["type"].(string)
    switch msgType {
    case "ping":
        conn.WriteJSON(map[string]any{"type": "pong"})
    default:
        conn.WriteJSON(map[string]any{"type": "echo", "payload": req})
    }
}

WriteJSON vs WriteMessage：WriteJSON 自动序列化为 JSON，省去手动 json.Marshal + WriteMessage。

第 4 步：Hub 单聊路由

回声只能验证”收发链路通”，IM 需要的是两个客户端之间的消息路由。这是 Mock Server 从”玩具”变成”可用”的关键一步。

4.1 数据结构

type Hub struct {
    mu      sync.RWMutex
    clients map[int64]*Client   // userID → 连接
}

type Client struct {
    conn   *websocket.Conn
    userID int64
}

为什么用 map[int64]*Client 而非 map[int64][]*Client？

Mock 阶段每个用户只允许一个连接——新连接会踢掉旧连接。这样做简化了路由逻辑，也避免了”发给哪个设备”的歧义。

func (h *Hub) register(c *Client) {
    h.mu.Lock()
    defer h.mu.Unlock()
    if old, ok := h.clients[c.userID]; ok {
        old.conn.Close()  // 踢掉旧连接
    }
    h.clients[c.userID] = c
    log.Printf("user %d registered (%d online)", c.userID, len(h.clients))
}

4.2 消息路由

case "msg":
    toUser := toFloat64(msg["to"])
    body, _ := msg["body"].(string)
    if toUser > 0 && body != "" {
        forward := map[string]any{
            "type": "msg",
            "from": userID,
            "body": body,
        }
        forwardBytes, _ := json.Marshal(forward)
        hub.sendToUser(int64(toUser), forwardBytes)
    }

toFloat64 的必要性：Go 的 json.Unmarshal 到 map[string]any 时，所有 JSON 数字都变成 float64。客户端发来的 "to": 1002 会被解析为 1002.0，必须用 float64 类型断言才能取到。

func toFloat64(v any) float64 {
    switch n := v.(type) {
    case float64:  return n              // 最常见：json.Unmarshal 产生
    case json.Number:                    // json.Decoder 产生
        f, _ := n.Float64()
        return f
    }
    return 0
}

4.3 完整消息流

A (1001) 发送 {"type":"msg","to":1002,"body":"hello"}
   │
   ▼
wsHandler → json.Unmarshal → switch "msg"
   │
   ▼
hub.sendToUser(1002, forwardBytes)
   │
   ▼
hub.mu.RLock() → clients[1002].conn.WriteMessage(...)
   │
   ▼
B (1002) 收到 {"type":"msg","from":1001,"body":"hello"}

踩坑：GOPROXY 被墙

go mod tidy
# go: module github.com/gorilla/websocket: 
# Get "https://proxy.golang.org/...": dial tcp 142.250.77.209:443: 
# connectex: A connection attempt failed

Go 默认从 proxy.golang.org 下载模块，国内无法访问。一行命令解决：

go env -w GOPROXY=https://goproxy.cn,direct

goproxy.cn 是七牛云维护的 Go 模块代理，同步延迟通常在 5 分钟以内。

调试技巧：`init()` 打印测试 token

func init() {
    tok, _ := makeJWT(1001, "admin")
    tok2, _ := makeJWT(1002, "zhangsan")
    log.Println(strings.Repeat("=", 60))
    log.Println("Test tokens (valid 2h):")
    log.Printf("  admin:    ?token=%s", tok)
    log.Printf("  zhangsan: ?token=%s", tok2)
    log.Println(strings.Repeat("=", 60))
}

服务器启动时自动打印两个测试用户的完整 token。直接复制粘贴到浏览器控制台就能用，省去先调 login 接口的步骤。

服务端并发安全速查

操作	锁	原因
`register()`	`Lock`	写 map
`unregister()`	`Lock`	写 map
`sendToUser()`	`RLock`	只读 map，写的是 conn（gorilla 内部线程安全）

只在 sendToUser 用 RLock，为什么 register/unregister 要 Lock？

因为 register 除了写 map 还调了 old.conn.Close()——这个副作用必须在写锁保护下执行，否则可能出现”旧连接被关闭的同时新连接在写入”的竞态。

总结对比：Mock Server vs 生产后端

维度	Mock Server（本实现）	生产后端
HTTP 路由	Go 1.22 标准 `mux`	Gin / Echo
JWT	`crypto/hmac` 手写	`golang-jwt/jwt`
WebSocket	单 goroutine 同步读写	goroutine-per-conn + channel
消息路由	`map[userID]*Client` + RWMutex	分布式消息队列 + Redis Pub/Sub
持久化	无	PostgreSQL
部署	`go run`	Docker + K8s

Mock Server 的价值不是替代生产后端，而是让客户端开发可以在后端不存在的情况下全速推进。120 行代码换来了 12 项客户端任务的独立验证能力——这是最划算的投资。

Go Mock Server 消息路由实战