From 4df3638527fed976f385f9a8a923cc44e9c0d60b Mon Sep 17 00:00:00 2001 From: Aldino Kemal Date: Wed, 23 Nov 2022 06:36:51 +0700 Subject: [PATCH] feat: move WebSocket to independent package --- src/cmd/root.go | 5 ++- .../rest/{send_controller.go => send_rest.go} | 0 .../rest/{user_controller.go => user_rest.go} | 0 .../{rest/helpers => websocket}/websocket.go | 39 +++++++++---------- src/utils/whatsapp.go | 4 +- 5 files changed, 24 insertions(+), 24 deletions(-) rename src/internal/rest/{send_controller.go => send_rest.go} (100%) rename src/internal/rest/{user_controller.go => user_rest.go} (100%) rename src/internal/{rest/helpers => websocket}/websocket.go (71%) diff --git a/src/cmd/root.go b/src/cmd/root.go index 9239cd0..7341138 100644 --- a/src/cmd/root.go +++ b/src/cmd/root.go @@ -6,6 +6,7 @@ import ( "github.com/aldinokemal/go-whatsapp-web-multidevice/config" "github.com/aldinokemal/go-whatsapp-web-multidevice/internal/rest" "github.com/aldinokemal/go-whatsapp-web-multidevice/internal/rest/helpers" + "github.com/aldinokemal/go-whatsapp-web-multidevice/internal/websocket" "github.com/aldinokemal/go-whatsapp-web-multidevice/middleware" "github.com/aldinokemal/go-whatsapp-web-multidevice/services" "github.com/aldinokemal/go-whatsapp-web-multidevice/utils" @@ -108,8 +109,8 @@ func runRest(_ *cobra.Command, _ []string) { }) }) - helpers.WsRegisterRoutes(app, cli) - go helpers.WsRunHub() + websocket.RegisterRoutes(app) + go websocket.RunHub() // Set auto reconnect to whatsapp server after booting go helpers.SetAutoConnectAfterBooting() diff --git a/src/internal/rest/send_controller.go b/src/internal/rest/send_rest.go similarity index 100% rename from src/internal/rest/send_controller.go rename to src/internal/rest/send_rest.go diff --git a/src/internal/rest/user_controller.go b/src/internal/rest/user_rest.go similarity index 100% rename from src/internal/rest/user_controller.go rename to src/internal/rest/user_rest.go diff --git a/src/internal/rest/helpers/websocket.go b/src/internal/websocket/websocket.go similarity index 71% rename from src/internal/rest/helpers/websocket.go rename to src/internal/websocket/websocket.go index 59ddf5c..bf5098f 100644 --- a/src/internal/rest/helpers/websocket.go +++ b/src/internal/websocket/websocket.go @@ -1,32 +1,31 @@ -package helpers +package websocket import ( "encoding/json" "github.com/gofiber/fiber/v2" "github.com/gofiber/websocket/v2" - "go.mau.fi/whatsmeow" "log" ) type client struct{} // Add more data to this type if needed -type WsBroadcastMessage struct { +type BroadcastMessage struct { Code string `json:"code"` Message string `json:"message"` } -var WsClients = make(map[*websocket.Conn]client) // Note: although large maps with pointer-like types (e.g. strings) as keys are slow, using pointers themselves as keys is acceptable and fast -var WsRegister = make(chan *websocket.Conn) -var WsBroadcast = make(chan WsBroadcastMessage) -var WsUnregister = make(chan *websocket.Conn) +var Clients = make(map[*websocket.Conn]client) // Note: although large maps with pointer-like types (e.g. strings) as keys are slow, using pointers themselves as keys is acceptable and fast +var Register = make(chan *websocket.Conn) +var Broadcast = make(chan BroadcastMessage) +var Unregister = make(chan *websocket.Conn) -func WsRunHub() { +func RunHub() { for { select { - case connection := <-WsRegister: - WsClients[connection] = client{} + case connection := <-Register: + Clients[connection] = client{} log.Println("connection registered") - case message := <-WsBroadcast: + case message := <-Broadcast: log.Println("message received:", message) marshalMessage, err := json.Marshal(message) if err != nil { @@ -35,7 +34,7 @@ func WsRunHub() { } // Send the message to all clients - for connection := range WsClients { + for connection := range Clients { if err := connection.WriteMessage(websocket.TextMessage, marshalMessage); err != nil { log.Println("write error:", err) @@ -49,20 +48,20 @@ func WsRunHub() { log.Println("close error:", err) return } - delete(WsClients, connection) + delete(Clients, connection) } } - case connection := <-WsUnregister: + case connection := <-Unregister: // Remove the client from the hub - delete(WsClients, connection) + delete(Clients, connection) log.Println("connection unregistered") } } } -func WsRegisterRoutes(app *fiber.App, cli *whatsmeow.Client) { +func RegisterRoutes(app *fiber.App) { app.Use("/ws", func(c *fiber.Ctx) error { if websocket.IsWebSocketUpgrade(c) { // Returns true if the client requested upgrade to the WebSocket protocol return c.Next() @@ -73,12 +72,12 @@ func WsRegisterRoutes(app *fiber.App, cli *whatsmeow.Client) { app.Get("/ws", websocket.New(func(c *websocket.Conn) { // When the function returns, unregister the client and close the connection defer func() { - WsUnregister <- c + Unregister <- c _ = c.Close() }() // Register the client - WsRegister <- c + Register <- c for { messageType, message, err := c.ReadMessage() @@ -91,13 +90,13 @@ func WsRegisterRoutes(app *fiber.App, cli *whatsmeow.Client) { if messageType == websocket.TextMessage { // Broadcast the received message - var messageData WsBroadcastMessage + var messageData BroadcastMessage err := json.Unmarshal(message, &messageData) if err != nil { log.Println("error unmarshal message:", err) return } - WsBroadcast <- messageData + Broadcast <- messageData } else { log.Println("websocket message received of type", messageType) } diff --git a/src/utils/whatsapp.go b/src/utils/whatsapp.go index dec7696..4bba2ac 100644 --- a/src/utils/whatsapp.go +++ b/src/utils/whatsapp.go @@ -6,7 +6,7 @@ import ( "encoding/json" "fmt" "github.com/aldinokemal/go-whatsapp-web-multidevice/config" - "github.com/aldinokemal/go-whatsapp-web-multidevice/internal/rest/helpers" + "github.com/aldinokemal/go-whatsapp-web-multidevice/internal/websocket" "go.mau.fi/whatsmeow" "go.mau.fi/whatsmeow/appstate" waProto "go.mau.fi/whatsmeow/binary/proto" @@ -136,7 +136,7 @@ func handler(rawEvt interface{}) { } } case *events.PairSuccess: - helpers.WsBroadcast <- helpers.WsBroadcastMessage{ + websocket.Broadcast <- websocket.BroadcastMessage{ Code: "LOGIN_SUCCESS", Message: fmt.Sprintf("Successfully pair with %s", evt.ID.String()), }