Merge pull request #7 from LubosD/feature/avoid-panic-on-disconnect

Avoid panic on disconnect, provide connection errors

Author: saml-dev

README.md

@@ -40,7 +40,7 @@ The general flow is
 import ga "saml.dev/gome-assistant"
 
 // replace with IP and port of your Home Assistant installation
-app := ga.NewApp("0.0.0.0:8123")
+app, err := ga.NewApp("0.0.0.0:8123")
 
 // create automations here (see next sections)
 

app.go

@@ -15,6 +15,9 @@ import (
 	ws "saml.dev/gome-assistant/internal/websocket"
 )
 
+// Returned by NewApp() if authentication fails
+var ErrInvalidToken = ws.ErrInvalidToken
+
 type App struct {
 	ctx       context.Context
 	ctxCancel context.CancelFunc
@@ -77,15 +80,19 @@ type NewAppRequest struct {
 NewApp establishes the websocket connection and returns an object
 you can use to register schedules and listeners.
 */
-func NewApp(request NewAppRequest) *App {
+func NewApp(request NewAppRequest) (*App, error) {
 	if request.IpAddress == "" || request.HAAuthToken == "" || request.HomeZoneEntityId == "" {
 		log.Fatalln("IpAddress, HAAuthToken, and HomeZoneEntityId are all required arguments in NewAppRequest.")
 	}
 	port := request.Port
 	if port == "" {
 		port = "8123"
 	}
-	conn, ctx, ctxCancel := ws.SetupConnection(request.IpAddress, port, request.HAAuthToken)
+	conn, ctx, ctxCancel, err := ws.SetupConnection(request.IpAddress, port, request.HAAuthToken)
+
+	if conn == nil {
+		return nil, err
+	}
 
 	httpClient := http.NewHttpClient(request.IpAddress, port, request.HAAuthToken)
 
@@ -105,7 +112,7 @@ func NewApp(request NewAppRequest) *App {
 		intervals:       pq.New(),
 		entityListeners: map[string][]*EntityListener{},
 		eventListeners:  map[string][]*EventListener{},
-	}
+	}, nil
 }
 
 func (a *App) Cleanup() {
@@ -263,9 +270,12 @@ func (a *App) Start() {
 	// entity listeners and event listeners
 	elChan := make(chan ws.ChanMsg)
 	go ws.ListenWebsocket(a.conn, a.ctx, elChan)
-	var msg ws.ChanMsg
+
 	for {
-		msg = <-elChan
+		msg, ok := <-elChan
+		if !ok {
+			break
+		}
 		if a.entityListenersId == msg.Id {
 			go callEntityListeners(a, msg.Raw)
 		} else {

example/example.go

@@ -10,11 +10,16 @@ import (
 )
 
 func main() {
-	app := ga.NewApp(ga.NewAppRequest{
+	app, err := ga.NewApp(ga.NewAppRequest{
 		IpAddress:        "192.168.86.67", // Replace with your Home Assistant IP Address
 		HAAuthToken:      os.Getenv("HA_AUTH_TOKEN"),
 		HomeZoneEntityId: "zone.home",
 	})
+
+	if err != nil {
+		log.Fatalln("Error connecting to HASS:", err)
+	}
+
 	defer app.Cleanup()
 
 	pantryDoor := ga.

internal/websocket/reader.go

@@ -3,6 +3,7 @@ package websocket
 import (
 	"context"
 	"encoding/json"
+	"log"
 
 	"github.com/gorilla/websocket"
 )
@@ -20,7 +21,14 @@ type ChanMsg struct {
 
 func ListenWebsocket(conn *websocket.Conn, ctx context.Context, c chan ChanMsg) {
 	for {
-		bytes, _ := ReadMessage(conn, ctx)
+		bytes, err := ReadMessage(conn, ctx)
+
+		if err != nil {
+			log.Default().Println("Error reading from websocket:", err)
+			close(c)
+			break
+		}
+
 		base := BaseMessage{}
 		json.Unmarshal(bytes, &base)
 		chanMsg := ChanMsg{

internal/websocket/websocket.go

@@ -17,6 +17,10 @@ import (
 	i "saml.dev/gome-assistant/internal"
 )
 
+var (
+	ErrInvalidToken = errors.New("invalid authentication token")
+)
+
 type AuthMessage struct {
 	MsgType     string `json:"type"`
 	AccessToken string `json:"access_token"`
@@ -47,39 +51,43 @@ func ReadMessage(conn *websocket.Conn, ctx context.Context) ([]byte, error) {
 	return msg, nil
 }
 
-func SetupConnection(ip, port, authToken string) (*websocket.Conn, context.Context, context.CancelFunc) {
+func SetupConnection(ip, port, authToken string) (*websocket.Conn, context.Context, context.CancelFunc, error) {
 	ctx, ctxCancel := context.WithTimeout(context.Background(), time.Second*3)
 
 	// Init websocket connection
 	dialer := websocket.DefaultDialer
 	conn, _, err := dialer.DialContext(ctx, fmt.Sprintf("ws://%s:%s/api/websocket", ip, port), nil)
 	if err != nil {
 		ctxCancel()
-		log.Fatalf("ERROR: Failed to connect to websocket at ws://%s:%s/api/websocket. Check IP address and port\n", ip, port)
+		log.Printf("ERROR: Failed to connect to websocket at ws://%s:%s/api/websocket. Check IP address and port\n", ip, port)
+		return nil, nil, nil, err
 	}
 
 	// Read auth_required message
 	_, err = ReadMessage(conn, ctx)
 	if err != nil {
 		ctxCancel()
-		log.Fatalf("Unknown error creating websocket client\n")
+		log.Printf("Unknown error creating websocket client\n")
+		return nil, nil, nil, err
 	}
 
 	// Send auth message
 	err = SendAuthMessage(conn, ctx, authToken)
 	if err != nil {
 		ctxCancel()
-		log.Fatalf("Unknown error creating websocket client\n")
+		log.Printf("Unknown error creating websocket client\n")
+		return nil, nil, nil, err
 	}
 
 	// Verify auth message was successful
 	err = VerifyAuthResponse(conn, ctx)
 	if err != nil {
 		ctxCancel()
-		log.Fatalf("ERROR: Auth token is invalid. Please double check it or create a new token in your Home Assistant profile\n")
+		log.Printf("ERROR: Auth token is invalid. Please double check it or create a new token in your Home Assistant profile\n")
+		return nil, nil, nil, err
 	}
 
-	return conn, ctx, ctxCancel
+	return conn, ctx, ctxCancel, nil
 }
 
 func SendAuthMessage(conn *websocket.Conn, ctx context.Context, token string) error {
@@ -105,7 +113,7 @@ func VerifyAuthResponse(conn *websocket.Conn, ctx context.Context) error {
 	json.Unmarshal(msg, &authResp)
 	// log.Println(authResp.MsgType)
 	if authResp.MsgType != "auth_ok" {
-		return errors.New("invalid auth token")
+		return ErrInvalidToken
 	}
 
 	return nil