Skip to content

Commit c2f3f83

Browse files
EmandMEmandM
committed
fix: [Backport] Ensure NetworkConnectionManager correctly handles multiple disconnect messages being sent
1 parent ef07caa commit c2f3f83

File tree

8 files changed

+159
-38
lines changed

8 files changed

+159
-38
lines changed

‎com.unity.netcode.gameobjects/Runtime/Connection/NetworkConnectionManager.cs‎

Lines changed: 30 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -317,45 +317,45 @@ internal void RemovePendingClient(ulong clientId)
317317
private ulong m_NextClientId = 1;
318318

319319
[MethodImpl(MethodImplOptions.AggressiveInlining)]
320-
internal ulong TransportIdToClientId(ulong transportId)
320+
internal (ulong, bool) TransportIdToClientId(ulong transportId)
321321
{
322322
if (transportId == GetServerTransportId())
323323
{
324-
return NetworkManager.ServerClientId;
324+
return (NetworkManager.ServerClientId, true);
325325
}
326326

327327
if (TransportIdToClientIdMap.TryGetValue(transportId, out var clientId))
328328
{
329-
return clientId;
329+
return (clientId, true);
330330
}
331331

332332
if (NetworkLog.CurrentLogLevel == LogLevel.Developer)
333333
{
334334
NetworkLog.LogWarning($"Trying to get the NGO client ID map for the transport ID ({transportId}) but did not find the map entry! Returning default transport ID value.");
335335
}
336336

337-
return default;
337+
return (0, false);
338338
}
339339

340340
[MethodImpl(MethodImplOptions.AggressiveInlining)]
341-
internal ulong ClientIdToTransportId(ulong clientId)
341+
internal (ulong, bool) ClientIdToTransportId(ulong clientId)
342342
{
343343
if (clientId == NetworkManager.ServerClientId)
344344
{
345-
return GetServerTransportId();
345+
return (GetServerTransportId(), true);
346346
}
347347

348348
if (ClientIdToTransportIdMap.TryGetValue(clientId, out var transportClientId))
349349
{
350-
return transportClientId;
350+
return (transportClientId, true);
351351
}
352352

353353
if (NetworkLog.CurrentLogLevel == LogLevel.Developer)
354354
{
355355
NetworkLog.LogWarning($"Trying to get the transport client ID map for the NGO client ID ({clientId}) but did not find the map entry! Returning default transport ID value.");
356356
}
357357

358-
return default;
358+
return (0, false);
359359
}
360360

361361
/// <summary>
@@ -384,19 +384,24 @@ private ulong GetServerTransportId()
384384
/// Handles cleaning up the transport id/client id tables after receiving a disconnect event from transport.
385385
/// </summary>
386386
[MethodImpl(MethodImplOptions.AggressiveInlining)]
387-
internal ulong TransportIdCleanUp(ulong transportId)
387+
internal (ulong, bool) TransportIdCleanUp(ulong transportId)
388388
{
389389
// This check is for clients that attempted to connect but failed.
390390
// When this happens, the client will not have an entry within the m_TransportIdToClientIdMap or m_ClientIdToTransportIdMap lookup tables so we exit early and just return 0 to be used for the disconnect event.
391391
if (!LocalClient.IsServer && !TransportIdToClientIdMap.ContainsKey(transportId))
392392
{
393-
return NetworkManager.LocalClientId;
393+
return (NetworkManager.LocalClientId, true);
394+
}
395+
396+
var (clientId, isConnectedClient) = TransportIdToClientId(transportId);
397+
if (!isConnectedClient)
398+
{
399+
return (default, false);
394400
}
395401

396-
var clientId = TransportIdToClientId(transportId);
397402
TransportIdToClientIdMap.Remove(transportId);
398403
ClientIdToTransportIdMap.Remove(clientId);
399-
return clientId;
404+
return (clientId, true);
400405
}
401406

402407
internal void PollAndHandleNetworkEvents()
@@ -502,8 +507,11 @@ internal void DataEventHandler(ulong transportClientId, ref ArraySegment<byte> p
502507
#if DEVELOPMENT_BUILD || UNITY_EDITOR
503508
s_HandleIncomingData.Begin();
504509
#endif
505-
var clientId = TransportIdToClientId(transportClientId);
506-
MessageManager.HandleIncomingData(clientId, payload, receiveTime);
510+
var (clientId, isConnectedClient) = TransportIdToClientId(transportClientId);
511+
if (isConnectedClient)
512+
{
513+
MessageManager.HandleIncomingData(clientId, payload, receiveTime);
514+
}
507515

508516
#if DEVELOPMENT_BUILD || UNITY_EDITOR
509517
s_HandleIncomingData.End();
@@ -515,10 +523,15 @@ internal void DataEventHandler(ulong transportClientId, ref ArraySegment<byte> p
515523
/// </summary>
516524
internal void DisconnectEventHandler(ulong transportClientId)
517525
{
526+
var (clientId, wasConnectedClient) = TransportIdCleanUp(transportClientId);
527+
if (!wasConnectedClient)
528+
{
529+
return;
530+
}
531+
518532
#if DEVELOPMENT_BUILD || UNITY_EDITOR
519533
s_TransportDisconnect.Begin();
520534
#endif
521-
var clientId = TransportIdCleanUp(transportClientId);
522535
if (NetworkLog.CurrentLogLevel <= LogLevel.Developer)
523536
{
524537
NetworkLog.LogInfo($"Disconnect Event From {clientId}");
@@ -1040,9 +1053,9 @@ internal void OnClientDisconnectFromServer(ulong clientId)
10401053
}
10411054

10421055
// If the client ID transport map exists
1043-
if (ClientIdToTransportIdMap.ContainsKey(clientId))
1056+
var (transportId, isConnected) = ClientIdToTransportId(clientId);
1057+
if (isConnected)
10441058
{
1045-
var transportId = ClientIdToTransportId(clientId);
10461059
NetworkManager.NetworkConfig.NetworkTransport.DisconnectRemoteClient(transportId);
10471060

10481061
InvokeOnClientDisconnectCallback(clientId);

‎com.unity.netcode.gameobjects/Runtime/Core/NetworkManager.cs‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1156,15 +1156,27 @@ private void HostServerInitialize()
11561156
/// Get the TransportId from the associated ClientId.
11571157
/// </summary>
11581158
/// <param name="clientId">The ClientId to get the TransportId from</param>
1159-
/// <returns>The TransportId associated with the given ClientId</returns>
1160-
public ulong GetTransportIdFromClientId(ulong clientId) => ConnectionManager.ClientIdToTransportId(clientId);
1159+
/// <returns>
1160+
/// The TransportId associated with the given ClientId if the given clientId is valid; otherwise <see cref="ulong.MaxValue"/>
1161+
/// </returns>
1162+
public ulong GetTransportIdFromClientId(ulong clientId)
1163+
{
1164+
var (id, success) = ConnectionManager.ClientIdToTransportId(clientId);
1165+
return success ? id : ulong.MaxValue;
1166+
}
11611167

11621168
/// <summary>
11631169
/// Get the ClientId from the associated TransportId.
11641170
/// </summary>
11651171
/// <param name="transportId">The TransportId to get the ClientId from</param>
1166-
/// <returns>The ClientId from the associated TransportId</returns>
1167-
public ulong GetClientIdFromTransportId(ulong transportId) => ConnectionManager.TransportIdToClientId(transportId);
1172+
/// <returns>
1173+
/// The ClientId from the associated TransportId if the given transportId is valid; otherwise <see cref="ulong.MaxValue"/>
1174+
/// </returns>
1175+
public ulong GetClientIdFromTransportId(ulong transportId)
1176+
{
1177+
var (id, success) = ConnectionManager.TransportIdToClientId(transportId);
1178+
return success ? id : ulong.MaxValue;
1179+
}
11681180

11691181
/// <summary>
11701182
/// Disconnects the remote client.

‎com.unity.netcode.gameobjects/Runtime/Messaging/DefaultMessageSender.cs‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,19 @@ public DefaultMessageSender(NetworkManager manager)
1919
public void Send(ulong clientId, NetworkDelivery delivery, FastBufferWriter batchData)
2020
{
2121
var sendBuffer = batchData.ToTempByteArray();
22+
var (transportId, clientExists) = m_ConnectionManager.ClientIdToTransportId(clientId);
2223

23-
m_NetworkTransport.Send(m_ConnectionManager.ClientIdToTransportId(clientId), sendBuffer, delivery);
24+
if (!clientExists)
25+
{
26+
if (m_ConnectionManager.NetworkManager.LogLevel <= LogLevel.Error)
27+
{
28+
NetworkLog.LogWarning("Trying to send a message to a client who doesn't have a transport connection");
29+
}
30+
31+
return;
32+
}
33+
34+
m_NetworkTransport.Send(transportId, sendBuffer, delivery);
2435
}
2536
}
2637
}

‎com.unity.netcode.gameobjects/Runtime/Transports/UTP/UnityTransport.cs‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -868,7 +868,11 @@ private void SendBatchedMessages(SendTarget sendTarget, BatchedSendQueue queue)
868868
var mtu = 0;
869869
if (NetworkManager)
870870
{
871-
var ngoClientId = NetworkManager.ConnectionManager.TransportIdToClientId(sendTarget.ClientId);
871+
var (ngoClientId, isConnectedClient) = NetworkManager.ConnectionManager.TransportIdToClientId(sendTarget.ClientId);
872+
if (!isConnectedClient)
873+
{
874+
return;
875+
}
872876
mtu = NetworkManager.GetPeerMTU(ngoClientId);
873877
}
874878

@@ -1278,7 +1282,7 @@ public override ulong GetCurrentRtt(ulong clientId)
12781282

12791283
if (NetworkManager != null)
12801284
{
1281-
var transportId = NetworkManager.ConnectionManager.ClientIdToTransportId(clientId);
1285+
var (transportId, _) = NetworkManager.ConnectionManager.ClientIdToTransportId(clientId);
12821286

12831287
var rtt = ExtractRtt(ParseClientId(transportId));
12841288
if (rtt > 0)
@@ -1329,9 +1333,9 @@ public NetworkEndpoint GetEndpoint(ulong clientId)
13291333
{
13301334
if (m_Driver.IsCreated && NetworkManager != null && NetworkManager.IsListening)
13311335
{
1332-
var transportId = NetworkManager.ConnectionManager.ClientIdToTransportId(clientId);
1336+
var (transportId, connectionExists) = NetworkManager.ConnectionManager.ClientIdToTransportId(clientId);
13331337
var networkConnection = ParseClientId(transportId);
1334-
if (m_Driver.GetConnectionState(networkConnection) == NetworkConnection.State.Connected)
1338+
if (connectionExists && m_Driver.GetConnectionState(networkConnection) == NetworkConnection.State.Connected)
13351339
{
13361340
return m_Driver.RemoteEndPoint(networkConnection);
13371341
}
@@ -1460,10 +1464,17 @@ public override void Send(ulong clientId, ArraySegment<byte> payload, NetworkDel
14601464
// If the message is sent reliably, then we're over capacity and we can't
14611465
// provide any reliability guarantees anymore. Disconnect the client since at
14621466
// this point they're bound to become desynchronized.
1467+
if (NetworkManager != null)
1468+
{
1469+
var (ngoClientId, isConnectedClient) = NetworkManager.ConnectionManager.TransportIdToClientId(clientId);
1470+
if (isConnectedClient)
1471+
{
1472+
clientId = ngoClientId;
1473+
}
14631474

1464-
var ngoClientId = NetworkManager?.ConnectionManager.TransportIdToClientId(clientId) ?? clientId;
1475+
}
14651476
Debug.LogError($"Couldn't add payload of size {payload.Count} to reliable send queue. " +
1466-
$"Closing connection {ngoClientId} as reliability guarantees can't be maintained.");
1477+
$"Closing connection {clientId} as reliability guarantees can't be maintained.");
14671478

14681479
if (clientId == m_ServerClientId)
14691480
{

‎com.unity.netcode.gameobjects/TestHelpers/Runtime/NetcodeIntegrationTest.cs‎

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -287,6 +287,17 @@ public enum HostOrServer
287287
/// </summary>
288288
protected virtual bool m_EnableTimeTravel => false;
289289

290+
/// <summary>
291+
/// When true, <see cref="CreateServerAndClients()"/> and <see cref="CreateNewClient"/> will use a <see cref="MockTransport"/>
292+
/// as the <see cref="NetworkConfig.NetworkTransport"/> on the created server and/or clients.
293+
/// When false, a <see cref="UnityTransport"/> is used.
294+
/// </summary>
295+
/// <remarks>
296+
/// This defaults to, and is required to be true when <see cref="m_EnableTimeTravel"/> is true.
297+
/// <see cref="m_EnableTimeTravel"/> will not work with the <see cref="UnityTransport"/> component.
298+
/// </remarks>
299+
protected virtual bool m_UseMockTransport => m_EnableTimeTravel;
300+
290301
/// <summary>
291302
/// If this is false, SetUp will call OnInlineSetUp instead of OnSetUp.
292303
/// This is a performance advantage when not using the coroutine functionality, as a coroutine that
@@ -407,7 +418,7 @@ public IEnumerator SetUp()
407418
{
408419
VerboseDebug($"Entering {nameof(SetUp)}");
409420
NetcodeLogAssert = new NetcodeLogAssert();
410-
if (m_EnableTimeTravel)
421+
if (m_UseMockTransport)
411422
{
412423
if (m_NetworkManagerInstatiationMode == NetworkManagerInstatiationMode.AllTests)
413424
{
@@ -417,9 +428,14 @@ public IEnumerator SetUp()
417428
{
418429
MockTransport.Reset();
419430
}
420-
// Setup the frames per tick for time travel advance to next tick
431+
}
432+
433+
// Setup the frames per tick for time travel advance to next tick
434+
if (m_EnableTimeTravel)
435+
{
421436
ConfigureFramesPerTick();
422437
}
438+
423439
if (m_SetupIsACoroutine)
424440
{
425441
yield return OnSetup();
@@ -577,7 +593,7 @@ protected virtual bool ShouldWaitForNewClientToConnect(NetworkManager networkMan
577593
/// <returns>An IEnumerator for use with Unity's coroutine system.</returns>
578594
protected IEnumerator CreateAndStartNewClient()
579595
{
580-
var networkManager = NetcodeIntegrationTestHelpers.CreateNewClient(m_ClientNetworkManagers.Length, m_EnableTimeTravel);
596+
var networkManager = NetcodeIntegrationTestHelpers.CreateNewClient(m_ClientNetworkManagers.Length, m_UseMockTransport);
581597
networkManager.NetworkConfig.PlayerPrefab = m_PlayerPrefab;
582598

583599
// Notification that the new client (NetworkManager) has been created
@@ -619,7 +635,7 @@ protected IEnumerator CreateAndStartNewClient()
619635
/// </summary>
620636
protected void CreateAndStartNewClientWithTimeTravel()
621637
{
622-
var networkManager = NetcodeIntegrationTestHelpers.CreateNewClient(m_ClientNetworkManagers.Length, m_EnableTimeTravel);
638+
var networkManager = NetcodeIntegrationTestHelpers.CreateNewClient(m_ClientNetworkManagers.Length, m_UseMockTransport);
623639
networkManager.NetworkConfig.PlayerPrefab = m_PlayerPrefab;
624640

625641
// Notification that the new client (NetworkManager) has been created
@@ -732,7 +748,7 @@ protected void CreateServerAndClients(int numberOfClients)
732748
}
733749

734750
// Create multiple NetworkManager instances
735-
if (!NetcodeIntegrationTestHelpers.Create(numberOfClients, out NetworkManager server, out NetworkManager[] clients, m_TargetFrameRate, m_CreateServerFirst, m_EnableTimeTravel))
751+
if (!NetcodeIntegrationTestHelpers.Create(numberOfClients, out NetworkManager server, out NetworkManager[] clients, m_TargetFrameRate, m_CreateServerFirst, m_UseMockTransport))
736752
{
737753
Debug.LogError("Failed to create instances");
738754
Assert.Fail("Failed to create instances");
@@ -1219,7 +1235,7 @@ public IEnumerator TearDown()
12191235
VerboseDebug($"Exiting {nameof(TearDown)}");
12201236
LogWaitForMessages();
12211237
NetcodeLogAssert.Dispose();
1222-
if (m_EnableTimeTravel)
1238+
if (m_UseMockTransport)
12231239
{
12241240
if (m_NetworkManagerInstatiationMode == NetworkManagerInstatiationMode.AllTests)
12251241
{

‎com.unity.netcode.gameobjects/Tests/Runtime/DisconnectTests.cs‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -108,16 +108,24 @@ private void OnConnectionEvent(NetworkManager networkManager, ConnectionEventDat
108108
/// </summary>
109109
private bool TransportIdCleanedUp()
110110
{
111-
if (m_ServerNetworkManager.ConnectionManager.TransportIdToClientId(m_TransportClientId) == m_ClientId)
111+
var (clientId, isConnected) = m_ServerNetworkManager.ConnectionManager.TransportIdToClientId(m_TransportClientId);
112+
if (isConnected)
112113
{
113114
return false;
114115
}
115116

116-
if (m_ServerNetworkManager.ConnectionManager.ClientIdToTransportId(m_ClientId) == m_TransportClientId)
117+
if (clientId == m_ClientId)
117118
{
118119
return false;
119120
}
120-
return true;
121+
122+
var (transportId, connectionExists) = m_ServerNetworkManager.ConnectionManager.ClientIdToTransportId(m_ClientId);
123+
if (connectionExists)
124+
{
125+
return false;
126+
}
127+
128+
return transportId != m_TransportClientId;
121129
}
122130

123131
/// <summary>
@@ -144,7 +152,9 @@ public IEnumerator ClientPlayerDisconnected([Values] ClientDisconnectType client
144152

145153
var serverSideClientPlayer = m_ServerNetworkManager.ConnectionManager.ConnectedClients[m_ClientId].PlayerObject;
146154

147-
m_TransportClientId = m_ServerNetworkManager.ConnectionManager.ClientIdToTransportId(m_ClientId);
155+
bool connectionExists;
156+
(m_TransportClientId, connectionExists) = m_ServerNetworkManager.ConnectionManager.ClientIdToTransportId(m_ClientId);
157+
Assert.IsTrue(connectionExists);
148158

149159
var clientManager = m_ClientNetworkManagers[0];
150160

‎com.unity.netcode.gameobjects/Tests/Runtime/Transports/TransportTests.cs‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
using System.Collections;
2+
using NUnit.Framework;
3+
using Unity.Netcode.TestHelpers.Runtime;
4+
using UnityEngine.TestTools;
5+
6+
namespace Unity.Netcode.RuntimeTests
7+
{
8+
[TestFixture(HostOrServer.Server)]
9+
[TestFixture(HostOrServer.Host)]
10+
internal class TransportTests : NetcodeIntegrationTest
11+
{
12+
protected override int NumberOfClients => 2;
13+
14+
protected override bool m_UseMockTransport => true;
15+
16+
public TransportTests(HostOrServer hostOrServer) : base(hostOrServer) { }
17+
18+
[UnityTest]
19+
public IEnumerator MultipleDisconnectEventsNoop()
20+
{
21+
var clientToDisconnect = m_ClientNetworkManagers[0];
22+
var clientTransport = clientToDisconnect.NetworkConfig.NetworkTransport;
23+
24+
var otherClient = m_ClientNetworkManagers[1];
25+
26+
// Send multiple disconnect events
27+
clientTransport.DisconnectLocalClient();
28+
clientTransport.DisconnectLocalClient();
29+
30+
// completely stop and clean up the client
31+
yield return StopOneClient(clientToDisconnect);
32+
33+
var expectedConnectedClients = m_UseHost ? NumberOfClients : NumberOfClients - 1;
34+
yield return WaitForConditionOrTimeOut(() => otherClient.ConnectedClientsIds.Count == expectedConnectedClients);
35+
AssertOnTimeout($"Incorrect number of connected clients. Expected: {expectedConnectedClients}, have: {otherClient.ConnectedClientsIds.Count}");
36+
37+
// Start a new client to ensure everything is still working
38+
yield return CreateAndStartNewClient();
39+
40+
var newExpectedClients = m_UseHost ? NumberOfClients + 1 : NumberOfClients;
41+
yield return WaitForConditionOrTimeOut(() => otherClient.ConnectedClientsIds.Count == newExpectedClients);
42+
AssertOnTimeout($"Incorrect number of connected clients. Expected: {newExpectedClients}, have: {otherClient.ConnectedClientsIds.Count}");
43+
}
44+
}
45+
}

0 commit comments

Comments
 (0)