fix: hopefully fix sockt

Signed-off-by: SIGMazer <[email protected]>
0xCator · May 10, 2024 · 08b90e4 · 08b90e4
1 parent bfea46c
commit 08b90e4
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 9 deletions.
diff --git a/Program.cs b/Program.cs
@@ -86,9 +86,9 @@
         var userID = context.Request.Query["id"];
         var webSocket = await context.WebSockets.AcceptWebSocketAsync();
 
-        Socket socket = new Socket(int.Parse(userID));
+        Socket socket = new Socket(int.Parse(userID), webSocket);
         SocketManager.Instance.AddSocket(socket);
-        await socket.Start(webSocket);
+        await socket.Start();
     }
     else
     {

diff --git a/WebSocket/Socket.cs b/WebSocket/Socket.cs
@@ -12,19 +12,19 @@ public class Socket
         public System.Net.WebSockets.WebSocket socket;
         public int id;
 
-        public Socket(int id)
+        public Socket(int id, System.Net.WebSockets.WebSocket socket)
         {
             this.id = id;
+            this.socket = socket;
             Console.WriteLine("Socket started with id: " + id);
         }
 
-        public async Task Start(System.Net.WebSockets.WebSocket socket)
+        public async Task Start()
         {
-            this.socket = socket;
             var buffer = new byte[1024 * 4];
             try
             {
-                while (socket.State == WebSocketState.Open)
+                while (this.socket.State == WebSocketState.Open)
                 {
                     WebSocketReceiveResult result = await socket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
                     if (result.MessageType == WebSocketMessageType.Close)
@@ -36,12 +36,14 @@ public async Task Start(System.Net.WebSockets.WebSocket socket)
                     else
                     {
                         var message = Encoding.UTF8.GetString(buffer, 0, result.Count);
+                        Array.Clear(buffer, 0, buffer.Length);
                     }
                 }
             }
             catch (WebSocketException)
             {
-
+                SocketManager.Instance.RemoveSocket(this);
+                Console.WriteLine("Socket closed with id: " + id);
             }
         }
     }

diff --git a/WebSocket/SocketManager.cs b/WebSocket/SocketManager.cs
@@ -70,9 +70,19 @@ public async Task TaskHasUpdate(ICollection<User> users)
             {
                 foreach (var user in users)
                 {
-                    if (socket.id == user.Id)
+                    if(socket.socket.State == WebSocketState.Open)
                     {
-                        await socket.socket.SendAsync(new ArraySegment<byte>(message, 0, message.Length), WebSocketMessageType.Text, true, CancellationToken.None);
+                        if (socket.id == user.Id)
+                        {
+                            try
+                            {
+                                await socket.socket.SendAsync(new ArraySegment<byte>(message, 0, message.Length), WebSocketMessageType.Text, true, CancellationToken.None);
+                            }
+                            catch (WebSocketException ex)
+                            {
+                                Console.WriteLine(ex.Message);
+                            }
+                        }
                     }
                 }
             }