fix: fixing use of network identity in network message (#955)

adding objectLocator to object locator so that it can be used to find objects when reading network messages

Assets/Mirage/Runtime/MessageHandler.cs

@@ -12,6 +12,7 @@ public class MessageHandler : IMessageReceiver
         static readonly ILogger logger = LogFactory.GetLogger(typeof(MessageHandler));
 
         readonly bool disconnectOnException;
+        readonly IObjectLocator objectLocator;
 
         /// <summary>
         /// Handles network messages on client and server
@@ -22,9 +23,10 @@ public class MessageHandler : IMessageReceiver
 
         internal readonly Dictionary<int, NetworkMessageDelegate> messageHandlers = new Dictionary<int, NetworkMessageDelegate>();
 
-        public MessageHandler(bool disconnectOnException)
+        public MessageHandler(IObjectLocator objectLocator, bool disconnectOnException)
         {
             this.disconnectOnException = disconnectOnException;
+            this.objectLocator = objectLocator;
         }
 
         private static NetworkMessageDelegate MessageWrapper<T>(MessageDelegateWithPlayer<T> handler)
@@ -109,6 +111,9 @@ public void HandleMessage(INetworkPlayer player, ArraySegment<byte> packet)
         {
             using (PooledNetworkReader networkReader = NetworkReaderPool.GetReader(packet))
             {
+                // set ObjectLocator so that message can use NetworkIdentity
+                networkReader.ObjectLocator = objectLocator;
+
                 // protect against attackers trying to send invalid data packets
                 // exception could be throw if:
                 // - invalid headers

Assets/Mirage/Runtime/NetworkClient.cs

@@ -114,7 +114,7 @@ public void Connect(string address = null, ushort? port = null)
             if (logger.LogEnabled()) logger.Log($"Client connecting to endpoint: {endPoint}");
 
             ISocket socket = SocketFactory.CreateClientSocket();
-            MessageHandler = new MessageHandler(DisconnectOnException);
+            MessageHandler = new MessageHandler(World, DisconnectOnException);
             var dataHandler = new DataHandler(MessageHandler);
             Metrics = EnablePeerMetrics ? new Metrics() : null;
 
@@ -190,7 +190,7 @@ internal void ConnectHost(NetworkServer server, IDataHandler serverDataHandler)
             World = server.World;
 
             // create local connection objects and connect them
-            MessageHandler = new MessageHandler(DisconnectOnException);
+            MessageHandler = new MessageHandler(World, DisconnectOnException);
             var dataHandler = new DataHandler(MessageHandler);
             (IConnection clientConn, IConnection serverConn) = PipePeerConnection.Create(dataHandler, serverDataHandler, OnHostDisconnected, null);
 

Assets/Mirage/Runtime/NetworkServer.cs

@@ -192,7 +192,7 @@ public void StartServer(NetworkClient localClient = null)
             SyncVarSender = new SyncVarSender();
 
             LocalClient = localClient;
-            MessageHandler = new MessageHandler(DisconnectOnException);
+            MessageHandler = new MessageHandler(World, DisconnectOnException);
             MessageHandler.RegisterHandler<NetworkPingMessage>(World.Time.OnServerPing);
 
             ISocket socket = SocketFactory.CreateServerSocket();

Assets/Tests/Runtime/ClientServer/NetworkIdentityMessageTest.cs

@@ -0,0 +1,29 @@
+using NUnit.Framework;
+
+namespace Mirage.Tests.Runtime.ClientServer
+{
+    public class NetworkIdentityMessageTest : ClientServerSetup<MockComponent>
+    {
+        [NetworkMessage]
+        public struct MyMessage
+        {
+            public NetworkIdentity player1;
+        }
+
+        [Test]
+        public void MessageFindsNetworkIdentities()
+        {
+            NetworkIdentity found = null;
+            client.MessageHandler.RegisterHandler((MyMessage msg) =>
+            {
+                found = msg.player1;
+            });
+            serverPlayer.Send(new MyMessage { player1 = serverPlayer.Identity });
+
+            server.Update();
+            client.Update();
+
+            Assert.That(found == clientPlayer.Identity, "Could not find client version of object");
+        }
+    }
+}

Assets/Tests/Runtime/Serialization/MessageHandlerTest.cs

@@ -23,7 +23,7 @@ public void SetUp()
             reader = new NetworkReader();
             reader.Reset(new byte[] { 1, 2, 3, 4 });
 
-            messageHandler = new MessageHandler(true);
+            messageHandler = new MessageHandler(null, true);
         }
 
 
@@ -44,7 +44,7 @@ public void InvokesMessageHandler()
         [TestCase(false)]
         public void DisconnectsIfHandlerHasException(bool disconnectOnThrow)
         {
-            messageHandler = new MessageHandler(disconnectOnThrow);
+            messageHandler = new MessageHandler(null, disconnectOnThrow);
 
             int invoked = 0;
             messageHandler.RegisterHandler<SceneReadyMessage>(_ => { invoked++; throw new InvalidOperationException("Fun Exception"); });