From 58d029f455e58123d1dba8a7e073c59a3ea37fae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Tue, 16 Mar 2021 10:45:22 +0100
Subject: [PATCH] WS upgrade no returning error.

---
 internal/http/manager.go      |  1 -
 internal/types/websocket.go   |  2 +-
 internal/websocket/manager.go | 35 +++++++++++++++++++++++++----------
 3 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/internal/http/manager.go b/internal/http/manager.go
index 185df234..6084eb52 100644
--- a/internal/http/manager.go
+++ b/internal/http/manager.go
@@ -44,7 +44,6 @@ func New(WebSocketManager types.WebSocketManager, ApiManager types.ApiManager, c
 	router.Route("/api", ApiManager.Route)
 
 	router.Get("/api/ws", func(w http.ResponseWriter, r *http.Request) {
-		//nolint
 		WebSocketManager.Upgrade(w, r, func(r *http.Request) bool {
 			return conf.AllowOrigin(r.Header.Get("Origin"))
 		})
diff --git a/internal/types/websocket.go b/internal/types/websocket.go
index 154d5b74..9953b598 100644
--- a/internal/types/websocket.go
+++ b/internal/types/websocket.go
@@ -15,5 +15,5 @@ type WebSocketManager interface {
 	Start()
 	Shutdown() error
 	AddHandler(handler HandlerFunction)
-	Upgrade(w http.ResponseWriter, r *http.Request, checkOrigin CheckOrigin) error
+	Upgrade(w http.ResponseWriter, r *http.Request, checkOrigin CheckOrigin)
 }
diff --git a/internal/websocket/manager.go b/internal/websocket/manager.go
index f8a95e58..e3a0b76d 100644
--- a/internal/websocket/manager.go
+++ b/internal/websocket/manager.go
@@ -138,7 +138,7 @@ func (manager *WebSocketManagerCtx) AddHandler(handler types.HandlerFunction) {
 	manager.handlers = append(manager.handlers, handler)
 }
 
-func (manager *WebSocketManagerCtx) Upgrade(w http.ResponseWriter, r *http.Request, checkOrigin types.CheckOrigin) error {
+func (manager *WebSocketManagerCtx) Upgrade(w http.ResponseWriter, r *http.Request, checkOrigin types.CheckOrigin) {
 	manager.logger.Debug().Msg("attempting to upgrade connection")
 
 	upgrader := websocket.Upgrader{
@@ -148,12 +148,12 @@ func (manager *WebSocketManagerCtx) Upgrade(w http.ResponseWriter, r *http.Reque
 	connection, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		manager.logger.Error().Err(err).Msg("failed to upgrade connection")
-		return err
+		return
 	}
 
 	session, err := manager.sessions.Authenticate(r)
 	if err != nil {
-		manager.logger.Warn().Err(err).Msg("authentication failed")
+		manager.logger.Debug().Err(err).Msg("authentication failed")
 
 		// TODO: Refactor, return error code.
 		if err = connection.WriteJSON(
@@ -161,36 +161,52 @@ func (manager *WebSocketManagerCtx) Upgrade(w http.ResponseWriter, r *http.Reque
 				Event:   event.SYSTEM_DISCONNECT,
 				Message: err.Error(),
 			}); err != nil {
-			manager.logger.Error().Err(err).Msg("failed to send disconnect")
+			manager.logger.Error().Err(err).Msg("failed to send disconnect event")
 		}
 
-		return connection.Close()
+		if err := connection.Close(); err != nil {
+			manager.logger.Warn().Err(err).Msg("connection closed with an error")
+		}
+
+		return
 	}
 
 	if !session.Profile().CanConnect {
+		manager.logger.Debug().Str("session_id", session.ID()).Msg("connection disabled")
+
 		// TODO: Refactor, return error code.
 		if err = connection.WriteJSON(
 			message.SystemDisconnect{
 				Event:   event.SYSTEM_DISCONNECT,
 				Message: "connection disabled",
 			}); err != nil {
-			manager.logger.Error().Err(err).Msg("failed to send disconnect")
+			manager.logger.Error().Err(err).Msg("failed to send disconnect event")
 		}
 
-		return connection.Close()
+		if err := connection.Close(); err != nil {
+			manager.logger.Warn().Err(err).Msg("connection closed with an error")
+		}
+
+		return
 	}
 
 	if session.State().IsConnected {
+		manager.logger.Debug().Str("session_id", session.ID()).Msg("already connected")
+
 		// TODO: Refactor, return error code.
 		if err = connection.WriteJSON(
 			message.SystemDisconnect{
 				Event:   event.SYSTEM_DISCONNECT,
 				Message: "already connected",
 			}); err != nil {
-			manager.logger.Error().Err(err).Msg("failed to send disconnect")
+			manager.logger.Error().Err(err).Msg("failed to send disconnect event")
 		}
 
-		return connection.Close()
+		if err := connection.Close(); err != nil {
+			manager.logger.Warn().Err(err).Msg("connection closed with an error")
+		}
+
+		return
 	}
 
 	session.SetWebSocketPeer(&WebSocketPeerCtx{
@@ -218,7 +234,6 @@ func (manager *WebSocketManagerCtx) Upgrade(w http.ResponseWriter, r *http.Reque
 	}()
 
 	manager.handle(connection, session)
-	return nil
 }
 
 func (manager *WebSocketManagerCtx) handle(connection *websocket.Conn, session types.Session) {