feat: disconnect reason (#2280)

* disconnect reason

Author: jeffreyrainy

com.unity.netcode.gameobjects/CHANGELOG.md

@@ -11,6 +11,7 @@ Additional documentation and release notes are available at [Multiplayer Documen
 
 ### Added
 - Added `NetworkObject` auto-add helper and Multiplayer Tools install reminder settings to Project Settings. (#2285)
+- Added `public string DisconnectReason` getter to `NetworkManager` and `string Reason` to `ConnectionApprovalResponse`. Allows connection approval to communicate back a reason. Also added `public void DisconnectClient(ulong clientId, string reason)` allowing setting a disconnection reason, when explicitly disconnecting a client. 
 
 ### Fixed
 

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

@@ -86,6 +86,12 @@ public NetworkPrefabHandler PrefabHandler
         private bool m_ShuttingDown;
         private bool m_StopProcessingMessages;
 
+        // <summary>
+        // When disconnected from the server, the server may send a reason. If a reason was sent, this property will
+        // tell client code what the reason was. It should be queried after the OnClientDisconnectCallback is called
+        // </summary>
+        public string DisconnectReason { get; internal set; }
+
         private class NetworkManagerHooks : INetworkHooks
         {
             private NetworkManager m_NetworkManager;
@@ -443,6 +449,11 @@ public class ConnectionApprovalResponse
             /// If the Approval decision cannot be made immediately, the client code can set Pending to true, keep a reference to the ConnectionApprovalResponse object and write to it later. Client code must exercise care to setting all the members to the value it wants before marking Pending to false, to indicate completion. If the field is set as Pending = true, we'll monitor the object until it gets set to not pending anymore and use the parameters then.
             /// </summary>
             public bool Pending;
+
+            // <summary>
+            // Optional reason. If Approved is false, this reason will be sent to the client so they know why they
+            // were not approved.
+            public string Reason;
         }
 
         /// <summary>
@@ -889,6 +900,7 @@ private void Initialize(bool server)
                 return;
             }
 
+            DisconnectReason = string.Empty;
             IsApproved = false;
 
             ComponentFactory.SetDefaults();
@@ -2004,12 +2016,31 @@ internal void HandleIncomingData(ulong clientId, ArraySegment<byte> payload, flo
         /// </summary>
         /// <param name="clientId">The ClientId to disconnect</param>
         public void DisconnectClient(ulong clientId)
+        {
+            DisconnectClient(clientId, null);
+        }
+
+        /// <summary>
+        /// Disconnects the remote client.
+        /// </summary>
+        /// <param name="clientId">The ClientId to disconnect</param>
+        /// <param name="reason">Disconnection reason. If set, client will receive a DisconnectReasonMessage and have the
+        /// reason available in the NetworkManager.DisconnectReason property</param>
+        public void DisconnectClient(ulong clientId, string reason)
         {
             if (!IsServer)
             {
                 throw new NotServerException($"Only server can disconnect remote clients. Please use `{nameof(Shutdown)}()` instead.");
             }
 
+            if (!string.IsNullOrEmpty(reason))
+            {
+                var disconnectReason = new DisconnectReasonMessage();
+                disconnectReason.Reason = reason;
+                SendMessage(ref disconnectReason, NetworkDelivery.Reliable, clientId);
+            }
+            MessagingSystem.ProcessSendQueues();
+
             OnClientDisconnectFromServer(clientId);
             DisconnectRemoteClient(clientId);
         }
@@ -2243,6 +2274,15 @@ internal void HandleConnectionApproval(ulong ownerClientId, ConnectionApprovalRe
             }
             else
             {
+                if (!string.IsNullOrEmpty(response.Reason))
+                {
+                    var disconnectReason = new DisconnectReasonMessage();
+                    disconnectReason.Reason = response.Reason;
+                    SendMessage(ref disconnectReason, NetworkDelivery.Reliable, ownerClientId);
+
+                    MessagingSystem.ProcessSendQueues();
+                }
+
                 PendingClients.Remove(ownerClientId);
                 DisconnectRemoteClient(ownerClientId);
             }

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

@@ -0,0 +1,38 @@
+namespace Unity.Netcode
+{
+    internal struct DisconnectReasonMessage : INetworkMessage
+    {
+        public string Reason;
+
+        public void Serialize(FastBufferWriter writer)
+        {
+            string reasonSent = Reason;
+            if (reasonSent == null)
+            {
+                reasonSent = string.Empty;
+            }
+
+            if (writer.TryBeginWrite(FastBufferWriter.GetWriteSize(reasonSent)))
+            {
+                writer.WriteValueSafe(reasonSent);
+            }
+            else
+            {
+                writer.WriteValueSafe(string.Empty);
+                NetworkLog.LogWarning(
+                    "Disconnect reason didn't fit. Disconnected without sending a reason. Consider shortening the reason string.");
+            }
+        }
+
+        public bool Deserialize(FastBufferReader reader, ref NetworkContext context)
+        {
+            reader.ReadValueSafe(out Reason);
+            return true;
+        }
+
+        public void Handle(ref NetworkContext context)
+        {
+            ((NetworkManager)context.SystemOwner).DisconnectReason = Reason;
+        }
+    };
+}

com.unity.netcode.gameobjects/Runtime/Messaging/DisconnectReasonMessage.cs.meta

@@ -0,0 +1,56 @@
+using NUnit.Framework;
+using Unity.Collections;
+
+namespace Unity.Netcode.EditorTests
+{
+    public class DisconnectMessageTests
+    {
+        [Test]
+        public void EmptyDisconnectReason()
+        {
+            var networkContext = new NetworkContext();
+            var writer = new FastBufferWriter(20, Allocator.Temp, 20);
+            var msg = new DisconnectReasonMessage();
+            msg.Reason = string.Empty;
+            msg.Serialize(writer);
+
+            var fbr = new FastBufferReader(writer, Allocator.Temp);
+            var recvMsg = new DisconnectReasonMessage();
+            recvMsg.Deserialize(fbr, ref networkContext);
+
+            Assert.IsEmpty(recvMsg.Reason);
+        }
+
+        [Test]
+        public void DisconnectReason()
+        {
+            var networkContext = new NetworkContext();
+            var writer = new FastBufferWriter(20, Allocator.Temp, 20);
+            var msg = new DisconnectReasonMessage();
+            msg.Reason = "Foo";
+            msg.Serialize(writer);
+
+            var fbr = new FastBufferReader(writer, Allocator.Temp);
+            var recvMsg = new DisconnectReasonMessage();
+            recvMsg.Deserialize(fbr, ref networkContext);
+
+            Assert.AreEqual("Foo", recvMsg.Reason);
+        }
+
+        [Test]
+        public void DisconnectReasonTooLong()
+        {
+            var networkContext = new NetworkContext();
+            var writer = new FastBufferWriter(20, Allocator.Temp, 20);
+            var msg = new DisconnectReasonMessage();
+            msg.Reason = "ThisStringIsWayLongerThanTwentyBytes";
+            msg.Serialize(writer);
+
+            var fbr = new FastBufferReader(writer, Allocator.Temp);
+            var recvMsg = new DisconnectReasonMessage();
+            recvMsg.Deserialize(fbr, ref networkContext);
+
+            Assert.IsEmpty(recvMsg.Reason);
+        }
+    }
+}

com.unity.netcode.gameobjects/Tests/Editor/DisconnectMessageTests.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections;
 using System.Text.RegularExpressions;
+using NUnit.Framework;
 using UnityEngine;
 using UnityEngine.TestTools;
 using Unity.Netcode.TestHelpers.Runtime;
@@ -25,16 +26,48 @@ protected override void OnServerAndClientsCreated()
         }
 
         private int m_DisconnectCount;
+        private bool m_ThrowOnDisconnect = false;
 
         public void OnClientDisconnectCallback(ulong clientId)
         {
             m_DisconnectCount++;
-            throw new SystemException("whatever");
+            if (m_ThrowOnDisconnect)
+            {
+                throw new SystemException("whatever");
+            }
+        }
+
+        [UnityTest]
+        public IEnumerator DisconnectReasonTest()
+        {
+            float startTime = Time.realtimeSinceStartup;
+            m_ThrowOnDisconnect = false;
+            m_DisconnectCount = 0;
+
+            // Add a callback for both clients, when they get disconnected
+            m_ClientNetworkManagers[0].OnClientDisconnectCallback += OnClientDisconnectCallback;
+            m_ClientNetworkManagers[1].OnClientDisconnectCallback += OnClientDisconnectCallback;
+
+            // Disconnect both clients, from the server
+            m_ServerNetworkManager.DisconnectClient(m_ClientNetworkManagers[0].LocalClientId, "Bogus reason 1");
+            m_ServerNetworkManager.DisconnectClient(m_ClientNetworkManagers[1].LocalClientId, "Bogus reason 2");
+
+            while (m_DisconnectCount < 2 && Time.realtimeSinceStartup < startTime + 10.0f)
+            {
+                yield return null;
+            }
+
+            Assert.AreEqual(m_ClientNetworkManagers[0].DisconnectReason, "Bogus reason 1");
+            Assert.AreEqual(m_ClientNetworkManagers[1].DisconnectReason, "Bogus reason 2");
+
+            Debug.Assert(m_DisconnectCount == 2);
         }
 
         [UnityTest]
         public IEnumerator DisconnectExceptionTest()
         {
+            m_ThrowOnDisconnect = true;
+            m_DisconnectCount = 0;
             float startTime = Time.realtimeSinceStartup;
 
             // Add a callback for first client, when they get disconnected

com.unity.netcode.gameobjects/Tests/Editor/DisconnectMessageTests.cs.meta

@@ -171,6 +171,11 @@ private IEnumerator ConnectionApprovalHandler(int numClients, int failureTestCou
                 }
             }
 
+            foreach (var c in clientsToClean)
+            {
+                Assert.AreEqual(c.DisconnectReason, "Some valid reason");
+            }
+
             foreach (var client in clients)
             {
                 // If a client failed, then it will already be shutdown
@@ -228,6 +233,14 @@ private void ConnectionApprovalCallback(NetworkManager.ConnectionApprovalRequest
                 response.Rotation = null;
                 response.PlayerPrefabHash = m_PrefabOverrideGlobalObjectIdHash;
             }
+            if (!response.Approved)
+            {
+                response.Reason = "Some valid reason";
+            }
+            else
+            {
+                response.Reason = string.Empty;
+            }
         }