internal/websocket/client_commands.go

@@ -10,15 +10,15 @@ import (
 )
 
 // Listens for commands from a client and handles them
-func handleClientCommands(c models.WebsocketClient) {
+func handleClientCommands(c models.Client) {
 	ws := c.Conn
 	for {
 		var clientCommand models.ClientCommand
 
 		_, message, err := ws.ReadMessage()
 		if err != nil {
-			log.Info("Client disconnected:", "msg", err, "host", ws.RemoteAddr(), "clients", len(Clients)-1)
+			log.Info("Client disconnected:", "msg", err, "host", ws.NetConn().RemoteAddr(), "clients", len(Clients)-1)
-			delete(Clients, ws.LocalAddr())
+			delete(Clients, ws)
 			break
 		}
 

internal/websocket/handle_connection.go

@@ -3,15 +3,14 @@ package websocket
 import (
 	"github.com/charmbracelet/log"
 	"github.com/gorilla/websocket"
-	"net"
 	"net/http"
 	"sync"
 
 	"git.smsvc.net/pomodoro/GoTomato/pkg/models"
 )
 
-// Clients is a map of connected WebSocket clients, where each client is represented by the WebsocketClient struct
+// Clients is a map of connected WebSocket clients, where each client is represented by the Client struct
-var Clients = make(map[net.Addr]*models.WebsocketClient)
+var Clients = make(map[*websocket.Conn]*models.Client)
 var mu sync.Mutex // Mutex to protect access to the Clients map
 
 // Upgrade HTTP requests to WebSocket connections
@@ -29,14 +28,14 @@ func HandleConnection(w http.ResponseWriter, r *http.Request) {
 	}
 	defer ws.Close()
 
-	log.Info("Client connected", "host", ws.RemoteAddr(), "clients", len(Clients)+1)
+	log.Info("Client connected", "host", ws.NetConn().RemoteAddr(), "clients", len(Clients)+1)
 
 	// Register the new client
-	client := models.WebsocketClient{
+	client := models.Client{
 		Conn: ws,
 	}
 	mu.Lock()
-	Clients[ws.LocalAddr()] = &client
+	Clients[ws] = &client
 	mu.Unlock()
 
 	// Listen for commands from the connected client

pkg/models/client.go

@@ -13,13 +13,13 @@ type ClientCommand struct {
 }
 
 // Represents a single client
-type WebsocketClient struct {
+type Client struct {
 	Conn *websocket.Conn
 }
 
 // Sends a message to the websocket.
 // Automatically locks and unlocks the client mutex, to ensure that only one goroutine can write at a time.
-func (c *WebsocketClient) SendMessage(messageType int, data []byte) error {
+func (c *Client) SendMessage(messageType int, data []byte) error {
 	err := c.Conn.WriteMessage(messageType, data)
 	if err != nil {
 		log.Error("Error writing to WebSocket:", "msg", err)