fix: Ensure NetworkConnectionManager correctly handles multiple disconnect messages being sent

Author: EmandM

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

@@ -1441,18 +1441,13 @@ private void HostServerInitialize()
                     }
                 }
 
-                response.Approved = true;
-                ConnectionManager.HandleConnectionApproval(ServerClientId, response);
+                ConnectionManager.HandleConnectionApproval(ServerClientId, response.CreatePlayerObject, response.PlayerPrefabHash, response.Position, response.Rotation);
             }
             else
             {
-                var response = new ConnectionApprovalResponse
-                {
-                    Approved = true,
-                    // Distributed authority always returns true since the client side handles spawning (whether automatically or manually)
-                    CreatePlayerObject = DistributedAuthorityMode || NetworkConfig.PlayerPrefab != null,
-                };
-                ConnectionManager.HandleConnectionApproval(ServerClientId, response);
+                // Distributed authority always tries to create the player object since the client side handles spawning (whether automatically or manually)
+                var createPlayerObject = DistributedAuthorityMode || NetworkConfig.PlayerPrefab != null;
+                ConnectionManager.HandleConnectionApproval(ServerClientId, createPlayerObject);
             }
 
             SpawnManager.ServerSpawnSceneObjectsOnStartSweep();
@@ -1474,14 +1469,22 @@ private void HostServerInitialize()
         /// </summary>
         /// <param name="clientId">The ClientId to get the TransportId from</param>
         /// <returns>The TransportId associated with the given ClientId</returns>
-        public ulong GetTransportIdFromClientId(ulong clientId) => ConnectionManager.ClientIdToTransportId(clientId);
+        public ulong GetTransportIdFromClientId(ulong clientId)
+        {
+            var (id, success) =  ConnectionManager.TransportIdToClientId(clientId);
+            return success ? id : ulong.MaxValue;
+        }
 
         /// <summary>
         /// Get the ClientId from the associated TransportId.
         /// </summary>
         /// <param name="transportId">The TransportId to get the ClientId from</param>
         /// <returns>The ClientId from the associated TransportId</returns>
-        public ulong GetClientIdFromTransportId(ulong transportId) => ConnectionManager.TransportIdToClientId(transportId);
+        public ulong GetClientIdFromTransportId(ulong transportId)
+        {
+            var (id, success) =  ConnectionManager.TransportIdToClientId(transportId);
+            return success ? id : ulong.MaxValue;
+        }
 
         /// <summary>
         /// Disconnects the remote client.

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

@@ -19,8 +19,19 @@ public DefaultMessageSender(NetworkManager manager)
         public void Send(ulong clientId, NetworkDelivery delivery, FastBufferWriter batchData)
         {
             var sendBuffer = batchData.ToTempByteArray();
+            var (transportId, clientExists) = m_ConnectionManager.ClientIdToTransportId(clientId);
 
-            m_NetworkTransport.Send(m_ConnectionManager.ClientIdToTransportId(clientId), sendBuffer, delivery);
+            if (!clientExists)
+            {
+                if (m_ConnectionManager.NetworkManager.LogLevel <= LogLevel.Error)
+                {
+                    NetworkLog.LogWarning("Trying to send a message to a client who doesn't have a transport connection");
+                }
+
+                return;
+            }
+
+            m_NetworkTransport.Send(transportId, sendBuffer, delivery);
         }
     }
 }

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

@@ -1,4 +1,5 @@
 using Unity.Collections;
+using UnityEngine;
 
 namespace Unity.Netcode
 {
@@ -210,12 +211,16 @@ public void Handle(ref NetworkContext context)
             }
             else
             {
-                var response = new NetworkManager.ConnectionApprovalResponse
+                var createPlayerObject = networkManager.NetworkConfig.PlayerPrefab != null;
+
+                // DAHost only:
+                // Never create the player object on the server if AutoSpawnPlayerPrefabClientSide is set.
+                if (networkManager.DistributedAuthorityMode && networkManager.AutoSpawnPlayerPrefabClientSide)
                 {
-                    Approved = true,
-                    CreatePlayerObject = networkManager.DistributedAuthorityMode && networkManager.AutoSpawnPlayerPrefabClientSide ? false : networkManager.NetworkConfig.PlayerPrefab != null
-                };
-                networkManager.ConnectionManager.HandleConnectionApproval(senderId, response);
+                    createPlayerObject = false;
+                }
+
+                networkManager.ConnectionManager.HandleConnectionApproval(senderId, createPlayerObject);
             }
         }
     }

com.unity.netcode.gameobjects/Runtime/Messaging/Messages/ConnectionRequestMessage.cs

@@ -909,7 +909,11 @@ private void SendBatchedMessages(SendTarget sendTarget, BatchedSendQueue queue)
             var mtu = 0;
             if (m_NetworkManager)
             {
-                var ngoClientId = m_NetworkManager.ConnectionManager.TransportIdToClientId(sendTarget.ClientId);
+                var (ngoClientId, isConnectedClient) = m_NetworkManager.ConnectionManager.TransportIdToClientId(sendTarget.ClientId);
+                if (!isConnectedClient)
+                {
+                    return;
+                }
                 mtu = m_NetworkManager.GetPeerMTU(ngoClientId);
             }
 
@@ -1321,7 +1325,7 @@ public override ulong GetCurrentRtt(ulong clientId)
 
             if (m_NetworkManager != null)
             {
-                var transportId = m_NetworkManager.ConnectionManager.ClientIdToTransportId(clientId);
+                var (transportId, _) = m_NetworkManager.ConnectionManager.ClientIdToTransportId(clientId);
 
                 var rtt = ExtractRtt(ParseClientId(transportId));
                 if (rtt > 0)
@@ -1347,13 +1351,14 @@ public NetworkEndpoint GetEndpoint(ulong clientId)
         {
             if (m_Driver.IsCreated && m_NetworkManager != null && m_NetworkManager.IsListening)
             {
-                var transportId = m_NetworkManager.ConnectionManager.ClientIdToTransportId(clientId);
+                var (transportId, connectionExists) = m_NetworkManager.ConnectionManager.ClientIdToTransportId(clientId);
                 var networkConnection = ParseClientId(transportId);
-                if (m_Driver.GetConnectionState(networkConnection) == NetworkConnection.State.Connected)
+                if (connectionExists && m_Driver.GetConnectionState(networkConnection) == NetworkConnection.State.Connected)
                 {
                     return m_Driver.GetRemoteEndpoint(networkConnection);
                 }
             }
+
             return new NetworkEndpoint();
         }
 
@@ -1447,10 +1452,18 @@ public override void Send(ulong clientId, ArraySegment<byte> payload, NetworkDel
                     // If the message is sent reliably, then we're over capacity and we can't
                     // provide any reliability guarantees anymore. Disconnect the client since at
                     // this point they're bound to become desynchronized.
+                    if (m_NetworkManager != null)
+                    {
+                        var (ngoClientId, isConnectedClient) = m_NetworkManager.ConnectionManager.TransportIdToClientId(clientId);
+                        if (isConnectedClient)
+                        {
+                            clientId = ngoClientId;
+                        }
+
+                    }
 
-                    var ngoClientId = m_NetworkManager?.ConnectionManager.TransportIdToClientId(clientId) ?? clientId;
                     Debug.LogError($"Couldn't add payload of size {payload.Count} to reliable send queue. " +
-                        $"Closing connection {ngoClientId} as reliability guarantees can't be maintained.");
+                        $"Closing connection {clientId} as reliability guarantees can't be maintained.");
 
                     if (clientId == m_ServerClientId)
                     {

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

@@ -108,16 +108,24 @@ private void OnConnectionEvent(NetworkManager networkManager, ConnectionEventDat
         /// </summary>
         private bool TransportIdCleanedUp()
         {
-            if (m_ServerNetworkManager.ConnectionManager.TransportIdToClientId(m_TransportClientId) == m_ClientId)
+            var (clientId, isConnected) = m_ServerNetworkManager.ConnectionManager.TransportIdToClientId(m_TransportClientId);
+            if (isConnected)
             {
                 return false;
             }
 
-            if (m_ServerNetworkManager.ConnectionManager.ClientIdToTransportId(m_ClientId) == m_TransportClientId)
+            if (clientId == m_ClientId)
             {
                 return false;
             }
-            return true;
+
+            var (transportId, connectionExists) = m_ServerNetworkManager.ConnectionManager.ClientIdToTransportId(m_ClientId);
+            if (connectionExists)
+            {
+                return false;
+            }
+
+            return transportId != m_TransportClientId;
         }
 
         /// <summary>
@@ -145,7 +153,9 @@ public IEnumerator ClientPlayerDisconnected([Values] ClientDisconnectType client
 
             var serverSideClientPlayer = m_ServerNetworkManager.ConnectionManager.ConnectedClients[m_ClientId].PlayerObject;
 
-            m_TransportClientId = m_ServerNetworkManager.ConnectionManager.ClientIdToTransportId(m_ClientId);
+            bool connectionExists;
+            (m_TransportClientId, connectionExists) = m_ServerNetworkManager.ConnectionManager.ClientIdToTransportId(m_ClientId);
+            Assert.IsTrue(connectionExists);
 
             if (clientDisconnectType == ClientDisconnectType.ServerDisconnectsClient)
             {

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

@@ -484,6 +484,8 @@ protected void SetDistributedAuthorityProperties(NetworkManager networkManager)
         /// </summary>
         protected virtual bool m_EnableTimeTravel => false;
 
+        protected virtual bool m_UseMockTransport => m_EnableTimeTravel;
+
         /// <summary>
         /// If this is false, SetUp will call OnInlineSetUp instead of OnSetUp.
         /// This is a performance advantage when not using the coroutine functionality, as a coroutine that
@@ -637,7 +639,7 @@ public IEnumerator SetUp()
             VerboseDebugLog.Clear();
             VerboseDebug($"Entering {nameof(SetUp)}");
             NetcodeLogAssert = new NetcodeLogAssert();
-            if (m_EnableTimeTravel)
+            if (m_UseMockTransport)
             {
                 if (m_NetworkManagerInstatiationMode == NetworkManagerInstatiationMode.AllTests)
                 {
@@ -780,7 +782,7 @@ protected void CreateServerAndClients(int numberOfClients)
             }
 
             // Create multiple NetworkManager instances
-            if (!NetcodeIntegrationTestHelpers.Create(numberOfClients, out NetworkManager server, out NetworkManager[] clients, m_TargetFrameRate, m_CreateServerFirst, m_EnableTimeTravel, m_UseCmbService))
+            if (!NetcodeIntegrationTestHelpers.Create(numberOfClients, out NetworkManager server, out NetworkManager[] clients, m_TargetFrameRate, m_CreateServerFirst, m_UseMockTransport, m_UseCmbService))
             {
                 Debug.LogError("Failed to create instances");
                 Assert.Fail("Failed to create instances");
@@ -871,7 +873,7 @@ protected virtual bool ShouldWaitForNewClientToConnect(NetworkManager networkMan
         /// <returns>The newly created <see cref="NetworkManager"/>.</returns>
         protected NetworkManager CreateNewClient()
         {
-            var networkManager = NetcodeIntegrationTestHelpers.CreateNewClient(m_ClientNetworkManagers.Length, m_EnableTimeTravel, m_UseCmbService);
+            var networkManager = NetcodeIntegrationTestHelpers.CreateNewClient(m_ClientNetworkManagers.Length, m_UseMockTransport, m_UseCmbService);
             networkManager.NetworkConfig.PlayerPrefab = m_PlayerPrefab;
             SetDistributedAuthorityProperties(networkManager);
 
@@ -980,7 +982,7 @@ private bool AllPlayerObjectClonesSpawned(NetworkManager joinedClient)
         /// </summary>
         protected void CreateAndStartNewClientWithTimeTravel()
         {
-            var networkManager = NetcodeIntegrationTestHelpers.CreateNewClient(m_ClientNetworkManagers.Length, m_EnableTimeTravel);
+            var networkManager = NetcodeIntegrationTestHelpers.CreateNewClient(m_ClientNetworkManagers.Length, m_UseMockTransport);
             networkManager.NetworkConfig.PlayerPrefab = m_PlayerPrefab;
             SetDistributedAuthorityProperties(networkManager);
 

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

@@ -0,0 +1,48 @@
+using System.Collections;
+using System.Linq;
+using NUnit.Framework;
+using Unity.Netcode.TestHelpers.Runtime;
+using UnityEngine.TestTools;
+
+namespace Unity.Netcode.RuntimeTests
+{
+    [TestFixture(HostOrServer.Server)]
+    [TestFixture(HostOrServer.Host)]
+    internal class TransportTests : NetcodeIntegrationTest
+    {
+        protected override int NumberOfClients => 2;
+
+        protected override bool m_UseMockTransport => true;
+
+        public TransportTests(HostOrServer hostOrServer) : base(hostOrServer) { }
+
+        [UnityTest]
+        public IEnumerator MultipleDisconnectEventsNoop()
+        {
+            var clientToDisconnect = GetNonAuthorityNetworkManager(0);
+            var clientTransport = clientToDisconnect.NetworkConfig.NetworkTransport;
+
+            var otherClient = GetNonAuthorityNetworkManager(1);
+
+            // Send multiple disconnect events
+            clientTransport.DisconnectLocalClient();
+            clientTransport.DisconnectLocalClient();
+
+            // completely stop and clean up the client 
+            yield return StopOneClient(clientToDisconnect);
+
+            var expectedConnectedClients = m_UseHost ? NumberOfClients : NumberOfClients - 1;
+            yield return WaitForConditionOrTimeOut(() => otherClient.ConnectedClientsIds.Count == expectedConnectedClients);
+            AssertOnTimeout($"Incorrect number of connected clients. Expected: {expectedConnectedClients}, have: {otherClient.ConnectedClientsIds.Count}");
+
+            // Start a new client to ensure everything is still working
+            yield return CreateAndStartNewClient();
+
+            var newExpectedClients = m_UseHost ? NumberOfClients + 1 : NumberOfClients;
+            yield return WaitForConditionOrTimeOut(() => otherClient.ConnectedClientsIds.Count == newExpectedClients);
+            AssertOnTimeout($"Incorrect number of connected clients. Expected: {newExpectedClients}, have: {otherClient.ConnectedClientsIds.Count}");
+
+
+        }
+    }
+}