From 390190cc35f39aedc5084cb2e12f93734c89beae Mon Sep 17 00:00:00 2001
From: DTieman <djurre.tieman@gmail.com>
Date: Mon, 13 Mar 2023 12:27:41 +0100
Subject: [PATCH] Added room functionality

---
 Controllers/WebSocketsController.cs | 52 +++++++++++++++++++++--------
 Program.cs                          |  3 ++
 RoomManager.cs                      |  6 ++++
 Websockets/WebsocketManager.cs      | 38 +++++++++++++++++++++
 4 files changed, 86 insertions(+), 13 deletions(-)
 create mode 100644 RoomManager.cs
 create mode 100644 Websockets/WebsocketManager.cs

diff --git a/Controllers/WebSocketsController.cs b/Controllers/WebSocketsController.cs
index 2f4944e..c786fef 100644
--- a/Controllers/WebSocketsController.cs
+++ b/Controllers/WebSocketsController.cs
@@ -1,5 +1,6 @@
 using System.Net.WebSockets;
 using System.Text;
+using MauMau_Server.Websockets;
 using Microsoft.AspNetCore.Mvc;
 
 namespace MauMau_Server.Controllers;
@@ -9,10 +10,13 @@ namespace MauMau_Server.Controllers;
 public class WebSocketsController : ControllerBase
 {
     private readonly ILogger<WebSocketsController> _logger;
+    private readonly IWebsocketManager _websocketManager;
+    private string _id = string.Empty;
 
-    public WebSocketsController(ILogger<WebSocketsController> logger)
+    public WebSocketsController(ILogger<WebSocketsController> logger, IWebsocketManager websocketManager)
     {
         _logger = logger;
+        _websocketManager = websocketManager;
     }
 
     [HttpGet("/ws")]
@@ -21,7 +25,8 @@ public class WebSocketsController : ControllerBase
         if (HttpContext.WebSockets.IsWebSocketRequest)
         {
             using var webSocket = await HttpContext.WebSockets.AcceptWebSocketAsync();
-            _logger.Log(LogLevel.Information, "WebSocket connection established");
+            _id = _websocketManager.AddConnection(webSocket);
+            _logger.Log(LogLevel.Information, $"WebSocket connection {_id} established");
             await Echo(webSocket);
         }
         else
@@ -32,22 +37,43 @@ public class WebSocketsController : ControllerBase
         
     private async Task Echo(WebSocket webSocket)
     {
-        var buffer = new byte[1024 * 4];
-        var result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
-        _logger.Log(LogLevel.Information, "Message received from Client");
+        var buffer = EmptyBuffer();
+        var result = await ReceiveAsync(webSocket, buffer);
 
         while (!result.CloseStatus.HasValue)
         {
-            var serverMsg = Encoding.UTF8.GetBytes($"Server: Hello. You said: {Encoding.UTF8.GetString(buffer)}");
-            await webSocket.SendAsync(new ArraySegment<byte>(serverMsg, 0, serverMsg.Length), result.MessageType, result.EndOfMessage, CancellationToken.None);
-            _logger.Log(LogLevel.Information, "Message sent to Client");
+            var message = Encoding.UTF8.GetString(buffer);
+            BroadcastAsync(message);
 
-            buffer = new byte[1024 * 4];
-            result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
-            _logger.Log(LogLevel.Information, "Message received from Client");
-                
+            buffer = EmptyBuffer();
+            result = await ReceiveAsync(webSocket, buffer);
         }
         await webSocket.CloseAsync(result.CloseStatus.Value, result.CloseStatusDescription, CancellationToken.None);
-        _logger.Log(LogLevel.Information, "WebSocket connection closed");
+        _websocketManager.RemoveConnection(_id);
+        _logger.Log(LogLevel.Information, $"WebSocket connection {_id} closed");
+    }
+    
+    private async Task<WebSocketReceiveResult?> ReceiveAsync(WebSocket webSocket, byte[] buffer)
+    {
+        var arraySegment = new ArraySegment<byte>(buffer);
+        var result = await webSocket.ReceiveAsync(arraySegment, CancellationToken.None);
+        _logger.Log(LogLevel.Information, $"Message received from Client {_id}");
+        return result;
+    }
+    
+    private void BroadcastAsync(string message)
+    {
+        var bytes = Encoding.Default.GetBytes($"{_id}: {message}");
+        var arraySegment = new ArraySegment<byte>(bytes);
+        foreach (var (id, socket) in _websocketManager.GetAllConnections())
+        {
+            socket.SendAsync(arraySegment, WebSocketMessageType.Text, true, CancellationToken.None);
+            _logger.Log(LogLevel.Information, $"Message from Client {_id} sent to Client {id}");
+        }
+    }
+    
+    private static byte[] EmptyBuffer()
+    {
+        return new byte[4096];
     }
 }
\ No newline at end of file
diff --git a/Program.cs b/Program.cs
index 1230db5..e45d0a5 100644
--- a/Program.cs
+++ b/Program.cs
@@ -1,3 +1,5 @@
+using MauMau_Server.Websockets;
+
 var builder = WebApplication.CreateBuilder(args);
 
 // Add services to the container.
@@ -6,6 +8,7 @@ services.AddControllers();
 // Learn more about configuring Swagger/OpenAPI at https://aka.ms/aspnetcore/swashbuckle
 services.AddEndpointsApiExplorer();
 services.AddSwaggerGen();
+services.AddScoped<IWebsocketManager, WebsocketManager>();
 
 var app = builder.Build();
 
diff --git a/RoomManager.cs b/RoomManager.cs
new file mode 100644
index 0000000..192bb86
--- /dev/null
+++ b/RoomManager.cs
@@ -0,0 +1,6 @@
+namespace MauMau_Server;
+
+public class RoomManager
+{
+    
+}
\ No newline at end of file
diff --git a/Websockets/WebsocketManager.cs b/Websockets/WebsocketManager.cs
new file mode 100644
index 0000000..bf909c6
--- /dev/null
+++ b/Websockets/WebsocketManager.cs
@@ -0,0 +1,38 @@
+using System.Net.WebSockets;
+
+namespace MauMau_Server.Websockets;
+
+public class WebsocketManager : IWebsocketManager
+{
+    private static readonly Dictionary<string, WebSocket> Connections = new();
+
+    public string AddConnection(WebSocket socket)
+    {
+        var id = Guid.NewGuid().ToString();
+        Connections.Add(id, socket);
+        return id;
+    }
+    
+    public WebSocket GetConnection(string id)
+    {
+        return Connections[id];
+    }
+    
+    public Dictionary<string, WebSocket> GetAllConnections()
+    {
+        return Connections;
+    }
+    
+    public void RemoveConnection(string id)
+    {
+        Connections.Remove(id);
+    }
+}
+
+public interface IWebsocketManager
+{
+    public string AddConnection(WebSocket socket);
+    public WebSocket GetConnection(string id);
+    public Dictionary<string, WebSocket> GetAllConnections();
+    public void RemoveConnection(string id);
+}
\ No newline at end of file