diff --git a/pkg/server/messages_ws.go b/pkg/server/messages_ws.go index a51475b..5fdcb58 100644 --- a/pkg/server/messages_ws.go +++ b/pkg/server/messages_ws.go @@ -5,14 +5,13 @@ import ( "git.solsynth.dev/hydrogen/messaging/pkg/services" "github.com/gofiber/contrib/websocket" jsoniter "github.com/json-iterator/go" - "github.com/samber/lo" ) func messageGateway(c *websocket.Conn) { user := c.Locals("principal").(models.Account) // Push connection - services.WsConn[user.ID] = append(services.WsConn[user.ID], c) + services.ClientRegister(user, c) // Event loop var task models.UnifiedCommand @@ -42,7 +41,5 @@ func messageGateway(c *websocket.Conn) { } // Pop connection - services.WsConn[user.ID] = lo.Filter(services.WsConn[user.ID], func(item *websocket.Conn, idx int) bool { - return item != c - }) + services.ClientUnregister(user, c) } diff --git a/pkg/services/connections.go b/pkg/services/connections.go index 53117a1..3c413d0 100644 --- a/pkg/services/connections.go +++ b/pkg/services/connections.go @@ -5,14 +5,24 @@ import ( "github.com/gofiber/contrib/websocket" ) -var WsConn = make(map[uint][]*websocket.Conn) +var wsConn = make(map[uint]map[*websocket.Conn]bool) -func CheckOnline(user models.Account) bool { - return len(WsConn[user.ID]) > 0 +func ClientRegister(user models.Account, conn *websocket.Conn) { + if wsConn[user.ID] == nil { + wsConn[user.ID] = make(map[*websocket.Conn]bool) + } + wsConn[user.ID][conn] = true +} + +func ClientUnregister(user models.Account, conn *websocket.Conn) { + if wsConn[user.ID] == nil { + wsConn[user.ID] = make(map[*websocket.Conn]bool) + } + delete(wsConn[user.ID], conn) } func PushCommand(userId uint, task models.UnifiedCommand) { - for _, conn := range WsConn[userId] { + for conn := range wsConn[userId] { _ = conn.WriteMessage(1, task.Marshal()) } }