From 24dc92e5644909cb7816dab0c35d61a0abec6c04 Mon Sep 17 00:00:00 2001
From: CWX <cwxdev@outlook.com>
Date: Mon, 20 Jan 2025 20:19:53 +0000
Subject: [PATCH] Fix WS issues

---
 Libraries/Core/Servers/WebSocketServer.cs        |  5 ++++-
 .../Servers/Ws/SptWebSocketConnectionHandler.cs  | 16 +++++++++++++---
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/Libraries/Core/Servers/WebSocketServer.cs b/Libraries/Core/Servers/WebSocketServer.cs
index 794ec638..538fb08d 100644
--- a/Libraries/Core/Servers/WebSocketServer.cs
+++ b/Libraries/Core/Servers/WebSocketServer.cs
@@ -34,7 +34,10 @@ public class WebSocketServer(
         foreach (var wsh in socketHandlers)
         {
             wsh.OnConnection(webSocket, context).Wait();
-            _logger.Info($"WebSocketHandler \"{wsh.GetSocketId()}\" connected");
+            if (webSocket.State == WebSocketState.Open)
+            {
+                _logger.Info($"WebSocketHandler \"{wsh.GetSocketId()}\" connected");
+            }
         }
 
         return Task.CompletedTask;
diff --git a/Libraries/Core/Servers/Ws/SptWebSocketConnectionHandler.cs b/Libraries/Core/Servers/Ws/SptWebSocketConnectionHandler.cs
index 560ff776..4708600c 100644
--- a/Libraries/Core/Servers/Ws/SptWebSocketConnectionHandler.cs
+++ b/Libraries/Core/Servers/Ws/SptWebSocketConnectionHandler.cs
@@ -145,13 +145,23 @@ public class SptWebSocketConnectionHandler(
                     sptWebSocketMessageHandler.OnSptMessage(sessionID, ws, readBytes.ToArray()).Wait();
                 }
             }
-            catch (OperationCanceledException e)
+            catch (OperationCanceledException _)
             {
                 _logger.Info("WebSocket disconnecting, receive task finalized...");
             }
-            catch (Exception e)
+            catch(Exception _)
             {
-                _logger.Error("Error reading data from WebSocket", e);
+                lock (_lockObject)
+                {
+                    _sockets.Remove(sessionID);
+                    _socketAliveTimers.Remove(sessionID);
+                    _receiveTasks.Remove(sessionID);
+                    var playerProfile = _profileHelper.GetFullProfile(sessionID);
+                    var playerInfoText = $"{playerProfile.ProfileInfo.Username} ({sessionID})";
+                    _logger.Info($"[ws] player: {playerInfoText} has disconnected");
+                }
+                
+                ws.CloseAsync(WebSocketCloseStatus.NormalClosure, "Client closed connection", CancellationToken.None);
             }
             finally
             {