Robust websockets

src/main.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"net/http"
 	"strconv"
+	"strings"
 	"sync"
 
 	"github.com/labstack/echo/v4"
@@ -98,33 +99,44 @@ func main() {
 		return c.Render(http.StatusOK, "article", article)
 	})
 
-	e.POST("/box/:id", updateBoxes)
 	e.GET("/ws", initWs)
 
 	e.Logger.Fatal(e.Start(":54321"))
 }
 
-func updateBoxes(c echo.Context) error {
-	id := c.Param("id")
-	checked := c.FormValue("checked") == "on"
+func sendUpdate(index int, checked bool) {
+	message := formatMessage(index, checked)
+	broadcastMessage(message)
+}
 
-	index, err := strconv.Atoi(id)
-	if err != nil {
-		return c.String(http.StatusBadRequest, "Invalid ID")
-	}
-
-	boxes[index] = checked
-
-	message := strconv.Itoa(index)
+func formatMessage(index int, checked bool) string {
+	message := "u:" + strconv.Itoa(index)
 	if checked {
 		message += ":+"
 	} else {
 		message += ":-"
 	}
+	return message
+}
 
-	broadcastMessage(message)
+func handleMessage(msg string, ws *websocket.Conn) error {
+	if strings.Contains(msg, "u:") {
+		parts := strings.Split(msg, ":")
+		index, err := strconv.Atoi(parts[1])
+		if err != nil {
+			return err
+		}
+		checked := parts[2] == "+"
 
-	return c.JSON(200, "[]")
+		boxes[index] = checked
+
+		sendUpdate(index, checked)
+	}
+	if strings.Contains(msg, "ping") {
+		len := len(wsConnections)
+		websocket.Message.Send(ws, "pong-"+strconv.Itoa(len))
+	}
+	return nil
 }
 
 func initWs(c echo.Context) error {
@@ -133,6 +145,18 @@ func initWs(c echo.Context) error {
 		wsConnections[ws] = true
 		wsMutex.Unlock()
 
+		initialState := "i:"
+		for index, checked := range boxes {
+			if checked {
+				initialState += strconv.Itoa(index) + ";"
+			}
+		}
+
+		if err := websocket.Message.Send(ws, initialState); err != nil {
+			log.Errorf("Failed to send initial data: %v", err)
+			return
+		}
+
 		for {
 			var msg string
 			if err := websocket.Message.Receive(ws, &msg); err != nil {
@@ -142,7 +166,10 @@ func initWs(c echo.Context) error {
 				ws.Close()
 				break
 			}
+			handleMessage(msg, ws)
+			log.Infof("Received message: %s", msg)
 		}
 	}).ServeHTTP(c.Response(), c.Request())
+
 	return nil
 }