diff --git a/MTSC-TestServer/Program.cs b/MTSC-TestServer/Program.cs index 28774b9..a6ccb5c 100644 --- a/MTSC-TestServer/Program.cs +++ b/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(); diff --git a/MTSC/Common/WebSockets/ServerModules/BroadcastModule.cs b/MTSC/Common/WebSockets/ServerModules/BroadcastModule.cs deleted file mode 100644 index 2c0f18b..0000000 --- a/MTSC/Common/WebSockets/ServerModules/BroadcastModule.cs +++ /dev/null @@ -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; - } - } -} diff --git a/MTSC/ServerSide/ClientData.cs b/MTSC/ServerSide/ClientData.cs new file mode 100644 index 0000000..41fae25 --- /dev/null +++ b/MTSC/ServerSide/ClientData.cs @@ -0,0 +1,60 @@ +using System; +using System.Net.Security; +using System.Net.Sockets; + +namespace MTSC.ServerSide +{ + /// + /// Structure containing client information. + /// + 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 + } +} diff --git a/MTSC/ServerSide/Handlers/EncryptionHandler.cs b/MTSC/ServerSide/Handlers/EncryptionHandler.cs index a3822d4..64c0991 100644 --- a/MTSC/ServerSide/Handlers/EncryptionHandler.cs +++ b/MTSC/ServerSide/Handlers/EncryptionHandler.cs @@ -23,7 +23,6 @@ namespace MTSC.ServerSide.Handlers public ClientState ClientState = ClientState.Initial; } #region Fields - private Dictionary additionalData; private RSACryptoServiceProvider rsa; private string privateKey, publicKey; #endregion @@ -37,7 +36,6 @@ namespace MTSC.ServerSide.Handlers /// Server object managed by the handler. public EncryptionHandler(RSACryptoServiceProvider rsa) { - additionalData = new Dictionary(); this.rsa = rsa; privateKey = HelperFunctions.ToXmlString(rsa, true); publicKey = HelperFunctions.ToXmlString(rsa, false); @@ -124,7 +122,7 @@ namespace MTSC.ServerSide.Handlers /// Client to be removed. void IHandler.ClientRemoved(Server server, ClientData client) { - additionalData.Remove(client); + } /// /// Handle a new client connection. @@ -133,7 +131,7 @@ namespace MTSC.ServerSide.Handlers /// False if an error occurred. bool IHandler.HandleClient(Server server, ClientData client) { - additionalData.Add(client, new AdditionalData()); + client.Resources.SetResource(new AdditionalData()); return false; } /// @@ -144,13 +142,14 @@ namespace MTSC.ServerSide.Handlers /// True if no other handler should handle current message. 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(); + 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 /// True if no other handlers should perform operations on current message. bool IHandler.PreHandleReceivedMessage(Server server, ClientData client, ref Message message) { - if (additionalData[client].ClientState == ClientState.Encrypted) + var additionalData = client.Resources.GetResource(); + 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 /// True if no other handler should process this message. bool IHandler.HandleSendMessage(Server server, ClientData client, ref Message message) { - if (additionalData[client].ClientState == ClientState.Encrypted) + var additionalData = client.Resources.GetResource(); + if (additionalData.ClientState == ClientState.Encrypted) { - message = CommunicationPrimitives.BuildMessage(EncryptBytes(additionalData[client].Key, message.MessageBytes)); + message = CommunicationPrimitives.BuildMessage(EncryptBytes(additionalData.Key, message.MessageBytes)); } return false; } diff --git a/MTSC/ServerSide/Handlers/WebsocketHandler.cs b/MTSC/ServerSide/Handlers/WebsocketHandler.cs index 13d1bac..9d66c19 100644 --- a/MTSC/ServerSide/Handlers/WebsocketHandler.cs +++ b/MTSC/ServerSide/Handlers/WebsocketHandler.cs @@ -26,7 +26,6 @@ namespace MTSC.ServerSide.Handlers Closed } #region Fields - public ConcurrentDictionary webSockets = new ConcurrentDictionary(); ConcurrentQueue> messageQueue = new ConcurrentQueue>(); List websocketModules = new List(); #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(); + 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 _); - } } } } diff --git a/MTSC/ServerSide/Handlers/WebsocketRoutingHandler.cs b/MTSC/ServerSide/Handlers/WebsocketRoutingHandler.cs index 6275170..6c9f88f 100644 --- a/MTSC/ServerSide/Handlers/WebsocketRoutingHandler.cs +++ b/MTSC/ServerSide/Handlers/WebsocketRoutingHandler.cs @@ -27,10 +27,8 @@ namespace MTSC.ServerSide.Handlers Closed } #region Fields - public ConcurrentDictionary webSockets = new ConcurrentDictionary(); Dictionary)> moduleDictionary = new Dictionary)>(); - ConcurrentDictionary routingTable = new ConcurrentDictionary(); ConcurrentQueue> messageQueue = new ConcurrentQueue>(); #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(); + 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().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(out var route)) { - routingTable[tuple.Item1].CallConnectionClosed(server, this, tuple.Item1); + route.CallConnectionClosed(server, this, tuple.Item1); } tuple.Item1.ToBeRemoved = true; } diff --git a/MTSC/ServerSide/ResourceDictionary.cs b/MTSC/ServerSide/ResourceDictionary.cs new file mode 100644 index 0000000..7b42290 --- /dev/null +++ b/MTSC/ServerSide/ResourceDictionary.cs @@ -0,0 +1,78 @@ +using System; +using System.Collections.Generic; + +namespace MTSC.ServerSide +{ + public class ResourceDictionary : IDisposable + { + private Dictionary Resources = new Dictionary(); + + public void SetResource(TValue value) + { + Resources[typeof(TValue)] = value; + } + + public TValue GetResource() + { + return (TValue)Resources[typeof(TValue)]; + } + + public bool Contains() + { + return Resources.ContainsKey(typeof(TValue)); + } + + public bool TryGetResource(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 + } +} diff --git a/MTSC/ServerSide/Server.cs b/MTSC/ServerSide/Server.cs index bac6113..47234b4 100644 --- a/MTSC/ServerSide/Server.cs +++ b/MTSC/ServerSide/Server.cs @@ -713,22 +713,4 @@ namespace MTSC.ServerSide } #endregion } - /// - /// Structure containing client information. - /// - 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; - } - } }