Skip to content

Commit c250595

Browse files
TwoTenPvPTwoTenPvP
committed
Added PendingClient structure to unify how pending clients are handled
1 parent 17de52b commit c250595

File tree

4 files changed

+67
-93
lines changed

4 files changed

+67
-93
lines changed

‎MLAPI/Data/PendingClient.cs‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
using MLAPI.Cryptography;
2+
3+
namespace MLAPI.Data
4+
{
5+
public class PendingClient
6+
{
7+
public uint ClientId;
8+
9+
internal EllipticDiffieHellman KeyExchange;
10+
11+
public byte[] AesKey;
12+
13+
public State ConnectionState;
14+
15+
public enum State
16+
{
17+
PendingHail,
18+
PendingConnection
19+
}
20+
}
21+
}

‎MLAPI/MLAPI.csproj‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@
9090
<Compile Include="Data\NetworkedCollections\NetworkedList.cs" />
9191
<Compile Include="Data\NetworkedVar.cs" />
9292
<Compile Include="Data\NetworkedVarMeta.cs" />
93+
<Compile Include="Data\PendingClient.cs" />
9394
<Compile Include="Data\Transports\ChannelType.cs" />
9495
<Compile Include="Data\FixedQueue.cs" />
9596
<Compile Include="Data\Transports\DefaultTransport.cs" />

‎MLAPI/MonoBehaviours/Core/NetworkingManager.cs‎

Lines changed: 38 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -81,8 +81,10 @@ internal set
8181
/// Gets a list of connected clients
8282
/// </summary>
8383
public readonly List<NetworkedClient> ConnectedClientsList = new List<NetworkedClient>();
84-
internal readonly HashSet<uint> connectionPendingClients = new HashSet<uint>();
85-
internal readonly HashSet<uint> hailPendingClients = new HashSet<uint>();
84+
/// <summary>
85+
/// Gets a dictionary of the clients that have been accepted by the transport but are still pending by the MLAPI.
86+
/// </summary>
87+
public readonly Dictionary<uint, PendingClient> PendingClients = new Dictionary<uint, PendingClient>();
8688
/// <summary>
8789
/// Gets wheter or not a server is running
8890
/// </summary>
@@ -192,12 +194,7 @@ public void SendCustomMessage(uint clientId, Stream stream, string channel = "ML
192194
InternalMessageHandler.Send(clientId, MLAPIConstants.MLAPI_CUSTOM_MESSAGE, channel, stream, new InternalSecuritySendOptions(false, false));
193195
}
194196

195-
196-
#if !DISABLE_CRYPTOGRAPHY
197-
internal readonly Dictionary<uint, EllipticDiffieHellman> pendingKeyExchanges = new Dictionary<uint, EllipticDiffieHellman>();
198197
internal byte[] clientAesKey;
199-
internal readonly Dictionary<uint, byte[]> pendingClientAesKeys = new Dictionary<uint, byte[]>();
200-
#endif
201198

202199
/// <summary>
203200
/// An inspector bool that acts as a Trigger for regenerating RSA keys. Should not be used outside Unity editor.
@@ -282,16 +279,12 @@ private object Init(bool server)
282279
lastEventTickTime = 0f;
283280
lastReceiveTickTime = 0f;
284281
eventOvershootCounter = 0f;
285-
connectionPendingClients.Clear();
286-
hailPendingClients.Clear();
287-
pendingClientAesKeys.Clear();
282+
PendingClients.Clear();
288283
ConnectedClients.Clear();
289284
ConnectedClientsList.Clear();
290285
messageBuffer = new byte[NetworkConfig.MessageBufferSize];
291286
encryptionBuffer = new byte[NetworkConfig.EncryptionBufferSize];
292-
#if !DISABLE_CRYPTOGRAPHY
293-
pendingKeyExchanges.Clear();
294-
#endif
287+
295288
MessageManager.channels.Clear();
296289
MessageManager.reverseChannels.Clear();
297290
SpawnManager.SpawnedObjects.Clear();
@@ -516,27 +509,16 @@ public void StopServer()
516509
NetworkConfig.NetworkTransport.DisconnectClient(pair.Key);
517510
}
518511
}
519-
foreach (uint clientId in connectionPendingClients)
520-
{
521-
if (!disconnectedIds.Contains(clientId))
522-
{
523-
disconnectedIds.Add(clientId);
524-
if (clientId == NetworkConfig.NetworkTransport.ServerClientId)
525-
continue;
526-
527-
NetworkConfig.NetworkTransport.DisconnectClient(clientId);
528-
}
529-
}
530-
531-
foreach (uint clientId in hailPendingClients)
512+
513+
foreach (KeyValuePair<uint, PendingClient> pair in PendingClients)
532514
{
533-
if (!disconnectedIds.Contains(clientId))
515+
if(!disconnectedIds.Contains(pair.Key))
534516
{
535-
disconnectedIds.Add(clientId);
536-
if (clientId == NetworkConfig.NetworkTransport.ServerClientId)
517+
disconnectedIds.Add(pair.Key);
518+
if (pair.Key == NetworkConfig.NetworkTransport.ServerClientId)
537519
continue;
538520

539-
NetworkConfig.NetworkTransport.DisconnectClient(clientId);
521+
NetworkConfig.NetworkTransport.DisconnectClient(pair.Key);
540522
}
541523
}
542524

@@ -710,7 +692,12 @@ private void Update()
710692
EllipticDiffieHellman diffieHellman = new EllipticDiffieHellman(EllipticDiffieHellman.DEFAULT_CURVE, EllipticDiffieHellman.DEFAULT_GENERATOR, EllipticDiffieHellman.DEFAULT_ORDER);
711693
byte[] diffieHellmanPublicPart = diffieHellman.GetPublicKey();
712694
hailWriter.WriteByteArray(diffieHellmanPublicPart);
713-
pendingKeyExchanges.Add(clientId, diffieHellman);
695+
PendingClients.Add(clientId, new PendingClient()
696+
{
697+
ClientId = clientId,
698+
ConnectionState = PendingClient.State.PendingHail,
699+
KeyExchange = diffieHellman
700+
});
714701

715702
if (NetworkConfig.SignKeyExchange)
716703
{
@@ -736,11 +723,14 @@ private void Update()
736723
// Send the hail
737724
InternalMessageHandler.Send(clientId, MLAPIConstants.MLAPI_CERTIFICATE_HAIL, "MLAPI_INTERNAL", hailStream, new InternalSecuritySendOptions(false, false), true);
738725
}
739-
hailPendingClients.Add(clientId);
740726
}
741727
else
742728
{
743-
connectionPendingClients.Add(clientId);
729+
PendingClients.Add(clientId, new PendingClient()
730+
{
731+
ClientId = clientId,
732+
ConnectionState = PendingClient.State.PendingConnection
733+
});
744734
}
745735
StartCoroutine(ApprovalTimeout(clientId));
746736
}
@@ -829,19 +819,12 @@ private IEnumerator ApprovalTimeout(uint clientId)
829819
{
830820
float timeStarted = NetworkTime;
831821
//We yield every frame incase a pending client disconnects and someone else gets its connection id
832-
while (NetworkTime - timeStarted < NetworkConfig.ClientConnectionBufferTimeout && (connectionPendingClients.Contains(clientId) || hailPendingClients.Contains(clientId)))
822+
while (NetworkTime - timeStarted < NetworkConfig.ClientConnectionBufferTimeout && PendingClients.ContainsKey(clientId))
833823
{
834824
yield return null;
835825
}
836826

837-
if (connectionPendingClients.Contains(clientId) && !ConnectedClients.ContainsKey(clientId))
838-
{
839-
// Timeout
840-
if (LogHelper.CurrentLogLevel <= LogLevel.Developer) LogHelper.LogInfo("Client " + clientId + " Handshake Timed Out");
841-
DisconnectClient(clientId);
842-
}
843-
844-
if (hailPendingClients.Contains(clientId) && !ConnectedClients.ContainsKey(clientId))
827+
if (PendingClients.ContainsKey(clientId) && !ConnectedClients.ContainsKey(clientId))
845828
{
846829
// Timeout
847830
if (LogHelper.CurrentLogLevel <= LogLevel.Developer) LogHelper.LogInfo("Client " + clientId + " Handshake Timed Out");
@@ -873,7 +856,7 @@ private void HandleIncomingData(uint clientId, byte[] data, int channelId, int t
873856
using (RijndaelManaged rijndael = new RijndaelManaged())
874857
{
875858
rijndael.Padding = PaddingMode.PKCS7;
876-
rijndael.Key = isServer ? (ConnectedClients.ContainsKey(clientId) ? ConnectedClients[clientId].AesKey : pendingClientAesKeys[clientId]) : clientAesKey;
859+
rijndael.Key = isServer ? (ConnectedClients.ContainsKey(clientId) ? ConnectedClients[clientId].AesKey : PendingClients[clientId].AesKey) : clientAesKey;
877860
rijndael.IV = IVBuffer;
878861
using (CryptoStream cryptoStream = new CryptoStream(bitStream, rijndael.CreateDecryptor(), CryptoStreamMode.Read))
879862
{
@@ -919,8 +902,8 @@ private void HandleIncomingData(uint clientId, byte[] data, int channelId, int t
919902
if (LogHelper.CurrentLogLevel <= LogLevel.Developer) LogHelper.LogInfo("Data Header: messageType=" + messageType);
920903

921904
//Client tried to send a network message that was not the connection request before he was accepted.
922-
if (isServer && (NetworkConfig.EnableEncryption && hailPendingClients.Contains(clientId) && messageType != MLAPIConstants.MLAPI_CERTIFICATE_HAIL_RESPONSE) ||
923-
(connectionPendingClients.Contains(clientId) && messageType != MLAPIConstants.MLAPI_CONNECTION_REQUEST))
905+
if (isServer && (NetworkConfig.EnableEncryption && PendingClients.ContainsKey(clientId) && PendingClients[clientId].ConnectionState == PendingClient.State.PendingHail && messageType != MLAPIConstants.MLAPI_CERTIFICATE_HAIL_RESPONSE) ||
906+
(PendingClients.ContainsKey(clientId) && PendingClients[clientId].ConnectionState == PendingClient.State.PendingConnection && messageType != MLAPIConstants.MLAPI_CONNECTION_REQUEST))
924907
{
925908
if (LogHelper.CurrentLogLevel <= LogLevel.Normal) LogHelper.LogWarning("Message recieved from clientId " + clientId + " before it has been accepted");
926909
return;
@@ -1013,42 +996,25 @@ internal void DisconnectClient(uint clientId)
1013996
if (!isServer)
1014997
return;
1015998

1016-
if (connectionPendingClients.Contains(clientId))
1017-
connectionPendingClients.Remove(clientId);
1018-
1019-
if (hailPendingClients.Contains(clientId))
1020-
hailPendingClients.Remove(clientId);
1021-
1022999
if (ConnectedClients.ContainsKey(clientId))
10231000
ConnectedClients.Remove(clientId);
10241001

1025-
if (pendingClientAesKeys.ContainsKey(clientId))
1026-
pendingClientAesKeys.Remove(clientId);
1002+
if (PendingClients.ContainsKey(clientId))
1003+
PendingClients.Remove(clientId);
10271004

10281005
for (int i = ConnectedClientsList.Count - 1; i > -1; i--)
10291006
{
10301007
if (ConnectedClientsList[i].ClientId == clientId)
10311008
ConnectedClientsList.RemoveAt(i);
10321009
}
10331010

1034-
#if !DISABLE_CRYPTOGRAPHY
1035-
if (pendingKeyExchanges.ContainsKey(clientId))
1036-
pendingKeyExchanges.Remove(clientId);
1037-
#endif
1038-
10391011
NetworkConfig.NetworkTransport.DisconnectClient(clientId);
10401012
}
10411013

10421014
internal void OnClientDisconnectFromServer(uint clientId)
10431015
{
1044-
if (connectionPendingClients.Contains(clientId))
1045-
connectionPendingClients.Remove(clientId);
1046-
1047-
if (hailPendingClients.Contains(clientId))
1048-
hailPendingClients.Remove(clientId);
1049-
1050-
if (pendingClientAesKeys.ContainsKey(clientId))
1051-
pendingClientAesKeys.Remove(clientId);
1016+
if (PendingClients.ContainsKey(clientId))
1017+
PendingClients.Remove(clientId);
10521018

10531019
if (ConnectedClients.ContainsKey(clientId))
10541020
{
@@ -1106,16 +1072,10 @@ internal void HandleApproval(uint clientId, int prefabId, bool approved, Vector3
11061072
{
11071073
if(approved)
11081074
{
1109-
// Inform new client it got approved
1110-
if (connectionPendingClients.Contains(clientId))
1111-
connectionPendingClients.Remove(clientId);
1112-
1113-
if (hailPendingClients.Contains(clientId))
1114-
hailPendingClients.Remove(clientId);
1115-
1116-
byte[] aesKey = pendingClientAesKeys.ContainsKey(clientId) ? pendingClientAesKeys[clientId] : null;
1117-
if (pendingClientAesKeys.ContainsKey(clientId))
1118-
pendingClientAesKeys.Remove(clientId);
1075+
// Inform new client it got approved
1076+
byte[] aesKey = PendingClients.ContainsKey(clientId) ? PendingClients[clientId].AesKey : null;
1077+
if (PendingClients.ContainsKey(clientId))
1078+
PendingClients.Remove(clientId);
11191079
NetworkedClient client = new NetworkedClient()
11201080
{
11211081
ClientId = clientId,
@@ -1230,16 +1190,8 @@ internal void HandleApproval(uint clientId, int prefabId, bool approved, Vector3
12301190
}
12311191
else
12321192
{
1233-
if (connectionPendingClients.Contains(clientId))
1234-
connectionPendingClients.Remove(clientId);
1235-
1236-
if (hailPendingClients.Contains(clientId))
1237-
hailPendingClients.Remove(clientId);
1238-
1239-
#if !DISABLE_CRYPTOGRAPHY
1240-
if (pendingKeyExchanges.ContainsKey(clientId))
1241-
pendingKeyExchanges.Remove(clientId);
1242-
#endif
1193+
if (PendingClients.ContainsKey(clientId))
1194+
PendingClients.Remove(clientId);
12431195

12441196
NetworkConfig.NetworkTransport.DisconnectClient(clientId);
12451197
}

‎MLAPI/NetworkingManagerComponents/Core/InternalMessageHandler.Receive.cs‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -123,15 +123,15 @@ internal static void HandleHailRequest(uint clientId, Stream stream, int channel
123123
// Ran on server
124124
internal static void HandleHailResponse(uint clientId, Stream stream, int channelId)
125125
{
126-
if (!NetworkingManager.singleton.hailPendingClients.Contains(clientId)) return;
127-
if (!NetworkingManager.singleton.NetworkConfig.EnableEncryption) return;
126+
if (!netManager.PendingClients.ContainsKey(clientId) || netManager.PendingClients[clientId].ConnectionState != PendingClient.State.PendingHail) return;
127+
if (!netManager.NetworkConfig.EnableEncryption) return;
128128

129129
using (PooledBitReader reader = PooledBitReader.Get(stream))
130130
{
131-
if (netManager.pendingKeyExchanges.ContainsKey(clientId))
131+
if (NetworkingManager.singleton.PendingClients[clientId].KeyExchange != null)
132132
{
133133
byte[] diffieHellmanPublic = reader.ReadByteArray();
134-
NetworkingManager.singleton.pendingClientAesKeys[clientId] = netManager.pendingKeyExchanges[clientId].GetSharedSecret(diffieHellmanPublic);
134+
netManager.PendingClients[clientId].AesKey = netManager.PendingClients[clientId].KeyExchange.GetSharedSecret(diffieHellmanPublic);
135135
if (netManager.NetworkConfig.SignKeyExchange)
136136
{
137137
byte[] diffieHellmanPublicSignature = reader.ReadByteArray();
@@ -166,9 +166,9 @@ internal static void HandleHailResponse(uint clientId, Stream stream, int channe
166166
}
167167
}
168168

169-
NetworkingManager.singleton.hailPendingClients.Remove(clientId);
170-
NetworkingManager.singleton.connectionPendingClients.Add(clientId);
171-
NetworkingManager.singleton.pendingKeyExchanges.Remove(clientId);
169+
netManager.PendingClients[clientId].ConnectionState = PendingClient.State.PendingConnection;
170+
netManager.PendingClients[clientId].KeyExchange = null; // Give to GC
171+
172172
// Send greetings, they have passed all the handshakes
173173
using (PooledBitStream outStream = PooledBitStream.Get())
174174
{

0 commit comments

Comments
 (0)