From 4d7fcbe23a1f4b84d484d01852740ff207bda507 Mon Sep 17 00:00:00 2001
From: Gabriel Adrian Samfira <gsamfira@cloudbasesolutions.com>
Date: Sat, 6 Jan 2024 14:15:52 +0000
Subject: [PATCH] Safely close the quit channel

Prevent accidental closure of an already closed channel.

Signed-off-by: Gabriel Adrian Samfira <gsamfira@cloudbasesolutions.com>
---
 websocket/websocket.go | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/websocket/websocket.go b/websocket/websocket.go
index a1791c09..a650088d 100644
--- a/websocket/websocket.go
+++ b/websocket/websocket.go
@@ -35,17 +35,19 @@ type Hub struct {
 	// Unregister requests from clients.
 	unregister chan *Client
 
-	mux sync.Mutex
+	mux  sync.Mutex
+	once sync.Once
 }
 
 func (h *Hub) run() {
+	defer func() {
+		close(h.closed)
+	}()
 	for {
 		select {
 		case <-h.quit:
-			close(h.closed)
 			return
 		case <-h.ctx.Done():
-			close(h.closed)
 			return
 		case client := <-h.register:
 			if client != nil {
@@ -116,8 +118,15 @@ func (h *Hub) Start() error {
 	return nil
 }
 
+func (h *Hub) Close() error {
+	h.once.Do(func() {
+		close(h.quit)
+	})
+	return nil
+}
+
 func (h *Hub) Stop() error {
-	close(h.quit)
+	h.Close()
 	select {
 	case <-h.closed:
 		return nil