From 74236258dbb2632f164e0474894fbf372e7c55ba Mon Sep 17 00:00:00 2001
From: Ralph Slooten <axllent@gmail.com>
Date: Sat, 11 Nov 2023 15:32:57 +1300
Subject: [PATCH] Fix: Correctly close websockets on client disconnect (#207)

---
 server/websockets/client.go | 25 +++++++++++++++++++++----
 server/websockets/hub.go    | 12 ++++++------
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/server/websockets/client.go b/server/websockets/client.go
index 51e7bb3..efbc634 100644
--- a/server/websockets/client.go
+++ b/server/websockets/client.go
@@ -53,7 +53,24 @@ type Client struct {
 	send chan []byte
 }
 
-// writePump pumps messages from the hub to the websocket connection.
+// ReadPump is used here solely to monitor the connection, not to actually receive messages.
+func (c *Client) readPump() {
+	defer func() {
+		c.hub.unregister <- c
+	}()
+
+	for {
+		_, _, err := c.conn.ReadMessage()
+		if err != nil {
+			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
+				logger.Log().Errorf("[websocket] error: %v", err)
+			}
+			break
+		}
+	}
+}
+
+// WritePump pumps messages from the hub to the websocket connection.
 //
 // A goroutine running writePump is started for each connection. The
 // application ensures that there is at most one writer to a connection by
@@ -62,7 +79,7 @@ func (c *Client) writePump() {
 	ticker := time.NewTicker(pingPeriod)
 	defer func() {
 		ticker.Stop()
-		_ = c.conn.Close()
+		c.hub.unregister <- c
 	}()
 	for {
 		select {
@@ -122,8 +139,8 @@ func ServeWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
 	client := &Client{hub: hub, conn: conn, send: make(chan []byte, 256)}
 	client.hub.register <- client
 
-	// Allow collection of memory referenced by the caller by doing all work in
-	// new goroutines.
+	// Allow collection of memory referenced by the caller by doing all work in new goroutines.
+	go client.readPump()
 	go client.writePump()
 }
 
diff --git a/server/websockets/hub.go b/server/websockets/hub.go
index 218a03e..cbaae88 100644
--- a/server/websockets/hub.go
+++ b/server/websockets/hub.go
@@ -1,7 +1,4 @@
-// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
-// Use of this source code is governed by a BSD-style
-// license that can be found in the LICENSE file.
-
+// Package websockets is used to broadcast messages to connected clients
 package websockets
 
 import (
@@ -47,14 +44,17 @@ func (h *Hub) Run() {
 	for {
 		select {
 		case client := <-h.register:
-			h.Clients[client] = true
+			if _, ok := h.Clients[client]; !ok {
+				logger.Log().Debugf("[websocket] client %s connected", client.conn.RemoteAddr().String())
+				h.Clients[client] = true
+			}
 		case client := <-h.unregister:
 			if _, ok := h.Clients[client]; ok {
+				logger.Log().Debugf("[websocket] client %s disconnected", client.conn.RemoteAddr().String())
 				delete(h.Clients, client)
 				close(client.send)
 			}
 		case message := <-h.Broadcast:
-			// logger.Log().Debugf("[broadcast] %s", message)
 			for client := range h.Clients {
 				select {
 				case client.send <- message: