Fix panics due to concurrent websocket writes

app.go

@@ -16,9 +16,13 @@ import (
 )
 
 type App struct {
-	ctx        context.Context
-	ctxCancel  context.CancelFunc
-	conn       *websocket.Conn
+	ctx       context.Context
+	ctxCancel context.CancelFunc
+	conn      *websocket.Conn
+
+	// Wraps the ws connection with added mutex locking
+	wsWriter *ws.WebsocketWriter
+
 	httpClient *http.HttpClient
 
 	service *Service
@@ -85,11 +89,13 @@ func NewApp(request NewAppRequest) *App {
 
 	httpClient := http.NewHttpClient(request.IpAddress, port, request.HAAuthToken)
 
-	service := newService(conn, ctx, httpClient)
+	wsWriter := &ws.WebsocketWriter{Conn: conn}
+	service := newService(wsWriter, ctx, httpClient)
 	state := newState(httpClient, request.HomeZoneEntityId)
 
 	return &App{
 		conn:            conn,
+		wsWriter:        wsWriter,
 		ctx:             ctx,
 		ctxCancel:       ctxCancel,
 		httpClient:      httpClient,
@@ -169,7 +175,7 @@ func (a *App) RegisterEventListeners(evls ...EventListener) {
 			if elList, ok := a.eventListeners[eventType]; ok {
 				a.eventListeners[eventType] = append(elList, &evl)
 			} else {
-				ws.SubscribeToEventType(eventType, a.conn, a.ctx)
+				ws.SubscribeToEventType(eventType, a.wsWriter, a.ctx)
 				a.eventListeners[eventType] = []*EventListener{&evl}
 			}
 		}
@@ -227,7 +233,7 @@ func (a *App) Start() {
 
 	// subscribe to state_changed events
 	id := internal.GetId()
-	ws.SubscribeToStateChangedEvents(id, a.conn, a.ctx)
+	ws.SubscribeToStateChangedEvents(id, a.wsWriter, a.ctx)
 	a.entityListenersId = id
 
 	// entity listeners runOnStartup

internal/services/alarm_control_panel.go

@@ -3,14 +3,13 @@ package services
 import (
 	"context"
 
-	"github.com/gorilla/websocket"
 	ws "saml.dev/gome-assistant/internal/websocket"
 )
 
 /* Structs */
 
 type AlarmControlPanel struct {
-	conn *websocket.Conn
+	conn *ws.WebsocketWriter
 	ctx  context.Context
 }
 
@@ -27,7 +26,7 @@ func (acp AlarmControlPanel) ArmAway(entityId string, serviceData ...map[string]
 		req.ServiceData = serviceData[0]
 	}
 
-	ws.WriteMessage(req, acp.conn, acp.ctx)
+	acp.conn.WriteMessage(req, acp.ctx)
 }
 
 // Send the alarm the command for arm away.
@@ -41,7 +40,7 @@ func (acp AlarmControlPanel) ArmWithCustomBypass(entityId string, serviceData ..
 		req.ServiceData = serviceData[0]
 	}
 
-	ws.WriteMessage(req, acp.conn, acp.ctx)
+	acp.conn.WriteMessage(req, acp.ctx)
 }
 
 // Send the alarm the command for arm home.
@@ -55,7 +54,7 @@ func (acp AlarmControlPanel) ArmHome(entityId string, serviceData ...map[string]
 		req.ServiceData = serviceData[0]
 	}
 
-	ws.WriteMessage(req, acp.conn, acp.ctx)
+	acp.conn.WriteMessage(req, acp.ctx)
 }
 
 // Send the alarm the command for arm night.
@@ -69,7 +68,7 @@ func (acp AlarmControlPanel) ArmNight(entityId string, serviceData ...map[string
 		req.ServiceData = serviceData[0]
 	}
 
-	ws.WriteMessage(req, acp.conn, acp.ctx)
+	acp.conn.WriteMessage(req, acp.ctx)
 }
 
 // Send the alarm the command for arm vacation.
@@ -83,7 +82,7 @@ func (acp AlarmControlPanel) ArmVacation(entityId string, serviceData ...map[str
 		req.ServiceData = serviceData[0]
 	}
 
-	ws.WriteMessage(req, acp.conn, acp.ctx)
+	acp.conn.WriteMessage(req, acp.ctx)
 }
 
 // Send the alarm the command for disarm.
@@ -97,7 +96,7 @@ func (acp AlarmControlPanel) Disarm(entityId string, serviceData ...map[string]a
 		req.ServiceData = serviceData[0]
 	}
 
-	ws.WriteMessage(req, acp.conn, acp.ctx)
+	acp.conn.WriteMessage(req, acp.ctx)
 }
 
 // Send the alarm the command for trigger.
@@ -111,5 +110,5 @@ func (acp AlarmControlPanel) Trigger(entityId string, serviceData ...map[string]
 		req.ServiceData = serviceData[0]
 	}
 
-	ws.WriteMessage(req, acp.conn, acp.ctx)
+	acp.conn.WriteMessage(req, acp.ctx)
 }

internal/services/cover.go

@@ -3,14 +3,13 @@ package services
 import (
 	"context"
 
-	"github.com/gorilla/websocket"
 	ws "saml.dev/gome-assistant/internal/websocket"
 )
 
 /* Structs */
 
 type Cover struct {
-	conn *websocket.Conn
+	conn *ws.WebsocketWriter
 	ctx  context.Context
 }
 
@@ -22,7 +21,7 @@ func (c Cover) Close(entityId string) {
 	req.Domain = "cover"
 	req.Service = "close_cover"
 
-	ws.WriteMessage(req, c.conn, c.ctx)
+	c.conn.WriteMessage(req, c.ctx)
 }
 
 // Close all or specified cover tilt. Takes an entityId.
@@ -31,7 +30,7 @@ func (c Cover) CloseTilt(entityId string) {
 	req.Domain = "cover"
 	req.Service = "close_cover_tilt"
 
-	ws.WriteMessage(req, c.conn, c.ctx)
+	c.conn.WriteMessage(req, c.ctx)
 }
 
 // Open all or specified cover. Takes an entityId.
@@ -40,7 +39,7 @@ func (c Cover) Open(entityId string) {
 	req.Domain = "cover"
 	req.Service = "open_cover"
 
-	ws.WriteMessage(req, c.conn, c.ctx)
+	c.conn.WriteMessage(req, c.ctx)
 }
 
 // Open all or specified cover tilt. Takes an entityId.
@@ -49,7 +48,7 @@ func (c Cover) OpenTilt(entityId string) {
 	req.Domain = "cover"
 	req.Service = "open_cover_tilt"
 
-	ws.WriteMessage(req, c.conn, c.ctx)
+	c.conn.WriteMessage(req, c.ctx)
 }
 
 // Move to specific position all or specified cover. Takes an entityId and an optional
@@ -62,7 +61,7 @@ func (c Cover) SetPosition(entityId string, serviceData ...map[string]any) {
 		req.ServiceData = serviceData[0]
 	}
 
-	ws.WriteMessage(req, c.conn, c.ctx)
+	c.conn.WriteMessage(req, c.ctx)
 }
 
 // Move to specific position all or specified cover tilt. Takes an entityId and an optional
@@ -75,7 +74,7 @@ func (c Cover) SetTiltPosition(entityId string, serviceData ...map[string]any) {
 		req.ServiceData = serviceData[0]
 	}
 
-	ws.WriteMessage(req, c.conn, c.ctx)
+	c.conn.WriteMessage(req, c.ctx)
 }
 
 // Stop a cover entity. Takes an entityId.
@@ -84,7 +83,7 @@ func (c Cover) Stop(entityId string) {
 	req.Domain = "cover"
 	req.Service = "stop_cover"
 
-	ws.WriteMessage(req, c.conn, c.ctx)
+	c.conn.WriteMessage(req, c.ctx)
 }
 
 // Stop a cover entity tilt. Takes an entityId.
@@ -93,7 +92,7 @@ func (c Cover) StopTilt(entityId string) {
 	req.Domain = "cover"
 	req.Service = "stop_cover_tilt"
 
-	ws.WriteMessage(req, c.conn, c.ctx)
+	c.conn.WriteMessage(req, c.ctx)
 }
 
 // Toggle a cover open/closed. Takes an entityId.
@@ -102,7 +101,7 @@ func (c Cover) Toggle(entityId string) {
 	req.Domain = "cover"
 	req.Service = "toggle"
 
-	ws.WriteMessage(req, c.conn, c.ctx)
+	c.conn.WriteMessage(req, c.ctx)
 }
 
 // Toggle a cover tilt open/closed. Takes an entityId.
@@ -111,5 +110,5 @@ func (c Cover) ToggleTilt(entityId string) {
 	req.Domain = "cover"
 	req.Service = "toggle_cover_tilt"
 
-	ws.WriteMessage(req, c.conn, c.ctx)
+	c.conn.WriteMessage(req, c.ctx)
 }

internal/services/homeassistant.go

@@ -3,12 +3,11 @@ package services
 import (
 	"context"
 
-	"github.com/gorilla/websocket"
 	ws "saml.dev/gome-assistant/internal/websocket"
 )
 
 type HomeAssistant struct {
-	conn *websocket.Conn
+	conn *ws.WebsocketWriter
 	ctx  context.Context
 }
 
@@ -22,7 +21,7 @@ func (ha *HomeAssistant) TurnOn(entityId string, serviceData ...map[string]any)
 		req.ServiceData = serviceData[0]
 	}
 
-	ws.WriteMessage(req, ha.conn, ha.ctx)
+	ha.conn.WriteMessage(req, ha.ctx)
 }
 
 // Toggle a Home Assistant entity. Takes an entityId and an optional
@@ -35,13 +34,13 @@ func (ha *HomeAssistant) Toggle(entityId string, serviceData ...map[string]any)
 		req.ServiceData = serviceData[0]
 	}
 
-	ws.WriteMessage(req, ha.conn, ha.ctx)
+	ha.conn.WriteMessage(req, ha.ctx)
 }
 
 func (ha *HomeAssistant) TurnOff(entityId string) {
 	req := NewBaseServiceRequest(entityId)
 	req.Domain = "homeassistant"
 	req.Service = "turn_off"
 
-	ws.WriteMessage(req, ha.conn, ha.ctx)
+	ha.conn.WriteMessage(req, ha.ctx)
 }

internal/services/input_boolean.go

@@ -3,14 +3,13 @@ package services
 import (
 	"context"
 
-	"github.com/gorilla/websocket"
 	ws "saml.dev/gome-assistant/internal/websocket"
 )
 
 /* Structs */
 
 type InputBoolean struct {
-	conn *websocket.Conn
+	conn *ws.WebsocketWriter
 	ctx  context.Context
 }
 
@@ -21,27 +20,27 @@ func (ib InputBoolean) TurnOn(entityId string) {
 	req.Domain = "input_boolean"
 	req.Service = "turn_on"
 
-	ws.WriteMessage(req, ib.conn, ib.ctx)
+	ib.conn.WriteMessage(req, ib.ctx)
 }
 
 func (ib InputBoolean) Toggle(entityId string) {
 	req := NewBaseServiceRequest(entityId)
 	req.Domain = "input_boolean"
 	req.Service = "toggle"
 
-	ws.WriteMessage(req, ib.conn, ib.ctx)
+	ib.conn.WriteMessage(req, ib.ctx)
 }
 
 func (ib InputBoolean) TurnOff(entityId string) {
 	req := NewBaseServiceRequest(entityId)
 	req.Domain = "input_boolean"
 	req.Service = "turn_off"
-	ws.WriteMessage(req, ib.conn, ib.ctx)
+	ib.conn.WriteMessage(req, ib.ctx)
 }
 
 func (ib InputBoolean) Reload() {
 	req := NewBaseServiceRequest("")
 	req.Domain = "input_boolean"
 	req.Service = "reload"
-	ws.WriteMessage(req, ib.conn, ib.ctx)
+	ib.conn.WriteMessage(req, ib.ctx)
 }

internal/services/input_button.go

@@ -3,14 +3,13 @@ package services
 import (
 	"context"
 
-	"github.com/gorilla/websocket"
 	ws "saml.dev/gome-assistant/internal/websocket"
 )
 
 /* Structs */
 
 type InputButton struct {
-	conn *websocket.Conn
+	conn *ws.WebsocketWriter
 	ctx  context.Context
 }
 
@@ -21,12 +20,12 @@ func (ib InputButton) Press(entityId string) {
 	req.Domain = "input_button"
 	req.Service = "press"
 
-	ws.WriteMessage(req, ib.conn, ib.ctx)
+	ib.conn.WriteMessage(req, ib.ctx)
 }
 
 func (ib InputButton) Reload() {
 	req := NewBaseServiceRequest("")
 	req.Domain = "input_button"
 	req.Service = "reload"
-	ws.WriteMessage(req, ib.conn, ib.ctx)
+	ib.conn.WriteMessage(req, ib.ctx)
 }

internal/services/input_datetime.go

@@ -5,14 +5,13 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/gorilla/websocket"
 	ws "saml.dev/gome-assistant/internal/websocket"
 )
 
 /* Structs */
 
 type InputDatetime struct {
-	conn *websocket.Conn
+	conn *ws.WebsocketWriter
 	ctx  context.Context
 }
 
@@ -26,12 +25,12 @@ func (ib InputDatetime) Set(entityId string, value time.Time) {
 		"timestamp": fmt.Sprint(value.Unix()),
 	}
 
-	ws.WriteMessage(req, ib.conn, ib.ctx)
+	ib.conn.WriteMessage(req, ib.ctx)
 }
 
 func (ib InputDatetime) Reload() {
 	req := NewBaseServiceRequest("")
 	req.Domain = "input_datetime"
 	req.Service = "reload"
-	ws.WriteMessage(req, ib.conn, ib.ctx)
+	ib.conn.WriteMessage(req, ib.ctx)
 }