From 1dab0bd8599e43b8b6a3ab1013d1c7573d672b95 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Thu, 2 Sep 2021 20:30:50 +0200
Subject: [PATCH] remove session from websocket peer.

---
 internal/websocket/manager.go | 37 ++++++++++-------------------------
 internal/websocket/peer.go    | 18 ++++++++++++++++-
 2 files changed, 27 insertions(+), 28 deletions(-)

diff --git a/internal/websocket/manager.go b/internal/websocket/manager.go
index 9f15fd3d..6e9c6476 100644
--- a/internal/websocket/manager.go
+++ b/internal/websocket/manager.go
@@ -156,43 +156,26 @@ func (manager *WebSocketManagerCtx) Upgrade(w http.ResponseWriter, r *http.Reque
 		return
 	}
 
+	// create new peer
+	peer := newPeer(connection)
+
 	session, err := manager.sessions.Authenticate(r)
 	if err != nil {
 		manager.logger.Warn().Err(err).Msg("authentication failed")
 
-		// TODO: Better handling...
-		raw, err := json.Marshal(message.SystemDisconnect{
-			Message: err.Error(),
-		})
-
-		if err != nil {
-			manager.logger.Err(err).Msg("failed to create disconnect event")
-		}
-
-		err = connection.WriteJSON(
-			types.WebSocketMessage{
-				Event:   event.SYSTEM_DISCONNECT,
-				Payload: raw,
+		peer.Send(
+			event.SYSTEM_DISCONNECT,
+			message.SystemDisconnect{
+				Message: err.Error(),
 			})
 
-		if err != nil {
-			manager.logger.Err(err).Msg("failed to send disconnect event")
-		}
-
-		err = connection.Close()
-		manager.logger.Err(err).Msg("connection closed")
+		peer.Destroy()
 		return
 	}
 
-	// use session id with defeault logger context
+	// add session id to all log messages
 	logger := manager.logger.With().Str("session_id", session.ID()).Logger()
-
-	// create new peer
-	peer := &WebSocketPeerCtx{
-		logger:     logger,
-		session:    session,
-		connection: connection,
-	}
+	peer.setSessionID(session.ID())
 
 	if !session.Profile().CanConnect {
 		logger.Warn().Msg("connection disabled")
diff --git a/internal/websocket/peer.go b/internal/websocket/peer.go
index 792574c0..7ac5790c 100644
--- a/internal/websocket/peer.go
+++ b/internal/websocket/peer.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/gorilla/websocket"
 	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
 
 	"demodesk/neko/internal/types"
 	"demodesk/neko/internal/types/event"
@@ -15,10 +16,25 @@ import (
 type WebSocketPeerCtx struct {
 	mu         sync.Mutex
 	logger     zerolog.Logger
-	session    types.Session
 	connection *websocket.Conn
 }
 
+func newPeer(connection *websocket.Conn) *WebSocketPeerCtx {
+	logger := log.With().
+		Str("module", "websocket").
+		Str("submodule", "peer").
+		Logger()
+
+	return &WebSocketPeerCtx{
+		logger:     logger,
+		connection: connection,
+	}
+}
+
+func (peer *WebSocketPeerCtx) setSessionID(sessionId string) {
+	peer.logger = peer.logger.With().Str("session_id", sessionId).Logger()
+}
+
 func (peer *WebSocketPeerCtx) Send(event string, payload interface{}) {
 	peer.mu.Lock()
 	defer peer.mu.Unlock()