apifiber/controllers/websocketController.go

120 lines
2.7 KiB
Go

package controllers
import (
"encoding/json"
"log"
"sync"
"github.com/gofiber/websocket/v2"
)
// WebSocketManager handles all active WebSocket connections
type WebSocketManager struct {
connections map[*websocket.Conn]bool
mutex sync.RWMutex
}
// WebSocketMessage represents the structure of messages
type WebSocketMessage struct {
Command string `json:"command"`
Channel string `json:"channel"`
Text string `json:"text"`
}
// Global instance of WebSocketManager
var WSManager = &WebSocketManager{
connections: make(map[*websocket.Conn]bool),
}
// BroadcastMessage sends a message to all connected clients
func (m *WebSocketManager) BroadcastMessage(command, channel, text string) {
message := WebSocketMessage{
Command: command,
Channel: channel,
Text: text,
}
jsonMessage, err := json.Marshal(message)
if err != nil {
log.Printf("Error marshaling message: %v", err)
return
}
m.mutex.RLock()
defer m.mutex.RUnlock()
for conn := range m.connections {
if err := conn.WriteMessage(websocket.TextMessage, jsonMessage); err != nil {
log.Printf("Error broadcasting message: %v", err)
// Remove failed connection
m.mutex.RUnlock()
m.removeConnection(conn)
m.mutex.RLock()
}
}
}
// SendMessageToClient sends a message to a specific client
func (m *WebSocketManager) SendMessageToClient(conn *websocket.Conn, command, channel, text string) error {
message := WebSocketMessage{
Command: command,
Channel: channel,
Text: text,
}
jsonMessage, err := json.Marshal(message)
if err != nil {
return err
}
m.mutex.RLock()
defer m.mutex.RUnlock()
if _, exists := m.connections[conn]; exists {
return conn.WriteMessage(websocket.TextMessage, jsonMessage)
}
return nil
}
func (m *WebSocketManager) addConnection(conn *websocket.Conn) {
m.mutex.Lock()
defer m.mutex.Unlock()
m.connections[conn] = true
}
func (m *WebSocketManager) removeConnection(conn *websocket.Conn) {
m.mutex.Lock()
defer m.mutex.Unlock()
delete(m.connections, conn)
conn.Close()
}
// WebsocketHandler handles WebSocket connections
func WebsocketHandler(c *websocket.Conn) {
// Add the connection to our manager
WSManager.addConnection(c)
defer WSManager.removeConnection(c)
// Handle incoming messages
for {
_, message, err := c.ReadMessage()
if err != nil {
log.Println("Error reading message:", err)
break
}
// Try to parse the incoming message as JSON
var wsMessage WebSocketMessage
if err := json.Unmarshal(message, &wsMessage); err != nil {
log.Printf("Error parsing message: %v", err)
continue
}
// Echo the message back to the client
if err := WSManager.SendMessageToClient(c, wsMessage.Command, wsMessage.Channel, wsMessage.Text); err != nil {
log.Println("Error writing message:", err)
break
}
}
}