From da1d073846c6c41a16a5a6fcf524b3bd86b25e40 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Tue, 23 Mar 2021 15:38:37 +0100
Subject: [PATCH] websocket connection change.

---
 internal/session/session.go   |  6 +++++-
 internal/types/session.go     |  2 +-
 internal/websocket/manager.go | 10 ++++++----
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/internal/session/session.go b/internal/session/session.go
index 7d4b8a1a..6497543d 100644
--- a/internal/session/session.go
+++ b/internal/session/session.go
@@ -67,7 +67,11 @@ func (session *SessionCtx) SetWebSocketPeer(websocketPeer types.WebSocketPeer) {
 	session.websocketPeer = websocketPeer
 }
 
-func (session *SessionCtx) SetWebSocketConnected(connected bool) {
+func (session *SessionCtx) SetWebSocketConnected(websocketPeer types.WebSocketPeer, connected bool) {
+	if websocketPeer != session.websocketPeer {
+		return
+	}
+
 	session.state.IsConnected = connected
 
 	if connected {
diff --git a/internal/types/session.go b/internal/types/session.go
index 94a6c68a..19d22031 100644
--- a/internal/types/session.go
+++ b/internal/types/session.go
@@ -15,7 +15,7 @@ type Session interface {
 
 	// websocket
 	SetWebSocketPeer(websocketPeer WebSocketPeer)
-	SetWebSocketConnected(connected bool)
+	SetWebSocketConnected(websocketPeer WebSocketPeer, connected bool)
 	Send(v interface{}) error
 	Disconnect(reason string) error
 
diff --git a/internal/websocket/manager.go b/internal/websocket/manager.go
index 9bf9bbab..22489928 100644
--- a/internal/websocket/manager.go
+++ b/internal/websocket/manager.go
@@ -207,11 +207,13 @@ func (manager *WebSocketManagerCtx) Upgrade(w http.ResponseWriter, r *http.Reque
 		return
 	}
 
-	session.SetWebSocketPeer(&WebSocketPeerCtx{
+	peer := &WebSocketPeerCtx{
 		session:    session,
 		manager:    manager,
 		connection: connection,
-	})
+	}
+
+	session.SetWebSocketPeer(peer)
 
 	manager.logger.
 		Debug().
@@ -219,7 +221,7 @@ func (manager *WebSocketManagerCtx) Upgrade(w http.ResponseWriter, r *http.Reque
 		Str("address", connection.RemoteAddr().String()).
 		Msg("connection started")
 
-	session.SetWebSocketConnected(true)
+	session.SetWebSocketConnected(peer, true)
 
 	defer func() {
 		manager.logger.
@@ -228,7 +230,7 @@ func (manager *WebSocketManagerCtx) Upgrade(w http.ResponseWriter, r *http.Reque
 			Str("address", connection.RemoteAddr().String()).
 			Msg("connection ended")
 
-		session.SetWebSocketConnected(false)
+		session.SetWebSocketConnected(peer, false)
 	}()
 
 	manager.handle(connection, session)