From 58165208162168c609c66e20f0219731553a1c98 Mon Sep 17 00:00:00 2001
From: Jason Ertel <jason.ertel@securityonionsolutions.com>
Date: Wed, 23 Jun 2021 11:29:42 -0400
Subject: [PATCH] Fix overly aggressive websocket cleanup. Resolves #4598

---
 web/host.go      |  5 +++--
 web/host_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/web/host.go b/web/host.go
index 49f3d68d..54df6e2c 100644
--- a/web/host.go
+++ b/web/host.go
@@ -114,12 +114,13 @@ func (host *Host) AddConnection(wsConn *websocket.Conn, ip string) *Connection {
 func (host *Host) RemoveConnection(wsConn *websocket.Conn) {
   host.lock.Lock();
   defer host.lock.Unlock()
-  host.connections = make([]*Connection, 0)
+  remaining := make([]*Connection, 0)
   for _, connection := range host.connections {
     if connection.websocket != wsConn {
-      host.connections = append(host.connections, connection)
+      remaining = append(remaining, connection)
     }
   }
+  host.connections = remaining
   log.WithField("Connections", len(host.connections)).Debug("Removed WebSocket connection")
 }
 
diff --git a/web/host_test.go b/web/host_test.go
index 3fcd63b0..12b7c513 100644
--- a/web/host_test.go
+++ b/web/host_test.go
@@ -43,6 +43,31 @@ func TestAddRemoveConnection(tester *testing.T) {
 	})
 }
 
+func TestMultipleConnections(tester *testing.T) {
+  host := NewHost("http://some.where/path", "/tmp/foo", 123, "unit test")
+	conn1 := &websocket.Conn{}
+	conn2 := &websocket.Conn{}
+	tester.Run("testing add multiple connections", func(t *testing.T) {
+	  host.AddConnection(conn1, "1.2.3.4");
+	  host.AddConnection(conn2, "1.2.3.4");
+		if len(host.connections) != 2 {
+			tester.Errorf("expected %d but got %d", 2, len(host.connections))
+		}
+	})
+	tester.Run("testing remove first connection", func(t *testing.T) {
+		host.RemoveConnection(conn1);
+		if len(host.connections) != 1 {
+			t.Errorf("final expected %d but got %d", 1, len(host.connections))
+		}
+	})
+	tester.Run("testing remove second connection", func(t *testing.T) {
+		host.RemoveConnection(conn2);
+		if len(host.connections) != 0 {
+			t.Errorf("final expected %d but got %d", 0, len(host.connections))
+		}
+	})
+}
+
 func TestManageConnections(tester *testing.T) {
   host := NewHost("http://some.where/path", "/tmp/foo", 123, "unit test")
 	conn := host.AddConnection(nil, "")