mirror of https://gitee.com/topnuomi/goweb
70 lines
1.5 KiB
Go
70 lines
1.5 KiB
Go
package core
|
|
|
|
import (
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
// WebSocketConnManager 管理WebSocket连接
|
|
type WebSocketConnManager struct {
|
|
ConnMap map[string]*websocket.Conn
|
|
}
|
|
|
|
// 创建一个新的WebSocketConnManager
|
|
func NewWebSocketConnManager() *WebSocketConnManager {
|
|
return &WebSocketConnManager{
|
|
ConnMap: make(map[string]*websocket.Conn),
|
|
}
|
|
}
|
|
|
|
// 添加连接
|
|
func (ws *WebSocketConnManager) AddConn(id string, conn *websocket.Conn) {
|
|
ws.ConnMap[id] = conn
|
|
}
|
|
|
|
// 根据ID移除连接
|
|
func (ws *WebSocketConnManager) RemoveConn(id string) {
|
|
delete(ws.ConnMap, id)
|
|
}
|
|
|
|
// 根据ID获取连接
|
|
func (ws *WebSocketConnManager) GetConn(id string) (*websocket.Conn, bool) {
|
|
conn, ok := ws.ConnMap[id]
|
|
return conn, ok
|
|
}
|
|
|
|
// 向所有连接发送消息
|
|
func (ws *WebSocketConnManager) SendMessageToAll(message []byte) {
|
|
for _, conn := range ws.ConnMap {
|
|
err := conn.WriteMessage(websocket.TextMessage, message)
|
|
if err != nil {
|
|
continue
|
|
}
|
|
}
|
|
}
|
|
|
|
// 向特定用户发送消息
|
|
func (ws *WebSocketConnManager) SendMessageToUser(id string, message []byte) {
|
|
conn, ok := ws.GetConn(id)
|
|
if !ok {
|
|
return
|
|
}
|
|
err := conn.WriteMessage(websocket.TextMessage, message)
|
|
if err != nil {
|
|
return
|
|
}
|
|
}
|
|
|
|
// 向多个用户发送消息
|
|
func (ws *WebSocketConnManager) SendMessageToUsers(ids []string, message []byte) {
|
|
for _, id := range ids {
|
|
conn, ok := ws.GetConn(id)
|
|
if !ok {
|
|
continue
|
|
}
|
|
err := conn.WriteMessage(websocket.TextMessage, message)
|
|
if err != nil {
|
|
continue
|
|
}
|
|
}
|
|
}
|