Fixed encryption issues with non seeking CryptoStream

Author: TwoTenPvP

MLAPI/Data/NetworkConfig.cs

@@ -68,6 +68,10 @@ public class NetworkConfig
         /// </summary>
         public int MessageBufferSize = 1024;
         /// <summary>
+        /// The size of the encryption buffer, this is the buffer where messages will be decrypted to.
+        /// </summary>
+        public int EncryptionBufferSize = 265;
+        /// <summary>
         /// Amount of times per second the receive queue is emptied and all messages inside are processed.
         /// </summary>
         public int ReceiveTickrate = 64;
@@ -205,6 +209,7 @@ public string ToBase64()
                     }
 
                     writer.WriteInt32Packed(config.MessageBufferSize);
+                    writer.WriteInt32Packed(config.EncryptionBufferSize);
                     writer.WriteInt32Packed(config.ReceiveTickrate);
                     writer.WriteInt32Packed(config.MaxReceiveEventsPerTickRate);
                     writer.WriteInt32Packed(config.SendTickrate);
@@ -285,6 +290,7 @@ public void FromBase64(string base64, bool createDummyObject = false)
                     }
 
                     config.MessageBufferSize = reader.ReadInt32Packed();
+                    config.EncryptionBufferSize = reader.ReadInt32Packed();
                     config.ReceiveTickrate = reader.ReadInt32Packed();
                     config.MaxReceiveEventsPerTickRate = reader.ReadInt32Packed();
                     config.SendTickrate = reader.ReadInt32Packed();

MLAPI/MonoBehaviours/Core/NetworkingManager.cs

@@ -100,6 +100,7 @@ internal set
         /// </summary>
         public bool isListening { get; internal set; }
         private byte[] messageBuffer;
+        private byte[] encryptionBuffer;
         /// <summary>
         /// Gets if we are connected as a client
         /// </summary>
@@ -282,9 +283,12 @@ private object Init(bool server)
             lastReceiveTickTime = 0f;
             eventOvershootCounter = 0f;
             connectionPendingClients.Clear();
+            hailPendingClients.Clear();
+            pendingClientAesKeys.Clear();
             ConnectedClients.Clear();
             ConnectedClientsList.Clear();
             messageBuffer = new byte[NetworkConfig.MessageBufferSize];
+            encryptionBuffer = new byte[NetworkConfig.EncryptionBufferSize];
 #if !DISABLE_CRYPTOGRAPHY
             pendingKeyExchanges.Clear();
 #endif
@@ -523,6 +527,19 @@ public void StopServer()
                     NetworkConfig.NetworkTransport.DisconnectClient(clientId);
                 }
             }
+
+            foreach (uint clientId in hailPendingClients)
+            {
+                if (!disconnectedIds.Contains(clientId))
+                {
+                    disconnectedIds.Add(clientId);
+                    if (clientId == NetworkConfig.NetworkTransport.ServerClientId)
+                        continue;
+
+                    NetworkConfig.NetworkTransport.DisconnectClient(clientId);
+                }
+            }
+            
             isServer = false;
             Shutdown();
         }
@@ -812,13 +829,21 @@ private IEnumerator ApprovalTimeout(uint clientId)
         {
             float timeStarted = NetworkTime;
             //We yield every frame incase a pending client disconnects and someone else gets its connection id
-            while (NetworkTime - timeStarted < NetworkConfig.ClientConnectionBufferTimeout && connectionPendingClients.Contains(clientId))
+            while (NetworkTime - timeStarted < NetworkConfig.ClientConnectionBufferTimeout && (connectionPendingClients.Contains(clientId) || hailPendingClients.Contains(clientId)))
             {
                 yield return null;
             }
-            if(connectionPendingClients.Contains(clientId) && !ConnectedClients.ContainsKey(clientId))
+            
+            if (connectionPendingClients.Contains(clientId) && !ConnectedClients.ContainsKey(clientId))
+            {
+                // Timeout
+                if (LogHelper.CurrentLogLevel <= LogLevel.Developer) LogHelper.LogInfo("Client " + clientId + " Handshake Timed Out");
+                DisconnectClient(clientId);
+            }
+
+            if (hailPendingClients.Contains(clientId) && !ConnectedClients.ContainsKey(clientId))
             {
-                //Timeout
+                // Timeout
                 if (LogHelper.CurrentLogLevel <= LogLevel.Developer) LogHelper.LogInfo("Client " + clientId + " Handshake Timed Out");
                 DisconnectClient(clientId);
             }
@@ -830,8 +855,7 @@ private void HandleIncomingData(uint clientId, byte[] data, int channelId, int t
         {
             using (BitStream bitStream = new BitStream(data))
             {
-                RijndaelManaged rijndael = null;
-                Stream stream = bitStream;
+                BitStream stream = bitStream;
                 try
                 {
                     if (LogHelper.CurrentLogLevel <= LogLevel.Developer) LogHelper.LogInfo("Unwrapping Data Header");
@@ -845,11 +869,20 @@ private void HandleIncomingData(uint clientId, byte[] data, int channelId, int t
                         {
                             headerReader.SkipPadBits();
                             headerReader.ReadByteArray(IVBuffer, 16);
-                            rijndael = new RijndaelManaged();
-                            rijndael.Padding = PaddingMode.PKCS7;
-                            rijndael.Key = isServer ? (ConnectedClients.ContainsKey(clientId) ? ConnectedClients[clientId].AesKey : pendingClientAesKeys[clientId]) : clientAesKey;
-                            rijndael.IV = IVBuffer;
-                            stream = new CryptoStream(bitStream, rijndael.CreateDecryptor(), CryptoStreamMode.Read);
+                            stream = new BitStream(encryptionBuffer);
+                            using (RijndaelManaged rijndael = new RijndaelManaged())
+                            {
+                                rijndael.Padding = PaddingMode.PKCS7;
+                                rijndael.Key = isServer ? (ConnectedClients.ContainsKey(clientId) ? ConnectedClients[clientId].AesKey : pendingClientAesKeys[clientId]) : clientAesKey;
+                                rijndael.IV = IVBuffer;
+                                using (CryptoStream cryptoStream = new CryptoStream(bitStream, rijndael.CreateDecryptor(), CryptoStreamMode.Read))
+                                {
+                                    int readByte = 0;
+                                    while ((readByte = cryptoStream.ReadByte()) != -1)
+                                        stream.WriteByte((byte)readByte);
+                                }
+                            }
+
                             using (PooledBitReader reader = PooledBitReader.Get(stream))
                             {
                                 messageType = reader.ReadByteDirect();
@@ -860,8 +893,8 @@ private void HandleIncomingData(uint clientId, byte[] data, int channelId, int t
                             headerReader.SkipPadBits();
                             using (HMACSHA256 hmac = new HMACSHA256(isServer ? ConnectedClients[clientId].AesKey : clientAesKey))
                             {
-                                // 1 is the size of the header. 32 is the size of the hmac
                                 headerReader.ReadByteArray(HMACBuffer, 32);
+                                // 1 is the size of the header. 32 is the size of the hmac
                                 byte[] hmacBytes = hmac.ComputeHash(bitStream.GetBuffer(), 1 + 32, totalSize - (1 + 32));
                                 for (int i = 0; i < hmacBytes.Length; i++)
                                 {
@@ -872,6 +905,7 @@ private void HandleIncomingData(uint clientId, byte[] data, int channelId, int t
                                     }
                                 }
                             }
+
                             messageType = headerReader.ReadByteDirect();
                         }
                         else
@@ -886,7 +920,7 @@ private void HandleIncomingData(uint clientId, byte[] data, int channelId, int t
 
                         //Client tried to send a network message that was not the connection request before he was accepted.
                         if (isServer && (NetworkConfig.EnableEncryption && hailPendingClients.Contains(clientId) && messageType != MLAPIConstants.MLAPI_CERTIFICATE_HAIL_RESPONSE) ||
-                                        (connectionPendingClients.Contains(clientId) && messageType != MLAPIConstants.MLAPI_CONNECTION_REQUEST))
+                            (connectionPendingClients.Contains(clientId) && messageType != MLAPIConstants.MLAPI_CONNECTION_REQUEST))
                         {
                             if (LogHelper.CurrentLogLevel <= LogLevel.Normal) LogHelper.LogWarning("Message recieved from clientId " + clientId + " before it has been accepted");
                             return;
@@ -967,7 +1001,6 @@ private void HandleIncomingData(uint clientId, byte[] data, int channelId, int t
                 finally
                 {
                     if (stream != bitStream) stream.Dispose();
-                    if (rijndael != null) rijndael.Clear();
                 }
             }
         }
@@ -983,9 +1016,15 @@ internal void DisconnectClient(uint clientId)
             if (connectionPendingClients.Contains(clientId))
                 connectionPendingClients.Remove(clientId);
 
+            if (hailPendingClients.Contains(clientId))
+                hailPendingClients.Remove(clientId);
+
             if (ConnectedClients.ContainsKey(clientId))
                 ConnectedClients.Remove(clientId);
 
+            if (pendingClientAesKeys.ContainsKey(clientId))
+                pendingClientAesKeys.Remove(clientId);
+
             for (int i = ConnectedClientsList.Count - 1; i > -1; i--)
             {
                 if (ConnectedClientsList[i].ClientId == clientId)
@@ -1004,6 +1043,13 @@ internal void OnClientDisconnectFromServer(uint clientId)
         {
             if (connectionPendingClients.Contains(clientId))
                 connectionPendingClients.Remove(clientId);
+
+            if (hailPendingClients.Contains(clientId))
+                hailPendingClients.Remove(clientId);
+
+            if (pendingClientAesKeys.ContainsKey(clientId))
+                pendingClientAesKeys.Remove(clientId);
+            
             if (ConnectedClients.ContainsKey(clientId))
             {
                 if(NetworkConfig.HandleObjectSpawning)
@@ -1060,12 +1106,16 @@ internal void HandleApproval(uint clientId, int prefabId, bool approved, Vector3
         {
             if(approved)
             {
-                //Inform new client it got approved
+                // Inform new client it got approved
                 if (connectionPendingClients.Contains(clientId))
                     connectionPendingClients.Remove(clientId);
 
+                if (hailPendingClients.Contains(clientId))
+                    hailPendingClients.Remove(clientId);
+                
                 byte[] aesKey = pendingClientAesKeys.ContainsKey(clientId) ? pendingClientAesKeys[clientId] : null;
-                pendingClientAesKeys.Remove(clientId);
+                if (pendingClientAesKeys.ContainsKey(clientId))
+                    pendingClientAesKeys.Remove(clientId);
                 NetworkedClient client = new NetworkedClient()
                 {
                     ClientId = clientId,
@@ -1183,6 +1233,9 @@ internal void HandleApproval(uint clientId, int prefabId, bool approved, Vector3
                 if (connectionPendingClients.Contains(clientId))
                     connectionPendingClients.Remove(clientId);
 
+                if (hailPendingClients.Contains(clientId))
+                    hailPendingClients.Remove(clientId);
+
 #if !DISABLE_CRYPTOGRAPHY
                 if (pendingKeyExchanges.ContainsKey(clientId))
                     pendingKeyExchanges.Remove(clientId);

MLAPI/NetworkingManagerComponents/Cryptography/CryptographyHelper.cs

@@ -14,11 +14,11 @@ public static class CryptographyHelper
     {
         public delegate bool VerifyCertificateDelegate(X509Certificate2 certificate, string hostname);
         public static VerifyCertificateDelegate OnValidateCertificateCallback = null;
-
+        
         public static bool VerifyCertificate(X509Certificate2 certificate, string hostname)
         {
             if (OnValidateCertificateCallback != null) return OnValidateCertificateCallback(certificate, hostname);
-            return certificate.Verify() && hostname == certificate.GetNameInfo(X509NameType.DnsName, false);
+            return certificate.Verify() && (hostname == certificate.GetNameInfo(X509NameType.DnsName, false) || hostname == "127.0.0.1");
         }
     }
 }