From 23f87fef8b7646e4646d98e6516f1c3e079b1608 Mon Sep 17 00:00:00 2001
From: Lubos Dolezel <lubos@dolezel.info>
Date: Tue, 7 Mar 2023 18:18:14 +0100
Subject: [PATCH] Avoid panic on disconnect

---
 app.go                       |  7 +++++--
 internal/websocket/reader.go | 10 +++++++++-
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/app.go b/app.go
index 4f6bda6..b06dc11 100644
--- a/app.go
+++ b/app.go
@@ -263,9 +263,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 {
diff --git a/internal/websocket/reader.go b/internal/websocket/reader.go
index ad0bdf0..d58e15b 100644
--- a/internal/websocket/reader.go
+++ b/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{