WebSocket #

一、WebSocket概述 #

1.1 什么是WebSocket #

WebSocket是一种在单个TCP连接上进行全双工通信的协议，适合实时应用。

1.2 应用场景 #

实时聊天
在线游戏
实时数据推送
协作编辑

二、基本使用 #

2.1 安装依赖 #

bash

go get github.com/gofiber/fiber/v2
go get github.com/gofiber/contrib/websocket

2.2 基本WebSocket #

package main

import (
    "log"
    
    "github.com/gofiber/fiber/v2"
    "github.com/gofiber/contrib/websocket"
)

func main() {
    app := fiber.New()
    
    app.Get("/ws", websocket.New(func(c *websocket.Conn) {
        var msg string
        
        for {
            // 读取消息
            mt, message, err := c.ReadMessage()
            if err != nil {
                break
            }
            
            log.Printf("Received: %s", message)
            
            // 发送消息
            err = c.WriteMessage(mt, message)
            if err != nil {
                break
            }
        }
    }))
    
    app.Listen(":3000")
}

三、连接管理 #

3.1 连接池 #

type Client struct {
    ID   string
    Conn *websocket.Conn
}

type Hub struct {
    Clients    map[string]*Client
    Register   chan *Client
    Unregister chan *Client
    Broadcast  chan []byte
}

func NewHub() *Hub {
    return &Hub{
        Clients:    make(map[string]*Client),
        Register:   make(chan *Client),
        Unregister: make(chan *Client),
        Broadcast:  make(chan []byte),
    }
}

func (h *Hub) Run() {
    for {
        select {
        case client := <-h.Register:
            h.Clients[client.ID] = client
            
        case client := <-h.Unregister:
            if _, ok := h.Clients[client.ID]; ok {
                delete(h.Clients, client.ID)
                client.Conn.Close()
            }
            
        case message := <-h.Broadcast:
            for _, client := range h.Clients {
                client.Conn.WriteMessage(websocket.TextMessage, message)
            }
        }
    }
}

3.2 使用Hub #

var hub = NewHub()

func main() {
    go hub.Run()
    
    app := fiber.New()
    
    app.Get("/ws/:id", websocket.New(func(c *websocket.Conn) {
        id := c.Params("id")
        
        client := &Client{
            ID:   id,
            Conn: c,
        }
        
        hub.Register <- client
        
        defer func() {
            hub.Unregister <- client
        }()
        
        for {
            mt, message, err := c.ReadMessage()
            if err != nil {
                break
            }
            
            hub.Broadcast <- message
        }
    }))
    
    app.Listen(":3000")
}

四、消息处理 #

4.1 JSON消息 #

type Message struct {
    Type    string      `json:"type"`
    Content interface{} `json:"content"`
}

app.Get("/ws", websocket.New(func(c *websocket.Conn) {
    for {
        mt, message, err := c.ReadMessage()
        if err != nil {
            break
        }
        
        var msg Message
        json.Unmarshal(message, &msg)
        
        switch msg.Type {
        case "chat":
            // 处理聊天消息
            hub.Broadcast <- message
            
        case "ping":
            // 心跳响应
            c.WriteMessage(mt, []byte(`{"type":"pong"}`))
        }
    }
}))

4.2 广播消息 #

func BroadcastMessage(hub *Hub, message []byte) {
    hub.Broadcast <- message
}

func BroadcastJSON(hub *Hub, msg interface{}) {
    data, _ := json.Marshal(msg)
    hub.Broadcast <- data
}

五、聊天室示例 #

package main

import (
    "encoding/json"
    "log"
    
    "github.com/gofiber/fiber/v2"
    "github.com/gofiber/contrib/websocket"
)

type Client struct {
    ID   string
    Name string
    Conn *websocket.Conn
    Hub  *Hub
}

type Hub struct {
    Clients    map[string]*Client
    Register   chan *Client
    Unregister chan *Client
    Broadcast  chan []byte
}

type Message struct {
    Type    string `json:"type"`
    Sender  string `json:"sender"`
    Content string `json:"content"`
}

func NewHub() *Hub {
    return &Hub{
        Clients:    make(map[string]*Client),
        Register:   make(chan *Client),
        Unregister: make(chan *Client),
        Broadcast:  make(chan []byte),
    }
}

func (h *Hub) Run() {
    for {
        select {
        case client := <-h.Register:
            h.Clients[client.ID] = client
            h.Broadcast <- []byte(`{"type":"join","sender":"` + client.Name + `"}`)
            
        case client := <-h.Unregister:
            if _, ok := h.Clients[client.ID]; ok {
                delete(h.Clients, client.ID)
                client.Conn.Close()
                h.Broadcast <- []byte(`{"type":"leave","sender":"` + client.Name + `"}`)
            }
            
        case message := <-h.Broadcast:
            for _, client := range h.Clients {
                client.Conn.WriteMessage(websocket.TextMessage, message)
            }
        }
    }
}

var hub = NewHub()

func main() {
    go hub.Run()
    
    app := fiber.New()
    
    app.Get("/ws/:id/:name", websocket.New(func(c *websocket.Conn) {
        id := c.Params("id")
        name := c.Params("name")
        
        client := &Client{
            ID:   id,
            Name: name,
            Conn: c,
            Hub:  hub,
        }
        
        hub.Register <- client
        
        defer func() {
            hub.Unregister <- client
        }()
        
        for {
            _, message, err := c.ReadMessage()
            if err != nil {
                break
            }
            
            var msg Message
            json.Unmarshal(message, &msg)
            msg.Sender = name
            
            data, _ := json.Marshal(msg)
            hub.Broadcast <- data
        }
    }))
    
    log.Fatal(app.Listen(":3000"))
}

六、总结 #

6.1 WebSocket要点 #

要点	说明
连接管理	维护客户端连接
消息处理	解析和处理消息
广播	向所有客户端发送消息
错误处理	处理断开连接

6.2 下一步 #

现在你已经掌握了WebSocket，接下来让我们学习限流控制，了解如何保护API！