From 8b3ff7ac260531d1fcd9024dc5c034bf458569a4 Mon Sep 17 00:00:00 2001
From: 0x0ff <0x0ff@onsats.org>
Date: Thu, 29 Dec 2022 15:02:59 +0100
Subject: [PATCH] fix: remove listeners when client disconnects

---
 handlers.go |  3 ++-
 listener.go | 14 +++++++++++++-
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/handlers.go b/handlers.go
index 43029d8..7875f20 100644
--- a/handlers.go
+++ b/handlers.go
@@ -60,6 +60,7 @@ func (s *Server) handleWebsocket(w http.ResponseWriter, r *http.Request) {
 			if _, ok := s.clients[conn]; ok {
 				conn.Close()
 				delete(s.clients, conn)
+				removeListener(ws)
 			}
 			s.clientsMu.Unlock()
 		}()
@@ -214,7 +215,7 @@ func (s *Server) handleWebsocket(w http.ResponseWriter, r *http.Request) {
 						return
 					}
 
-					removeListener(ws, id)
+					removeListenerId(ws, id)
 					break
 				default:
 					if cwh, ok := s.relay.(CustomWebSocketHandler); ok {
diff --git a/listener.go b/listener.go
index 3c20ca5..2d3ab6a 100644
--- a/listener.go
+++ b/listener.go
@@ -63,7 +63,8 @@ func setListener(id string, ws *WebSocket, filters nostr.Filters) {
 	}
 }
 
-func removeListener(ws *WebSocket, id string) {
+// Remove a specific subscription id from listeners for a given ws client
+func removeListenerId(ws *WebSocket, id string) {
 	listenersMutex.Lock()
 	defer func() {
 		listenersMutex.Unlock()
@@ -78,6 +79,17 @@ func removeListener(ws *WebSocket, id string) {
 	}
 }
 
+// Remove WebSocket conn from listeners
+func removeListener(ws *WebSocket) {
+	listenersMutex.Lock()
+	defer listenersMutex.Unlock()
+
+	_, ok := listeners[ws]
+	if ok {
+		delete(listeners, ws)
+	}
+}
+
 func notifyListeners(event *nostr.Event) {
 	listenersMutex.Lock()
 	defer func() {