From a02a1b2ab6cc3e3eb2f08632995c0dab56228c2f Mon Sep 17 00:00:00 2001
From: Boran Car <boran.car@gmail.com>
Date: Fri, 26 Aug 2022 19:21:23 +0300
Subject: [PATCH] Clean up connections after websocket close

Make sure we clean up connections by just going through the map and
cleaning all the verkeys that match the connection. This is a suboptimal
solution, but a connection close is not something happennig often.
---
 pkg/didcomm/transport/ws/pool.go | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/pkg/didcomm/transport/ws/pool.go b/pkg/didcomm/transport/ws/pool.go
index e558a881c..bdcd54679 100644
--- a/pkg/didcomm/transport/ws/pool.go
+++ b/pkg/didcomm/transport/ws/pool.go
@@ -79,9 +79,7 @@ func (d *connPool) remove(verKey string) {
 }
 
 func (d *connPool) listener(conn *websocket.Conn, outbound bool) {
-	verKeys := []string{}
-
-	defer d.close(conn, verKeys)
+	defer d.close(conn)
 
 	go keepConnAlive(conn, outbound, pingFrequency)
 
@@ -154,14 +152,25 @@ func (d *connPool) addKey(unpackMsg *transport.Envelope, trans *decorator.Transp
 	}
 }
 
-func (d *connPool) close(conn *websocket.Conn, verKeys []string) {
+func (d *connPool) close(conn *websocket.Conn) {
 	if err := conn.Close(websocket.StatusNormalClosure,
 		"closing the connection"); websocket.CloseStatus(err) != websocket.StatusNormalClosure {
 		logger.Errorf("connection close error")
 	}
 
-	for _, v := range verKeys {
-		d.remove(v)
+	d.Lock()
+	defer d.Unlock()
+
+	var toRemove []string
+
+	for k, v := range d.connMap {
+		if v == conn {
+			toRemove = append(toRemove, k)
+		}
+	}
+
+	for _, v := range toRemove {
+		delete(d.connMap, v)
 	}
 }