From 84de38de261415f9c4745ffd2c6d5977e4bac86e Mon Sep 17 00:00:00 2001 From: nilo Date: Wed, 21 May 2025 11:32:09 -0300 Subject: [PATCH] websocket first version --- controllers/websocketController.go | 119 +++++++++++++++++++++++++++++ go.mod | 3 + go.sum | 6 ++ routes/routes.go | 12 +++ 4 files changed, 140 insertions(+) create mode 100644 controllers/websocketController.go diff --git a/controllers/websocketController.go b/controllers/websocketController.go new file mode 100644 index 0000000..805087c --- /dev/null +++ b/controllers/websocketController.go @@ -0,0 +1,119 @@ +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 + } + } +} diff --git a/go.mod b/go.mod index 67ace5b..69c4d71 100644 --- a/go.mod +++ b/go.mod @@ -4,6 +4,7 @@ go 1.22 require ( github.com/gofiber/fiber/v2 v2.50.0 + github.com/gofiber/websocket/v2 v2.2.1 github.com/shirou/gopsutil/v3 v3.23.10 github.com/spf13/viper v1.17.0 golang.org/x/crypto v0.15.0 @@ -14,6 +15,7 @@ require ( require ( github.com/andybalholm/brotli v1.0.6 // indirect + github.com/fasthttp/websocket v1.5.3 // indirect github.com/fsnotify/fsnotify v1.7.0 // indirect github.com/go-ole/go-ole v1.3.0 // indirect github.com/google/uuid v1.4.0 // indirect @@ -36,6 +38,7 @@ require ( github.com/rivo/uniseg v0.4.4 // indirect github.com/sagikazarmark/locafero v0.3.0 // indirect github.com/sagikazarmark/slog-shim v0.1.0 // indirect + github.com/savsgio/gotils v0.0.0-20230208104028-c358bd845dee // indirect github.com/shoenig/go-m1cpu v0.1.6 // indirect github.com/sourcegraph/conc v0.3.0 // indirect github.com/spf13/afero v1.10.0 // indirect diff --git a/go.sum b/go.sum index 363f97a..bab32d8 100644 --- a/go.sum +++ b/go.sum @@ -58,6 +58,8 @@ github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1m github.com/envoyproxy/go-control-plane v0.9.7/go.mod h1:cwu0lG7PUMfa9snN8LXBig5ynNVH9qI8YYLbd1fK2po= github.com/envoyproxy/go-control-plane v0.9.9-0.20201210154907-fd9021fe5dad/go.mod h1:cXg6YxExXjJnVBQHBLXeUAgxn2UodCpnH306RInaBQk= github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= +github.com/fasthttp/websocket v1.5.3 h1:TPpQuLwJYfd4LJPXvHDYPMFWbLjsT91n3GpWtCQtdek= +github.com/fasthttp/websocket v1.5.3/go.mod h1:46gg/UBmTU1kUaTcwQXpUxtRwG2PvIZYeA8oL6vF3Fs= github.com/frankban/quicktest v1.14.4 h1:g2rn0vABPOOXmZUj+vbmUp0lPoXEMuhTpIluN0XL9UY= github.com/frankban/quicktest v1.14.4/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= github.com/fsnotify/fsnotify v1.7.0 h1:8JEhPFa5W2WU7YfeZzPNqzMP6Lwt7L2715Ggo0nosvA= @@ -70,6 +72,8 @@ github.com/go-ole/go-ole v1.3.0 h1:Dt6ye7+vXGIKZ7Xtk4s6/xVdGDQynvom7xCFEdWr6uE= github.com/go-ole/go-ole v1.3.0/go.mod h1:5LS6F96DhAwUc7C+1HLexzMXY1xGRSryjyPPKW6zv78= github.com/gofiber/fiber/v2 v2.50.0 h1:ia0JaB+uw3GpNSCR5nvC5dsaxXjRU5OEu36aytx+zGw= github.com/gofiber/fiber/v2 v2.50.0/go.mod h1:21eytvay9Is7S6z+OgPi7c7n4++tnClWmhpimVHMimw= +github.com/gofiber/websocket/v2 v2.2.1 h1:C9cjxvloojayOp9AovmpQrk8VqvVnT8Oao3+IUygH7w= +github.com/gofiber/websocket/v2 v2.2.1/go.mod h1:Ao/+nyNnX5u/hIFPuHl28a+NIkrqK7PRimyKaj4JxVU= github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q= github.com/golang/groupcache v0.0.0-20190702054246-869f871628b6/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc= github.com/golang/groupcache v0.0.0-20191227052852-215e87163ea7/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc= @@ -196,6 +200,8 @@ github.com/sagikazarmark/locafero v0.3.0 h1:zT7VEGWC2DTflmccN/5T1etyKvxSxpHsjb9c github.com/sagikazarmark/locafero v0.3.0/go.mod h1:w+v7UsPNFwzF1cHuOajOOzoq4U7v/ig1mpRjqV+Bu1U= github.com/sagikazarmark/slog-shim v0.1.0 h1:diDBnUNK9N/354PgrxMywXnAwEr1QZcOr6gto+ugjYE= github.com/sagikazarmark/slog-shim v0.1.0/go.mod h1:SrcSrq8aKtyuqEI1uvTDTK1arOWRIczQRv+GVI1AkeQ= +github.com/savsgio/gotils v0.0.0-20230208104028-c358bd845dee h1:8Iv5m6xEo1NR1AvpV+7XmhI4r39LGNzwUL4YpMuL5vk= +github.com/savsgio/gotils v0.0.0-20230208104028-c358bd845dee/go.mod h1:qwtSXrKuJh/zsFQ12yEE89xfCrGKK63Rr7ctU/uCo4g= github.com/shirou/gopsutil/v3 v3.23.10 h1:/N42opWlYzegYaVkWejXWJpbzKv2JDy3mrgGzKsh9hM= github.com/shirou/gopsutil/v3 v3.23.10/go.mod h1:JIE26kpucQi+innVlAUnIEOSBhBUkirr5b44yr55+WE= github.com/shoenig/go-m1cpu v0.1.6 h1:nxdKQNcEB6vzgA2E2bvzKIYRuNj7XNJ4S/aRSwKzFtM= diff --git a/routes/routes.go b/routes/routes.go index 693dd0d..a6de13f 100644 --- a/routes/routes.go +++ b/routes/routes.go @@ -4,6 +4,7 @@ import ( "api/controllers" "github.com/gofiber/fiber/v2" + "github.com/gofiber/websocket/v2" ) // Setup sets up the routes @@ -28,4 +29,15 @@ func Setup(app *fiber.App) { app.Post("/on_pub_start", controllers.OnPubStart) app.Post("/on_pub_stop", controllers.OnPubStop) // app.Post("/on_sub_start", controllers.OnSubStart) + + // WebSocket route + app.Use("/ws", func(c *fiber.Ctx) error { + if websocket.IsWebSocketUpgrade(c) { + c.Locals("allowed", true) + return c.Next() + } + return fiber.ErrUpgradeRequired + }) + + app.Get("/ws", websocket.New(controllers.WebsocketHandler)) }