Introduced typed resources for client

MTSC-TestServer/Program.cs

@@ -27,7 +27,7 @@ namespace MTSC_TestServer
                 .AddServerUsageMonitor(new TickrateEnforcer().SetTicksPerSecond(60))
                 .AddExceptionHandler(new ExceptionConsoleLogger())
                 //.AddHandler(new BroadcastHandler())
-                .AddHandler(new WebsocketHandler().AddWebsocketHandler(new BroadcastModule()))
+                .AddHandler(new WebsocketHandler())
                 .AddHandler(new HttpHandler().AddHttpModule(new FileServerModule())
                                             .AddHttpModule(new PostModule()))
                 .Run();

MTSC/Common/WebSockets/ServerModules/BroadcastModule.cs

@@ -1,29 +0,0 @@
-using MTSC.ServerSide;
-using MTSC.ServerSide.Handlers;
-
-namespace MTSC.Common.WebSockets.ServerModules
-{
-    public sealed class BroadcastModule : IWebsocketModule
-    {
-        void IWebsocketModule.ConnectionClosed(ServerSide.Server server, WebsocketHandler handler, ClientData client)
-        {
-            
-        }
-
-        void IWebsocketModule.ConnectionInitialized(ServerSide.Server server, WebsocketHandler handler, ClientData client)
-        {
-            
-        }
-
-        bool IWebsocketModule.HandleReceivedMessage(ServerSide.Server server, WebsocketHandler handler, ClientData client, WebsocketMessage receivedMessage)
-        {
-            receivedMessage.Masked = false;
-            //receivedMessage.Opcode = WebsocketMessage.Opcodes.Text;
-            foreach(ClientData otherClient in handler.webSockets.Keys)
-            {
-                handler.QueueMessage(otherClient, receivedMessage);
-            }
-            return false;
-        }
-    }
-}

MTSC/ServerSide/ClientData.cs

@@ -0,0 +1,60 @@
+using System;
+using System.Net.Security;
+using System.Net.Sockets;
+
+namespace MTSC.ServerSide
+{
+    /// <summary>
+    /// Structure containing client information.
+    /// </summary>
+    public class ClientData : IDisposable
+    {
+        public TcpClient TcpClient;
+        public DateTime LastMessageTime = DateTime.Now;
+        public bool ToBeRemoved = false;
+        public SslStream SslStream = null;
+        public ResourceDictionary Resources = new ResourceDictionary();
+
+        public ClientData(TcpClient client)
+        {
+            this.TcpClient = client;
+        }
+
+        #region IDisposable Support
+        private bool disposedValue = false; // To detect redundant calls
+
+        protected virtual void Dispose(bool disposing)
+        {
+            if (!disposedValue)
+            {
+                if (disposing)
+                {
+                    // TODO: dispose managed state (managed objects).
+                }
+
+                TcpClient?.Dispose();
+                SslStream?.Dispose();
+                Resources?.Dispose();
+
+                disposedValue = true;
+            }
+        }
+
+        // TODO: override a finalizer only if Dispose(bool disposing) above has code to free unmanaged resources.
+        // ~ClientData()
+        // {
+        //   // Do not change this code. Put cleanup code in Dispose(bool disposing) above.
+        //   Dispose(false);
+        // }
+
+        // This code added to correctly implement the disposable pattern.
+        public void Dispose()
+        {
+            // Do not change this code. Put cleanup code in Dispose(bool disposing) above.
+            Dispose(true);
+            // TODO: uncomment the following line if the finalizer is overridden above.
+            // GC.SuppressFinalize(this);
+        }
+        #endregion
+    }
+}

MTSC/ServerSide/Handlers/EncryptionHandler.cs

@@ -23,7 +23,6 @@ namespace MTSC.ServerSide.Handlers
             public ClientState ClientState = ClientState.Initial;
         }
         #region Fields
-        private Dictionary<ClientData, AdditionalData> additionalData;
         private RSACryptoServiceProvider rsa;
         private string privateKey, publicKey;
         #endregion
@@ -37,7 +36,6 @@ namespace MTSC.ServerSide.Handlers
         /// <param name="server">Server object managed by the handler.</param>
         public EncryptionHandler(RSACryptoServiceProvider rsa)
         {
-            additionalData = new Dictionary<ClientData, AdditionalData>();
             this.rsa = rsa;
             privateKey = HelperFunctions.ToXmlString(rsa, true);
             publicKey = HelperFunctions.ToXmlString(rsa, false);
@@ -124,7 +122,7 @@ namespace MTSC.ServerSide.Handlers
         /// <param name="client">Client to be removed.</param>
         void IHandler.ClientRemoved(Server server, ClientData client)
         {
-            additionalData.Remove(client);
+
         }
         /// <summary>
         /// Handle a new client connection.
@@ -133,7 +131,7 @@ namespace MTSC.ServerSide.Handlers
         /// <returns>False if an error occurred.</returns>
         bool IHandler.HandleClient(Server server, ClientData client)
         {
-            additionalData.Add(client, new AdditionalData());
+            client.Resources.SetResource(new AdditionalData());
             return false;
         }
         /// <summary>
@@ -144,13 +142,14 @@ namespace MTSC.ServerSide.Handlers
         /// <returns>True if no other handler should handle current message.</returns>
         bool IHandler.HandleReceivedMessage(Server server, ClientData client, Message message)
         {
-            if (additionalData[client].ClientState == ClientState.Initial || additionalData[client].ClientState == ClientState.Negotiating)
+            var additionalData = client.Resources.GetResource<AdditionalData>();
+            if (additionalData.ClientState == ClientState.Initial || additionalData.ClientState == ClientState.Negotiating)
             {
                 string asciiMessage = ASCIIEncoding.ASCII.GetString(message.MessageBytes);
                 if (asciiMessage == CommunicationPrimitives.RequestPublicKey)
                 {
                     server.QueueMessage(client, ASCIIEncoding.ASCII.GetBytes(CommunicationPrimitives.SendPublicKey + ":" + publicKey));
-                    additionalData[client].ClientState = ClientState.Negotiating;
+                    additionalData.ClientState = ClientState.Negotiating;
                     return true;
                 }
                 else if (asciiMessage.Contains(CommunicationPrimitives.SendEncryptionKey))
@@ -158,9 +157,9 @@ namespace MTSC.ServerSide.Handlers
                     byte[] encryptedKey = new byte[message.MessageLength - CommunicationPrimitives.SendEncryptionKey.Length - 1];
                     Array.Copy(message.MessageBytes, CommunicationPrimitives.SendEncryptionKey.Length + 1, encryptedKey, 0, encryptedKey.Length);
                     byte[] decryptedKey = rsa.Decrypt(encryptedKey, false);
-                    additionalData[client].Key = decryptedKey;
+                    additionalData.Key = decryptedKey;
                     server.QueueMessage(client, ASCIIEncoding.ASCII.GetBytes(CommunicationPrimitives.AcceptEncryptionKey));
-                    additionalData[client].ClientState = ClientState.Encrypted;
+                    additionalData.ClientState = ClientState.Encrypted;
                     return true;
                 }
                 else
@@ -181,13 +180,14 @@ namespace MTSC.ServerSide.Handlers
         /// <returns>True if no other handlers should perform operations on current message.</returns>
         bool IHandler.PreHandleReceivedMessage(Server server, ClientData client, ref Message message)
         {
-            if (additionalData[client].ClientState == ClientState.Encrypted)
+            var additionalData = client.Resources.GetResource<AdditionalData>();
+            if (additionalData.ClientState == ClientState.Encrypted)
             {
                 /*
                  * Decrypt message before returning.
                  */
                 byte[] encryptedBytes = message.MessageBytes;
-                byte[] decryptedBytes = DecryptBytes(additionalData[client].Key, encryptedBytes);
+                byte[] decryptedBytes = DecryptBytes(additionalData.Key, encryptedBytes);
                 message = new Message((uint)decryptedBytes.Length, decryptedBytes);
                 return false;
             }
@@ -214,9 +214,10 @@ namespace MTSC.ServerSide.Handlers
         /// <returns>True if no other handler should process this message.</returns>
         bool IHandler.HandleSendMessage(Server server, ClientData client, ref Message message)
         {
-            if (additionalData[client].ClientState == ClientState.Encrypted)
+            var additionalData = client.Resources.GetResource<AdditionalData>();
+            if (additionalData.ClientState == ClientState.Encrypted)
             {
-                message = CommunicationPrimitives.BuildMessage(EncryptBytes(additionalData[client].Key, message.MessageBytes));
+                message = CommunicationPrimitives.BuildMessage(EncryptBytes(additionalData.Key, message.MessageBytes));
             }
             return false;
         }

MTSC/ServerSide/Handlers/WebsocketHandler.cs

@@ -26,7 +26,6 @@ namespace MTSC.ServerSide.Handlers
             Closed
         }
         #region Fields
-        public ConcurrentDictionary<ClientData, SocketState> webSockets = new ConcurrentDictionary<ClientData, SocketState>();
         ConcurrentQueue<Tuple<ClientData,WebsocketMessage>> messageQueue = new ConcurrentQueue<Tuple<ClientData, WebsocketMessage>>();
         List<IWebsocketModule> websocketModules = new List<IWebsocketModule>();
         #endregion
@@ -80,22 +79,19 @@ namespace MTSC.ServerSide.Handlers
         #region Handler Implementation
         void IHandler.ClientRemoved(Server server, ClientData client)
         {
-            SocketState state = SocketState.Initial;
-            while (webSockets.ContainsKey(client))
-            {
-                webSockets.TryRemove(client, out state);
-            }
+            
         }
 
         bool IHandler.HandleClient(Server server, ClientData client)
         {
-            webSockets[client] = SocketState.Initial;
+            client.Resources.SetResource(SocketState.Initial);
             return false;
         }
 
         bool IHandler.HandleReceivedMessage(Server server, ClientData client, Message message)
         {
-            if (webSockets[client] == SocketState.Initial)
+            var socketState = client.Resources.GetResource<SocketState>();
+            if (socketState == SocketState.Initial)
             {
                 PartialHttpRequest request = new PartialHttpRequest(message.MessageBytes);
                 if(request.Method == HttpMessage.HttpMethods.Get && request.Headers.ContainsHeader(HttpMessage.GeneralHeaders.Connection) &&
@@ -119,7 +115,7 @@ namespace MTSC.ServerSide.Handlers
                     response.Headers[HttpMessage.GeneralHeaders.Connection] = "Upgrade";
                     response.Headers[WebsocketHeaderAcceptKey] = returnBase64Key;
                     server.QueueMessage(client, response.GetPackedResponse(false));
-                    webSockets[client] = SocketState.Established;
+                    client.Resources.SetResource(SocketState.Established);
                     server.LogDebug("Websocket initialized " + client.TcpClient.Client.RemoteEndPoint.ToString());
                     foreach (IWebsocketModule websocketModule in websocketModules)
                     {
@@ -128,7 +124,7 @@ namespace MTSC.ServerSide.Handlers
                     return true;
                 }
             }
-            else if(webSockets[client] == SocketState.Established)
+            else if(socketState == SocketState.Established)
             {
                 WebsocketMessage receivedMessage = new WebsocketMessage(message.MessageBytes);
                 if (receivedMessage.Opcode == WebsocketMessage.Opcodes.Close)
@@ -138,10 +134,6 @@ namespace MTSC.ServerSide.Handlers
                         websocketModule.ConnectionClosed(server, this, client);
                     }
                     client.ToBeRemoved = true;
-                    while (webSockets.ContainsKey(client))
-                    {
-                        webSockets.TryRemove(client, out SocketState _);
-                    }
                 }
                 else
                 {
@@ -182,10 +174,6 @@ namespace MTSC.ServerSide.Handlers
                             websocketModule.ConnectionClosed(server, this, tuple.Item1);
                         }
                         tuple.Item1.ToBeRemoved = true;
-                        while (webSockets.ContainsKey(tuple.Item1))
-                        {
-                            webSockets.TryRemove(tuple.Item1, out _);
-                        }
                     }
                 }
             }

MTSC/ServerSide/Handlers/WebsocketRoutingHandler.cs

@@ -27,10 +27,8 @@ namespace MTSC.ServerSide.Handlers
             Closed
         }
         #region Fields
-        public ConcurrentDictionary<ClientData, SocketState> webSockets = new ConcurrentDictionary<ClientData, SocketState>();
         Dictionary<string, (WebsocketRouteBase, Func<Server, HttpRequest, ClientData, RouteEnablerResponse>)> moduleDictionary =
             new Dictionary<string, (WebsocketRouteBase, Func<Server, HttpRequest, ClientData, RouteEnablerResponse>)>();
-        ConcurrentDictionary<ClientData, WebsocketRouteBase> routingTable = new ConcurrentDictionary<ClientData, WebsocketRouteBase>();
         ConcurrentQueue<Tuple<ClientData, WebsocketMessage>> messageQueue = new ConcurrentQueue<Tuple<ClientData, WebsocketMessage>>();
         #endregion
         #region Public Methods
@@ -83,29 +81,22 @@ namespace MTSC.ServerSide.Handlers
         #region Handler Implementation
         void IHandler.ClientRemoved(Server server, ClientData client)
         {
-            while (webSockets.ContainsKey(client))
-            {
-                webSockets.TryRemove(client, out SocketState state);
-            }
-            if (routingTable.TryGetValue(client, out WebsocketRouteBase route))
+            if (client.Resources.TryGetResource(out WebsocketRouteBase route))
             {
                 route.CallConnectionClosed(server, this, client);
             }
-            while (routingTable.ContainsKey(client))
-            {
-                routingTable.TryRemove(client, out _);
-            }
         }
 
         bool IHandler.HandleClient(Server server, ClientData client)
         {
-            webSockets[client] = SocketState.Initial;
+            client.Resources.SetResource(SocketState.Initial);
             return false;
         }
 
         bool IHandler.HandleReceivedMessage(Server server, ClientData client, Message message)
         {
-            if (webSockets[client] == SocketState.Initial)
+            var socketState = client.Resources.GetResource<SocketState>();
+            if (socketState == SocketState.Initial)
             {
                 PartialHttpRequest request;
                 try
@@ -154,9 +145,9 @@ namespace MTSC.ServerSide.Handlers
                     response.Headers[HttpMessage.GeneralHeaders.Connection] = "Upgrade";
                     response.Headers[WebsocketHeaderAcceptKey] = returnBase64Key;
                     server.QueueMessage(client, response.GetPackedResponse(false));
-                    webSockets[client] = SocketState.Established;
+                    client.Resources.SetResource(SocketState.Established);
                     server.LogDebug("Websocket initialized " + client.TcpClient.Client.RemoteEndPoint.ToString());
-                    routingTable[client] = module;
+                    client.Resources.SetResource(module);
                     module.CallConnectionInitialized(server, this, client);
                     return true;
                 }
@@ -165,7 +156,7 @@ namespace MTSC.ServerSide.Handlers
                     return false;
                 }
             }
-            else if (webSockets[client] == SocketState.Established)
+            else if (socketState == SocketState.Established)
             {
                 WebsocketMessage receivedMessage = null;
                 try
@@ -181,10 +172,6 @@ namespace MTSC.ServerSide.Handlers
                 if (receivedMessage.Opcode == WebsocketMessage.Opcodes.Close)
                 {
                     client.ToBeRemoved = true;
-                    while (webSockets.ContainsKey(client))
-                    {
-                        webSockets.TryRemove(client, out SocketState _);
-                    }
                     WebsocketMessage closeFrame = new WebsocketMessage();
                     closeFrame.Opcode = WebsocketMessage.Opcodes.Close;
                     QueueMessage(client, closeFrame);
@@ -194,7 +181,7 @@ namespace MTSC.ServerSide.Handlers
                 {
                     try 
                     {
-                        routingTable[client].CallHandleReceivedMessage(server, this, client, receivedMessage);
+                        client.Resources.GetResource<WebsocketRouteBase>().CallHandleReceivedMessage(server, this, client, receivedMessage);
                         return true;
                     }
                     catch(Exception e)
@@ -233,9 +220,9 @@ namespace MTSC.ServerSide.Handlers
                     server.QueueMessage(tuple.Item1, tuple.Item2.GetMessageBytes());
                     if (tuple.Item2.Opcode == WebsocketMessage.Opcodes.Close)
                     {
-                        if (routingTable.ContainsKey(tuple.Item1))
+                        if (tuple.Item1.Resources.TryGetResource<WebsocketRouteBase>(out var route))
                         {
-                            routingTable[tuple.Item1].CallConnectionClosed(server, this, tuple.Item1);
+                            route.CallConnectionClosed(server, this, tuple.Item1);
                         }
                         tuple.Item1.ToBeRemoved = true;
                     }

MTSC/ServerSide/ResourceDictionary.cs

@@ -0,0 +1,78 @@
+using System;
+using System.Collections.Generic;
+
+namespace MTSC.ServerSide
+{
+    public class ResourceDictionary : IDisposable
+    {
+        private Dictionary<Type, object> Resources = new Dictionary<Type, object>();
+
+        public void SetResource<TValue>(TValue value)
+        {
+            Resources[typeof(TValue)] = value;
+        }
+
+        public TValue GetResource<TValue>()
+        {
+            return (TValue)Resources[typeof(TValue)];
+        }
+
+        public bool Contains<TValue>()
+        {
+            return Resources.ContainsKey(typeof(TValue));
+        }
+
+        public bool TryGetResource<TValue>(out TValue value)
+        {
+            if (Resources.TryGetValue(typeof(TValue), out object objValue))
+            {
+                value = (TValue)objValue;
+                return true;
+            }
+            else
+            {
+                value = default(TValue);
+                return false;
+            }
+        }
+
+        #region IDisposable Support
+        private bool disposedValue = false; // To detect redundant calls
+
+        protected virtual void Dispose(bool disposing)
+        {
+            if (!disposedValue)
+            {
+                if (disposing)
+                {
+                    // TODO: dispose managed state (managed objects).
+                }
+
+                foreach(var value in Resources.Values)
+                {
+                    (value as IDisposable)?.Dispose();
+                }
+                Resources.Clear();
+                Resources = null;
+                disposedValue = true;
+            }
+        }
+
+        // TODO: override a finalizer only if Dispose(bool disposing) above has code to free unmanaged resources.
+        // ~ResourceDictionary()
+        // {
+        //   // Do not change this code. Put cleanup code in Dispose(bool disposing) above.
+        //   Dispose(false);
+        // }
+
+        // This code added to correctly implement the disposable pattern.
+        public void Dispose()
+        {
+            // Do not change this code. Put cleanup code in Dispose(bool disposing) above.
+            Dispose(true);
+            // TODO: uncomment the following line if the finalizer is overridden above.
+            // GC.SuppressFinalize(this);
+        }
+        #endregion
+    }
+}

MTSC/ServerSide/Server.cs

@@ -713,22 +713,4 @@ namespace MTSC.ServerSide
         }
         #endregion
     }
-    /// <summary>
-    /// Structure containing client information.
-    /// </summary>
-    public class ClientData
-    {
-        public TcpClient TcpClient;
-        public DateTime LastMessageTime;
-        public bool ToBeRemoved;
-        public SslStream SslStream;
-
-        public ClientData(TcpClient client)
-        {
-            this.TcpClient = client;
-            this.LastMessageTime = DateTime.Now;
-            ToBeRemoved = false;
-            SslStream = null;
-        }
-    }
 }