refactor: websockets into 'connect' module, rename & adjust generally

internal/connect/connection.go

@@ -2,7 +2,7 @@
 // websocket API. All HA interaction is done via websocket
 // except for cases explicitly called out in http package
 // documentation.
-package websocket
+package connect
 
 import (
 	"context"
@@ -20,46 +20,60 @@ import (
 
 var ErrInvalidToken = errors.New("invalid authentication token")
 
-type AuthMessage struct {
-	MsgType     string `json:"type"`
-	AccessToken string `json:"access_token"`
-}
-
-type WebsocketWriter struct {
-	Conn  *websocket.Conn
+// HAConnection is a wrapper around a websocket connection that provides a mutex for thread safety.
+type HAConnection struct {
+	Conn  *websocket.Conn // Note: this is not thread safe except for Close() and WriteControl()
 	mutex sync.Mutex
 }
 
-func (w *WebsocketWriter) WriteMessage(msg any) error {
+// WriteMessage writes a message to the websocket connection.
+func (w *HAConnection) WriteMessage(msg any) error {
 	w.mutex.Lock()
 	defer w.mutex.Unlock()
 
 	return w.Conn.WriteJSON(msg)
 }
 
-func ReadMessage(conn *websocket.Conn) ([]byte, error) {
+// ReadMessageRaw reads a raw message from the websocket connection.
+func ReadMessageRaw(conn *websocket.Conn) ([]byte, error) {
 	_, msg, err := conn.ReadMessage()
 	if err != nil {
-		return []byte{}, err
+		return nil, err
 	}
 	return msg, nil
 }
 
-func ConnectionFromUri(baseURL *url.URL, authToken string) (*websocket.Conn, context.Context, context.CancelFunc, error) {
+// ReadMessage reads a message from the websocket connection and unmarshals it into the given type.
+func ReadMessage[T any](conn *websocket.Conn) (T, error) {
+	var result T
+	_, msg, err := conn.ReadMessage()
+	if err != nil {
+		return result, err
+	}
+
+	err = json.Unmarshal(msg, &result)
+	if err != nil {
+		return result, err
+	}
+
+	return result, nil
+}
+
+// ConnectionFromUri creates a new websocket connection from the given base URL and authentication token.
+func ConnectionFromUri(baseUrl *url.URL, token string) (*HAConnection, context.Context, context.CancelFunc, error) {
+	// Build the websocket URL
+	urlWebsockets := *baseUrl
+	urlWebsockets.Path = "/api/websocket"
+	scheme, err := internal.GetEquivalentWebsocketScheme(baseUrl.Scheme)
+	if err != nil {
+		return nil, nil, nil, fmt.Errorf("failed to build websocket URL: %w", err)
+	}
+	urlWebsockets.Scheme = scheme
+
 	// Create a short timeout context for the connection only
 	connCtx, connCtxCancel := context.WithTimeout(context.Background(), time.Second*3)
 	defer connCtxCancel() // Always cancel the connection context when we're done
 
-	// Shallow copy the URL to avoid modifying the original
-	urlWebsockets := *baseURL
-	urlWebsockets.Path = "/api/websocket"
-	if baseURL.Scheme == "http" {
-		urlWebsockets.Scheme = "ws"
-	}
-	if baseURL.Scheme == "https" {
-		urlWebsockets.Scheme = "wss"
-	}
-
 	// Init websocket connection
 	dialer := websocket.DefaultDialer
 	conn, _, err := dialer.DialContext(connCtx, urlWebsockets.String(), nil)
@@ -69,14 +83,19 @@ func ConnectionFromUri(baseURL *url.URL, authToken string) (*websocket.Conn, con
 	}
 
 	// Read auth_required message
-	_, err = ReadMessage(conn)
+	msg, err := ReadMessage[struct {
+		MsgType string `json:"type"`
+	}](conn)
 	if err != nil {
 		slog.Error("Unknown error creating websocket client\n")
 		return nil, nil, nil, err
+	} else if msg.MsgType != "auth_required" {
+		slog.Error("Expected auth_required message, got", "msgType", msg.MsgType)
+		return nil, nil, nil, fmt.Errorf("expected auth_required message, got %s", msg.MsgType)
 	}
 
 	// Send auth message
-	err = SendAuthMessage(conn, connCtx, authToken)
+	err = SendAuthMessage(conn, connCtx, token)
 	if err != nil {
 		slog.Error("Unknown error creating websocket client\n")
 		return nil, nil, nil, err
@@ -92,51 +111,54 @@ func ConnectionFromUri(baseURL *url.URL, authToken string) (*websocket.Conn, con
 	// Create a new background context for the application lifecycle (no timeout)
 	appCtx, appCtxCancel := context.WithCancel(context.Background())
 
-	return conn, appCtx, appCtxCancel, nil
+	return &HAConnection{Conn: conn}, appCtx, appCtxCancel, nil
 }
 
+// SendAuthMessage sends an auth message to the websocket connection.
 func SendAuthMessage(conn *websocket.Conn, ctx context.Context, token string) error {
+	type AuthMessage struct {
+		MsgType     string `json:"type"`
+		AccessToken string `json:"access_token"`
+	}
+
 	err := conn.WriteJSON(AuthMessage{MsgType: "auth", AccessToken: token})
 	if err != nil {
 		return err
 	}
+
 	return nil
 }
 
-type authResponse struct {
-	MsgType string `json:"type"`
-	Message string `json:"message"`
-}
-
+// VerifyAuthResponse verifies that the auth response is valid.
 func VerifyAuthResponse(conn *websocket.Conn, ctx context.Context) error {
-	msg, err := ReadMessage(conn)
+	msg, err := ReadMessage[struct {
+		MsgType string `json:"type"`
+		Message string `json:"message"`
+	}](conn)
 	if err != nil {
 		return err
 	}
 
-	var authResp authResponse
-	err = json.Unmarshal(msg, &authResp)
-	if err != nil {
-		return err
-	}
-	if authResp.MsgType != "auth_ok" {
+	if msg.MsgType != "auth_ok" {
 		return ErrInvalidToken
 	}
 
 	return nil
 }
 
-type SubEvent struct {
-	Id        int64  `json:"id"`
-	Type      string `json:"type"`
-	EventType string `json:"event_type"`
-}
-
-func SubscribeToStateChangedEvents(id int64, conn *WebsocketWriter, ctx context.Context) {
+func SubscribeToStateChangedEvents(id int64, conn *HAConnection, ctx context.Context) {
 	SubscribeToEventType("state_changed", conn, ctx, id)
 }
 
-func SubscribeToEventType(eventType string, conn *WebsocketWriter, ctx context.Context, id ...int64) {
+// TODO: Instead of using variadic arguments, just use a nillable pointer for the id
+func SubscribeToEventType(eventType string, conn *HAConnection, ctx context.Context, id ...int64) {
+	type SubEvent struct {
+		Id        int64  `json:"id"`
+		Type      string `json:"type"`
+		EventType string `json:"event_type"`
+	}
+
+	// If no id is provided, generate a new one
 	var finalId int64
 	if len(id) == 0 {
 		finalId = internal.NextId()
@@ -148,12 +170,12 @@ func SubscribeToEventType(eventType string, conn *WebsocketWriter, ctx context.C
 		Type:      "subscribe_events",
 		EventType: eventType,
 	}
+
 	err := conn.WriteMessage(e)
+	// TODO: Handle errors better
 	if err != nil {
 		wrappedErr := fmt.Errorf("error writing to websocket: %w", err)
 		slog.Error(wrappedErr.Error())
 		panic(wrappedErr)
 	}
-	// m, _ := ReadMessage(conn, ctx)
-	// log.Default().Println(string(m))
 }

internal/connect/listen.go

@@ -1,4 +1,4 @@
-package websocket
+package connect
 
 import (
 	"encoding/json"
@@ -7,22 +7,26 @@ import (
 	"github.com/gorilla/websocket"
 )
 
+// BaseMessage is the base message type for all messages sent by the websocket server.
 type BaseMessage struct {
 	Type    string `json:"type"`
 	Id      int64  `json:"id"`
-	Success bool   `json:"success"`
+	Success bool   `json:"success"` // not present in all messages
 }
 
-type ChanMsg struct {
+type ChannelMessage struct {
 	Id      int64
 	Type    string
 	Success bool
 	Raw     []byte
 }
 
-func ListenWebsocket(conn *websocket.Conn, c chan ChanMsg) {
+// ListenWebsocket reads messages from the websocket connection and sends them to the channel.
+// It will close the channel if it encounters an error, or if the channel is full, and return.
+// It ignores errors in deserialization.
+func ListenWebsocket(conn *websocket.Conn, c chan ChannelMessage) {
 	for {
-		bytes, err := ReadMessage(conn)
+		raw, err := ReadMessageRaw(conn)
 		if err != nil {
 			slog.Error("Error reading from websocket", "err", err)
 			close(c)
@@ -33,20 +37,26 @@ func ListenWebsocket(conn *websocket.Conn, c chan ChanMsg) {
 			// default to true for messages that don't include "success" at all
 			Success: true,
 		}
-		_ = json.Unmarshal(bytes, &base)
-		if !base.Success {
-			slog.Warn("Received unsuccessful response", "response", string(bytes))
+		err = json.Unmarshal(raw, &base)
+		if err != nil {
+			slog.Error("Error unmarshalling message", "err", err, "message", string(raw))
+			continue
 		}
-		chanMsg := ChanMsg{
+		if !base.Success {
+			slog.Warn("Received unsuccessful response", "response", string(raw))
+		}
+
+		// Create a channel message from the raw message
+		channelMessage := ChannelMessage{
 			Type:    base.Type,
 			Id:      base.Id,
 			Success: base.Success,
-			Raw:     bytes,
+			Raw:     raw,
 		}
 
 		// Use non-blocking send to avoid hanging on closed channel
 		select {
-		case c <- chanMsg:
+		case c <- channelMessage:
 			// Message sent successfully
 		default:
 			// Channel is full or closed, break out of loop

internal/http.go

@@ -4,6 +4,7 @@
 package internal
 
 import (
+	"context"
 	"errors"
 	"net/url"
 	"time"
@@ -12,40 +13,44 @@ import (
 )
 
 type HttpClient struct {
-	client *resty.Client
+	client      *resty.Client
+	baseRequest *resty.Request
 }
 
-func NewHttpClient(url *url.URL, token string) *HttpClient {
+func NewHttpClient(ctx context.Context, baseUrl *url.URL, token string) *HttpClient {
 	// Shallow copy the URL to avoid modifying the original
-	u := *url
+	u := *baseUrl
 	u.Path = "/api"
-	if u.Scheme == "ws" {
-		u.Scheme = "http"
-	}
-	if u.Scheme == "wss" {
-		u.Scheme = "https"
-	}
 
 	// Create resty client with configuration
 	client := resty.New().
 		SetBaseURL(u.String()).
-		SetHeader("Authorization", "Bearer "+token).
-		SetTimeout(30 * time.Second).
+		SetTimeout(30*time.Second).
 		SetRetryCount(3).
-		SetRetryWaitTime(1 * time.Second).
-		SetRetryMaxWaitTime(5 * time.Second).
+		SetRetryWaitTime(1*time.Second).
+		SetRetryMaxWaitTime(5*time.Second).
 		AddRetryConditions(func(r *resty.Response, err error) bool {
-			return err != nil || r.StatusCode() >= 500
-		})
+			return err != nil || (r.StatusCode() >= 500 && r.StatusCode() != 403)
+		}).
+		SetHeader("User-Agent", "go-ha/"+currentVersion).
+		SetContext(ctx)
 
 	return &HttpClient{
 		client: client,
+		baseRequest: client.R().
+			SetContentType("application/json").
+			SetHeader("Accept", "application/json").
+			SetAuthToken(token),
 	}
 }
 
+// getRequest returns a new request
+func (c *HttpClient) getRequest() *resty.Request {
+	return c.baseRequest.Clone(c.client.Context())
+}
+
 func (c *HttpClient) GetState(entityId string) ([]byte, error) {
-	resp, err := c.client.R().
-		Get("/states/" + entityId)
+	resp, err := c.getRequest().Get("/states/" + entityId)
 
 	if err != nil {
 		return nil, errors.New("Error making HTTP request: " + err.Error())
@@ -58,9 +63,9 @@ func (c *HttpClient) GetState(entityId string) ([]byte, error) {
 	return resp.Bytes(), nil
 }
 
-func (c *HttpClient) States() ([]byte, error) {
-	resp, err := c.client.R().
-		Get("/states")
+// GetStates returns the states of all entities.
+func (c *HttpClient) GetStates() ([]byte, error) {
+	resp, err := c.getRequest().Get("/states")
 
 	if err != nil {
 		return nil, errors.New("Error making HTTP request: " + err.Error())

internal/misc.go

@@ -1,6 +1,7 @@
 package internal
 
 import (
+	"fmt"
 	"reflect"
 	"runtime"
 	"sync/atomic"
@@ -12,6 +13,10 @@ type EnabledDisabledInfo struct {
 	RunOnError bool
 }
 
+var (
+	currentVersion = "0.7.0"
+)
+
 var (
 	id atomic.Int64 // default value is 0
 )
@@ -26,3 +31,20 @@ func NextId() int64 {
 func GetFunctionName(i interface{}) string {
 	return runtime.FuncForPC(reflect.ValueOf(i).Pointer()).Name()
 }
+
+// GetEquivalentWebsocketScheme returns the equivalent websocket scheme for the given scheme.
+// If the scheme is http or https, it returns ws or wss respectively.
+// If the scheme is ws or wss, it returns the same scheme.
+// If the scheme is not any of the above, it returns an error.
+func GetEquivalentWebsocketScheme(scheme string) (string, error) {
+	switch scheme {
+	case "http":
+		return "ws", nil
+	case "https":
+		return "wss", nil
+	case "ws", "wss":
+		return scheme, nil
+	default:
+		return "", fmt.Errorf("unexpected scheme: %s", scheme)
+	}
+}

internal/services/adaptive_lighting.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	"github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type AdaptiveLighting struct {
-	conn *ws.WebsocketWriter
+	conn *connect.HAConnection
 }
 
 /* Public API */

internal/services/alarm_control_panel.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type AlarmControlPanel struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 /* Public API */

internal/services/climate.go

@@ -1,14 +1,14 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	"github.com/Xevion/go-ha/internal/connect"
 	"github.com/Xevion/go-ha/types"
 )
 
 /* Structs */
 
 type Climate struct {
-	conn *ws.WebsocketWriter
+	conn *connect.HAConnection
 }
 
 /* Public API */

internal/services/cover.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	"github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type Cover struct {
-	conn *ws.WebsocketWriter
+	conn *connect.HAConnection
 }
 
 /* Public API */

internal/services/event.go

@@ -2,11 +2,11 @@ package services
 
 import (
 	"github.com/Xevion/go-ha/internal"
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	"github.com/Xevion/go-ha/internal/connect"
 )
 
 type Event struct {
-	conn *ws.WebsocketWriter
+	conn *connect.HAConnection
 }
 
 // Fire an event

internal/services/home_assistant.go

@@ -1,11 +1,11 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	"github.com/Xevion/go-ha/internal/connect"
 )
 
 type HomeAssistant struct {
-	conn *ws.WebsocketWriter
+	conn *connect.HAConnection
 }
 
 // TurnOn a Home Assistant entity. Takes an entityId and an optional

internal/services/input_boolean.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type InputBoolean struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 /* Public API */

internal/services/input_button.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type InputButton struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 /* Public API */

internal/services/input_datetime.go

@@ -4,13 +4,13 @@ import (
 	"fmt"
 	"time"
 
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type InputDatetime struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 /* Public API */

internal/services/input_number.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type InputNumber struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 /* Public API */

internal/services/input_text.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	"github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type InputText struct {
-	conn *ws.WebsocketWriter
+	conn *connect.HAConnection
 }
 
 /* Public API */

internal/services/light.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type Light struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 /* Public API */

internal/services/lock.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type Lock struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 /* Public API */

internal/services/media_player.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type MediaPlayer struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 /* Public API */

internal/services/notify.go

@@ -1,12 +1,12 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 	"github.com/Xevion/go-ha/types"
 )
 
 type Notify struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 // Notify sends a notification. Takes a types.NotifyRequest.

internal/services/number.go

@@ -1,11 +1,11 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 type Number struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 func (ib Number) SetValue(entityId string, value float32) error {

internal/services/scene.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type Scene struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 /* Public API */

internal/services/script.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type Script struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 /* Public API */

internal/services/services.go

@@ -2,7 +2,7 @@ package services
 
 import (
 	"github.com/Xevion/go-ha/internal"
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 func BuildService[
@@ -29,7 +29,7 @@ func BuildService[
 		Timer |
 		Vacuum |
 		ZWaveJS,
-](conn *ws.WebsocketWriter) *T {
+](conn *ws.HAConnection) *T {
 	return &T{conn: conn}
 }
 

internal/services/switch.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type Switch struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 /* Public API */

internal/services/timer.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type Timer struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 /* Public API */

internal/services/tts.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	connect "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type TTS struct {
-	conn *ws.WebsocketWriter
+	conn *connect.HAConnection
 }
 
 /* Public API */

internal/services/vacuum.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type Vacuum struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 /* Public API */

internal/services/zwavejs.go

@@ -1,13 +1,13 @@
 package services
 
 import (
-	ws "github.com/Xevion/go-ha/internal/websocket"
+	ws "github.com/Xevion/go-ha/internal/connect"
 )
 
 /* Structs */
 
 type ZWaveJS struct {
-	conn *ws.WebsocketWriter
+	conn *ws.HAConnection
 }
 
 /* Public API */