package initialize import ( "confrontation-training/global" "encoding/json" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" log "github.com/sirupsen/logrus" "net/http" ) // ClientManager is a websocket manager type ClientManager struct { Clients map[string]*Client Broadcast chan []byte Register chan *Client Unregister chan *Client } // Client is a websocket client type Client struct { ID string 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[string]*Client), } // Start is 项目运行前, 协程开启start -> go Manager.Start() func (manager *ClientManager) Start() { for { log.Infoln("<---管道通信--->") select { case conn := <-Manager.Register: log.Infof("新用户加入:%v", conn.ID) Manager.Clients[conn.ID] = conn jsonMessage, _ := json.Marshal(&Message{Content: "Successful connection to socket service"}) conn.Send <- jsonMessage case conn := <-Manager.Unregister: log.Infof("用户离开:%v", conn.ID) if _, ok := Manager.Clients[conn.ID]; ok { jsonMessage, _ := json.Marshal(&Message{Content: "A socket has disconnected"}) conn.Send <- jsonMessage close(conn.Send) delete(Manager.Clients, conn.ID) } case message := <-Manager.Broadcast: MessageStruct := Message{} err := json.Unmarshal(message, &MessageStruct) if err != nil { panic(err) return } for id, conn := range Manager.Clients { if id != creatId(MessageStruct.Recipient, MessageStruct.Sender) { continue } else { select { case conn.Send <- message: default: close(conn.Send) delete(Manager.Clients, conn.ID) } } } } } } func creatId(uid, touid string) string { return uid + "_" + touid } func (c *Client) Read() { //defer func() { // Manager.Unregister <- c // c.Socket.Close() //}() for { c.Socket.PongHandler() _, message, err := c.Socket.ReadMessage() if err != nil { Manager.Unregister <- c c.Socket.Close() break } //log.Infof("读取到客户端的信息:%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 } global.Log4J.Info("发送到到客户端的信息:", string(message)) c.Socket.WriteMessage(websocket.TextMessage, message) } } } // WsHandler TestHandler socket 连接 中间件 作用:升级协议,用户验证,自定义信息等 func WsHandler(c *gin.Context) { uid := c.Query("uid") touid := c.Query("to_uid") conn, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil) if err != nil { http.NotFound(c.Writer, c.Request) return } //可以添加用户信息验证 client := &Client{ ID: creatId(uid, touid), Socket: conn, Send: make(chan []byte), } Manager.Register <- client go client.Read() go client.Write() }