fix: logging warning after disconnect instead of invoking other handler

handler might try access server/client after disconnect and throw exceptions
MirageNet · Aug 31, 2024 · 05a72fa · 05a72fa
1 parent c355245
commit 05a72fa
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 6 deletions.
diff --git a/Assets/Mirage/Runtime/NetworkClient.cs b/Assets/Mirage/Runtime/NetworkClient.cs
@@ -144,7 +144,7 @@ public void Connect(string address = null, ushort? port = null)
             var socket = SocketFactory.CreateClientSocket();
             var maxPacketSize = SocketFactory.MaxPacketSize;
             MessageHandler = new MessageHandler(World, DisconnectOnException, RethrowException);
-            var dataHandler = new DataHandler(MessageHandler);
+            var dataHandler = new DataHandler(this, MessageHandler);
             Metrics = EnablePeerMetrics ? new Metrics(MetricsSize) : null;
 
             var config = PeerConfig ?? new Config();
@@ -231,7 +231,7 @@ internal void ConnectHost(NetworkServer server, IDataHandler serverDataHandler)
 
             // create local connection objects and connect them
             MessageHandler = new MessageHandler(World, DisconnectOnException, RethrowException);
-            var dataHandler = new DataHandler(MessageHandler);
+            var dataHandler = new DataHandler(this, MessageHandler);
             (var clientConn, var serverConn) = PipePeerConnection.Create(dataHandler, serverDataHandler, OnHostDisconnected, null);
 
             // set up client before connecting to server, server could invoke handlers
@@ -419,10 +419,12 @@ internal class DataHandler : IDataHandler
         {
             private IConnection _connection;
             private INetworkPlayer _player;
+            private readonly NetworkClient _client;
             private readonly IMessageReceiver _messageHandler;
 
-            public DataHandler(IMessageReceiver messageHandler)
+            public DataHandler(NetworkClient networkClient, IMessageReceiver messageHandler)
             {
+                _client = networkClient;
                 _messageHandler = messageHandler;
             }
 
@@ -434,6 +436,11 @@ public void SetConnection(IConnection connection, INetworkPlayer player)
 
             public void ReceiveMessage(IConnection connection, ArraySegment<byte> message)
             {
+                if (!_client.Active)
+                {
+                    if (logger.WarnEnabled()) logger.LogWarning("Received message after disconnect");
+                    return;
+                }
                 logger.Assert(_connection == connection);
                 _messageHandler.HandleMessage(_player, message);
             }

diff --git a/Assets/Mirage/Runtime/NetworkServer.cs b/Assets/Mirage/Runtime/NetworkServer.cs
@@ -251,7 +251,7 @@ public void StartServer(NetworkClient localClient = null)
             // create after MessageHandler, SyncVarReceiver uses it 
             _syncVarReceiver = new SyncVarReceiver(this, World);
 
-            var dataHandler = new DataHandler(MessageHandler, _connections);
+            var dataHandler = new DataHandler(this, MessageHandler, _connections);
             Metrics = EnablePeerMetrics ? new Metrics(MetricsSize) : null;
 
             var config = PeerConfig;
@@ -478,7 +478,7 @@ internal void AddLocalConnection(NetworkClient client, IConnection connection)
 
         [Obsolete("Use SendToAll(msg, authenticatedOnly, excludeLocalPlayer, channelId) instead")]
         public void SendToAll<T>(T msg, bool excludeLocalPlayer, Channel channelId = Channel.Reliable) => SendToAll(msg, authenticatedOnly: false, excludeLocalPlayer, channelId);
-        
+
         public void SendToAll<T>(T msg, bool authenticatedOnly, bool excludeLocalPlayer, Channel channelId = Channel.Reliable)
         {
             if (authenticatedOnly)
@@ -614,17 +614,25 @@ private void OnDisconnected(INetworkPlayer player)
         /// </summary>
         private sealed class DataHandler : IDataHandler
         {
+            private readonly NetworkServer _server;
             private readonly IMessageReceiver _messageHandler;
             private readonly Dictionary<IConnection, INetworkPlayer> _players;
 
-            public DataHandler(IMessageReceiver messageHandler, Dictionary<IConnection, INetworkPlayer> connections)
+            public DataHandler(NetworkServer server, IMessageReceiver messageHandler, Dictionary<IConnection, INetworkPlayer> connections)
             {
+                _server = server;
                 _messageHandler = messageHandler;
                 _players = connections;
             }
 
             public void ReceiveMessage(IConnection connection, ArraySegment<byte> message)
             {
+                if (!_server.Active)
+                {
+                    if (logger.WarnEnabled()) logger.LogWarning("Received message after disconnect");
+                    return;
+                }
+
                 if (_players.TryGetValue(connection, out var player))
                 {
                     _messageHandler.HandleMessage(player, message);