实现功能: websocket+chan 通信

使用包:

github.com/gin-gonic/gin
github.com/gorilla/websocket

代码

package websocket

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

    "github.com/gin-gonic/gin"
    "github.com/gorilla/websocket"
)

// ClientManager is a websocket manager
type ClientManager struct {
    Clients    map[*Client]bool
    Broadcast  chan []byte
    Register   chan *Client
    Unregister chan *Client
}

// Client is a websocket client
type Client struct {
    ID     int
    Socket *websocket.Conn
    Send   chan []byte
}

// Message is return msg
type Message struct {
    Sender    string `json:"sender,omitempty"`
    Recipient string `json:"recipient,omitempty"`
    Content   string `json:"content,omitempty"`
}

// Manager define a ws server manager
var Manager = ClientManager{
    Broadcast:  make(chan []byte),
    Register:   make(chan *Client),
    Unregister: make(chan *Client),
    Clients:    make(map[*Client]bool),
}

// Start is  项目运行前, 协程开启start -> go Manager.Start()
func (manager *ClientManager) Start() {
    for {
        log.Println("<---管道通信--->")

        select {
        case conn := <-Manager.Register:
            log.Printf("新用户加入:%v", conn.ID)
            Manager.Clients[conn] = true
            jsonMessage, _ := json.Marshal(&Message{Content: "Successful connection to socket service"})
            Manager.Send(jsonMessage, conn)
        case conn := <-Manager.Unregister:
            log.Printf("用户离开:%v", conn.ID)
            if _, ok := Manager.Clients[conn]; ok {
                close(conn.Send)
                delete(Manager.Clients, conn)
                jsonMessage, _ := json.Marshal(&Message{Content: "A socket has disconnected"})
                Manager.Send(jsonMessage, conn)
            }

        case message := <-Manager.Broadcast:

            jsonMessage, _ := json.Marshal(&Message{Content: string(message)})
            for conn := range Manager.Clients {
                select {
                case conn.Send <- jsonMessage:
                default:
                    close(conn.Send)
                    delete(Manager.Clients, conn)
                }
            }
        }
    }
}

// Send is to send ws message to ws client
func (manager *ClientManager) Send(message []byte, ignore *Client) {

    for conn := range manager.Clients {
        // if conn != ignore { //向除了自己的socket 用户发送
        conn.Send <- message
        // }
    }
}

func (c *Client) Read() {
    defer func() {
        Manager.Unregister <- c
        c.Socket.Close()
    }()

    for {

        _, message, err := c.Socket.ReadMessage()
        if err != nil {
            Manager.Unregister <- c
            c.Socket.Close()
            break
        }
        log.Printf("读取到客户端的信息:%s", string(message))
        Manager.Broadcast <- message
    }
}

func (c *Client) Write() {
    defer func() {
        c.Socket.Close()
    }()

    for {
        select {
        case message, ok := <-c.Send:
            if !ok {
                c.Socket.WriteMessage(websocket.CloseMessage, []byte{})
                return
            }
            log.Printf("发送到到客户端的信息:%s", string(message))

            c.Socket.WriteMessage(websocket.TextMessage, message)
        }
    }

}

//TestHandler socket 连接 中间件 作用:升级协议,用户验证,自定义信息等
func TestHandler(c *gin.Context) {
    conn, err := websocket.Upgrade(c.Writer, c.Request, nil, 1024, 1024)
    if err != nil {
        http.NotFound(c.Writer, c.Request)
        return
    }
    //可以添加用户信息验证
    userID := 1
    client := &Client{
        ID:     userID,
        Socket: conn,
        Send:   make(chan []byte),
    }
    Manager.Register <- client
    go client.Read()
    go client.Write()
}

使用

定义个ws请求路径路由 , 路由验证 使用TestHandler 中间件, 用于升级服务,验证用户信息
例如 :

 route.Get("/ws",TestHandler)

==main 入口 写入 go Manager.Start() 协程开启服务==
如果文件不在main包里, 需要写上包名 例如: 包名为websocket
go websocket.Manager.Start()

每个方法的作用

Start():启动websocket服务
Send():向连接websocket的管道chan写入数据
Read():读取在websocket管道中的数据
Write():通过websocket协议向连接到ws的客户端发送数据
TestHandler(): ws链接交互的中间件, 用于协议升级, 用户信息验证等

前端测试 我使用的是 http://easyswoole.com/wstool.html 这个测试工具

代码原处:https://www.jianshu.com/p/f058fdbdea58 只是稍微修改下代码