99 lines
1.9 KiB
Go
99 lines
1.9 KiB
Go
package server
|
|
|
|
import (
|
|
"log"
|
|
"sync"
|
|
|
|
"github.com/gofiber/websocket/v2"
|
|
)
|
|
|
|
type hubClient struct{}
|
|
|
|
type MessageHub struct {
|
|
Broadcast chan string
|
|
|
|
clients map[*websocket.Conn]hubClient
|
|
register chan *websocket.Conn
|
|
unregister chan *websocket.Conn
|
|
lastMessage string
|
|
|
|
running bool
|
|
runningMtx sync.Mutex
|
|
}
|
|
|
|
func NewHub() *MessageHub {
|
|
return &MessageHub{
|
|
clients: make(map[*websocket.Conn]hubClient),
|
|
register: make(chan *websocket.Conn),
|
|
Broadcast: make(chan string, 5),
|
|
unregister: make(chan *websocket.Conn),
|
|
}
|
|
}
|
|
|
|
func (hub *MessageHub) sendMessage(conn *websocket.Conn, message string) {
|
|
if err := conn.WriteMessage(
|
|
websocket.TextMessage, []byte(message)); err != nil {
|
|
log.Println("write error:", err)
|
|
|
|
_ = conn.WriteMessage(websocket.CloseMessage, []byte{})
|
|
_ = conn.Close()
|
|
delete(hub.clients, conn)
|
|
}
|
|
}
|
|
|
|
func (hub *MessageHub) Run() {
|
|
hub.runningMtx.Lock()
|
|
isRunning := hub.running
|
|
hub.running = true
|
|
hub.runningMtx.Unlock()
|
|
|
|
if isRunning {
|
|
return
|
|
}
|
|
|
|
for {
|
|
select {
|
|
case conn := <-hub.register:
|
|
hub.clients[conn] = hubClient{}
|
|
log.Println("connection registered")
|
|
|
|
case message := <-hub.Broadcast:
|
|
log.Println("message received:", message)
|
|
hub.lastMessage = message
|
|
|
|
// Send the message to all clients
|
|
for conn := range hub.clients {
|
|
hub.sendMessage(conn, message)
|
|
}
|
|
|
|
case conn := <-hub.unregister:
|
|
// Remove the client from the hub
|
|
delete(hub.clients, conn)
|
|
|
|
log.Println("connection unregistered")
|
|
}
|
|
}
|
|
}
|
|
|
|
func (hub *MessageHub) Handler(conn *websocket.Conn) {
|
|
// When the function returns, unregister the client and close the connection
|
|
defer func() {
|
|
hub.unregister <- conn
|
|
conn.Close()
|
|
}()
|
|
|
|
// Register the client
|
|
hub.register <- conn
|
|
|
|
if hub.lastMessage != "" {
|
|
hub.sendMessage(conn, hub.lastMessage)
|
|
}
|
|
|
|
for {
|
|
_, _, err := conn.ReadMessage()
|
|
if err != nil {
|
|
return // Calls the deferred function, i.e. closes the connection on error
|
|
}
|
|
}
|
|
}
|