什么是 WebSocket？

WebSocket 是一种在单个 TCP 连接上进行全双工通信的协议。它于 2011 年被 IETF 定为标准 RFC 6455，设计目标是在 Web 浏览器和服务器之间提供实时、双向的数据传输能力。

为什么需要 WebSocket？

核心特性

• 🔄 全双工通信：客户端和服务器可以同时发送和接收数据
• 🌐 基于 TCP：建立在 TCP 协议之上，默认端口 80 (WS) 或 443 (WSS)
• 🔗 持久连接：连接建立后保持开放状态，避免重复握手
• 📦 轻量级：协议头部小，数据传输效率高
• 🎯 跨域支持：天然支持跨域通信
• 🔐 安全支持：支持 WSS (WebSocket Secure) 加密传输

协议层次

WebSocket 握手流程

WebSocket 连接建立通过 HTTP 升级机制完成，客户端发送一个带有特殊头部字段的 HTTP 请求。

客户端请求示例

GET /chat HTTP/1.1
Host: server.example.com
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
Origin: http://example.com
Sec-WebSocket-Version: 13

服务器响应示例

HTTP/1.1 101 Switching Protocols
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=

关键头部字段说明

WebSocket 帧格式

WebSocket 数据传输使用帧（Frame）作为基本单位，每个帧包含头部和负载两部分。

字段详解

消息分片示例

大消息可以被拆分成多个帧传输：

控制帧

控制帧用于连接管理，不能分片：

• Close (0x8): 优雅关闭连接，可包含状态码和原因
• Ping (0x9): 心跳检测，期望收到 Pong 响应
• Pong (0xA): 对 Ping 的响应

服务器端实现

使用 websockets 库创建简单的 Echo 服务器：

# 安装依赖: pip install websockets
import asyncio
import websockets
import json

async def echo_server(websocket):
    """处理客户端连接"""
    print(f"客户端已连接: {websocket.remote_address}")
    
    try:
        async for message in websocket:
            # 解析消息
            try:
                data = json.loads(message)
                print(f"收到消息: {data}")
                
                # 处理消息
                response = {
                    "type": "echo",
                    "data": data,
                    "timestamp": asyncio.get_event_loop().time()
                }
                
                # 发送响应
                await websocket.send(json.dumps(response))
                
            except json.JSONDecodeError:
                # 非JSON消息，直接回显
                await websocket.send(f"Echo: {message}")
                
    except websockets.exceptions.ConnectionClosed:
        print(f"客户端断开连接: {websocket.remote_address}")

async def main():
    """启动服务器"""
    print("WebSocket 服务器启动在 ws://localhost:8765")
    
    async with websockets.serve(echo_server, "localhost", 8765):
        await asyncio.Future()  # 永久运行

if __name__ == "__main__":
    asyncio.run(main())

客户端实现

import asyncio
import websockets
import json

async def websocket_client():
    """WebSocket 客户端示例"""
    uri = "ws://localhost:8765"
    
    try:
        async with websockets.connect(uri) as websocket:
            print(f"已连接到服务器: {uri}")
            
            # 发送消息
            messages = [
                {"type": "greeting", "message": "Hello, Server!"},
                {"type": "data", "value": 42},
                {"type": "ping", "timestamp": 1234567890}
            ]
            
            for msg in messages:
                await websocket.send(json.dumps(msg))
                print(f"发送: {msg}")
                
                # 接收响应
                response = await websocket.recv()
                print(f"接收: {response}\n")
                
                # 短暂延迟
                await asyncio.sleep(1)
                
    except Exception as e:
        print(f"连接错误: {e}")

if __name__ == "__main__":
    asyncio.run(websocket_client())

广播服务器

实现消息广播功能，向所有连接的客户端发送消息：

import asyncio
import websockets
import json

# 存储所有连接的客户端
connected_clients = set()

async def broadcast(message):
    """向所有客户端广播消息"""
    if connected_clients:
        websockets.broadcast(connected_clients, message)
        print(f"广播消息给 {len(connected_clients)} 个客户端")

async def handle_client(websocket):
    """处理客户端连接"""
    connected_clients.add(websocket)
    print(f"客户端连接: {websocket.remote_address}, 总数: {len(connected_clients)}")
    
    try:
        async for message in websocket:
            data = json.loads(message)
            print(f"收到消息: {data}")
            
            # 广播给所有客户端
            await broadcast(json.dumps({
                "type": "broadcast",
                "from": str(websocket.remote_address),
                "data": data
            }))
            
    except websockets.exceptions.ConnectionClosed:
        print(f"客户端断开: {websocket.remote_address}")
    finally:
        connected_clients.remove(websocket)
        print(f"客户端移除, 剩余: {len(connected_clients)}")

async def main():
    """启动广播服务器"""
    print("广播服务器启动在 ws://localhost:8765")
    
    async with websockets.serve(handle_client, "localhost", 8765):
        await asyncio.Future()

if __name__ == "__main__":
    asyncio.run(main())

服务器端实现

使用 gorilla/websocket 库创建 WebSocket 服务器：

// 安装依赖: go get github.com/gorilla/websocket
package main

import (
    "encoding/json"
    "log"
    "net/http"
    "time"

    "github.com/gorilla/websocket"
)

// WebSocket 升级器
var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool {
        return true // 生产环境需要验证 Origin
    },
}

// 消息结构
type Message struct {
    Type      string      `json:"type"`
    Data      interface{} `json:"data"`
    Timestamp float64     `json:"timestamp"`
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
    // 升级 HTTP 连接为 WebSocket
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Printf("升级失败: %v", err)
        return
    }
    defer conn.Close()

    log.Printf("客户端已连接: %s", conn.RemoteAddr())

    for {
        // 读取消息
        messageType, message, err := conn.ReadMessage()
        if err != nil {
            log.Printf("读取错误: %v", err)
            break
        }

        log.Printf("收到消息: %s", string(message))

        // 解析 JSON 消息
        var msg Message
        if err := json.Unmarshal(message, &msg); err == nil {
            // 构造响应
            response := Message{
                Type:      "echo",
                Data:      msg.Data,
                Timestamp: float64(time.Now().UnixNano()) / 1e9,
            }

            responseBytes, _ := json.Marshal(response)
            
            // 发送响应
            if err := conn.WriteMessage(messageType, responseBytes); err != nil {
                log.Printf("发送错误: %v", err)
                break
            }
        } else {
            // 非 JSON 消息，直接回显
            if err := conn.WriteMessage(messageType, message); err != nil {
                log.Printf("发送错误: %v", err)
                break
            }
        }
    }

    log.Printf("客户端断开连接: %s", conn.RemoteAddr())
}

func main() {
    http.HandleFunc("/ws", handleWebSocket)
    
    log.Println("WebSocket 服务器启动在 ws://localhost:8080/ws")
    if err := http.ListenAndServe(":8080", nil); err != nil {
        log.Fatal("服务器启动失败: ", err)
    }
}

客户端实现

package main

import (
    "encoding/json"
    "log"
    "net/url"
    "time"

    "github.com/gorilla/websocket"
)

type Message struct {
    Type      string      `json:"type"`
    Data      interface{} `json:"data"`
    Timestamp float64     `json:"timestamp"`
}

func main() {
    // 建立 WebSocket 连接
    u := url.URL{Scheme: "ws", Host: "localhost:8080", Path: "/ws"}
    log.Printf("连接到服务器: %s", u.String())

    conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
    if err != nil {
        log.Fatal("连接失败: ", err)
    }
    defer conn.Close()

    // 发送消息
    messages := []Message{
        {Type: "greeting", Data: "Hello, Server!"},
        {Type: "data", Data: 42},
        {Type: "ping", Data: time.Now().Unix()},
    }

    for _, msg := range messages {
        // 序列化消息
        messageBytes, err := json.Marshal(msg)
        if err != nil {
            log.Printf("序列化错误: %v", err)
            continue
        }

        // 发送消息
        if err := conn.WriteMessage(websocket.TextMessage, messageBytes); err != nil {
            log.Printf("发送错误: %v", err)
            break
        }

        log.Printf("发送: %+v", msg)

        // 接收响应
        _, response, err := conn.ReadMessage()
        if err != nil {
            log.Printf("接收错误: %v", err)
            break
        }

        log.Printf("接收: %s\n", string(response))
        time.Sleep(1 * time.Second)
    }
}

Hub 模式 - 广播服务器

package main

import (
    "encoding/json"
    "log"
    "net/http"
    "sync"

    "github.com/gorilla/websocket"
)

// Hub 管理所有客户端连接
type Hub struct {
    clients    map[*Client]bool
    broadcast  chan []byte
    register   chan *Client
    unregister chan *Client
    mu         sync.RWMutex
}

// Client 表示一个 WebSocket 客户端
type Client struct {
    hub  *Hub
    conn *websocket.Conn
    send chan []byte
}

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

// NewHub 创建新的 Hub
func NewHub() *Hub {
    return &Hub{
        clients:    make(map[*Client]bool),
        broadcast:  make(chan []byte),
        register:   make(chan *Client),
        unregister: make(chan *Client),
    }
}

// Run 运行 Hub
func (h *Hub) Run() {
    for {
        select {
        case client := <-h.register:
            h.mu.Lock()
            h.clients[client] = true
            h.mu.Unlock()
            log.Printf("客户端连接, 总数: %d", len(h.clients))

        case client := <-h.unregister:
            h.mu.Lock()
            if _, ok := h.clients[client]; ok {
                delete(h.clients, client)
                close(client.send)
            }
            h.mu.Unlock()
            log.Printf("客户端断开, 剩余: %d", len(h.clients))

        case message := <-h.broadcast:
            h.mu.RLock()
            for client := range h.clients {
                select {
                case client.send <- message:
                default:
                    close(client.send)
                    delete(h.clients, client)
                }
            }
            h.mu.RUnlock()
        }
    }
}

// readPump 从客户端读取消息
func (c *Client) readPump() {
    defer func() {
        c.hub.unregister <- c
        c.conn.Close()
    }()

    for {
        _, message, err := c.conn.ReadMessage()
        if err != nil {
            break
        }

        // 广播消息
        var msg map[string]interface{}
        if err := json.Unmarshal(message, &msg); err == nil {
            msg["from"] = c.conn.RemoteAddr().String()
            broadcastMsg, _ := json.Marshal(msg)
            c.hub.broadcast <- broadcastMsg
        }
    }
}

// writePump 向客户端发送消息
func (c *Client) writePump() {
    defer c.conn.Close()

    for message := range c.send {
        err := c.conn.WriteMessage(websocket.TextMessage, message)
        if err != nil {
            break
        }
    }
}

var hub = NewHub()

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Printf("升级失败: %v", err)
        return
    }

    client := &Client{
        hub:  hub,
        conn: conn,
        send: make(chan []byte, 256),
    }

    hub.register <- client

    go client.writePump()
    go client.readPump()
}

func main() {
    go hub.Run()

    http.HandleFunc("/ws", handleWebSocket)
    
    log.Println("广播服务器启动在 ws://localhost:8080/ws")
    if err := http.ListenAndServe(":8080", nil); err != nil {
        log.Fatal("服务器启动失败: ", err)
    }
}

性能优化

• 使用连接池管理 WebSocket 连接
• 实现心跳机制检测连接状态
• 合理设置消息缓冲区大小
• 使用二进制帧传输大数据
• 启用消息压缩（permessage-deflate）

安全考虑

• 生产环境使用 WSS (WebSocket Secure)
• 验证 Origin 头部防止 CSRF 攻击
• 实现连接数限制和速率限制
• 输入验证，防止注入攻击
• 定期更新依赖库

错误处理

• 实现优雅的连接关闭机制
• 处理网络中断和重连逻辑
• 记录详细的错误日志
• 设置合理的超时时间
• 实现状态码和错误消息规范