fix: Add size validation to named and unnamed message sending [Up Port] (#3049)

* fix

Adding the fixes from #3043

* test

Adding the tests from #3043

* update

Adding changelog entry

* fix

Check for CustomMessageManager existing before invoking.

Author: NoelStephensUnity

com.unity.netcode.gameobjects/CHANGELOG.md

@@ -10,6 +10,7 @@ Additional documentation and release notes are available at [Multiplayer Documen
 
 ### Added
 
+- Added message size validation to named and unnamed message sending functions for better error messages. (#3049)
 - Added "Check for NetworkObject Component" property to the Multiplayer->Netcode for GameObjects project settings. When disabled, this will bypass the in-editor `NetworkObject` check on `NetworkBehaviour` components. (#3031)
 - Added `NetworkTransform.SwitchTransformSpaceWhenParented` property that, when enabled, will handle the world to local, local to world, and local to local transform space transitions when interpolation is enabled. (#3013)
 - Added `NetworkTransform.TickSyncChildren` that, when enabled, will tick synchronize nested and/or child `NetworkTransform` components to eliminate any potential visual jittering that could occur if the `NetworkTransform` instances get into a state where their state updates are landing on different network ticks. (#3013)

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

@@ -95,6 +95,7 @@ public void SendUnnamedMessage(IReadOnlyList<ulong> clientIds, FastBufferWriter
                 return;
             }
 
+            ValidateMessageSize(messageBuffer, networkDelivery, isNamed: false);
 
             if (m_NetworkManager.IsHost)
             {
@@ -131,6 +132,8 @@ public void SendUnnamedMessage(IReadOnlyList<ulong> clientIds, FastBufferWriter
         /// <param name="networkDelivery">The delivery type (QoS) to send data with</param>
         public void SendUnnamedMessage(ulong clientId, FastBufferWriter messageBuffer, NetworkDelivery networkDelivery = NetworkDelivery.ReliableSequenced)
         {
+            ValidateMessageSize(messageBuffer, networkDelivery, isNamed: false);
+
             if (m_NetworkManager.IsHost)
             {
                 if (clientId == m_NetworkManager.LocalClientId)
@@ -286,6 +289,8 @@ public void SendNamedMessageToAll(string messageName, FastBufferWriter messageSt
         /// <param name="networkDelivery">The delivery type (QoS) to send data with</param>
         public void SendNamedMessage(string messageName, ulong clientId, FastBufferWriter messageStream, NetworkDelivery networkDelivery = NetworkDelivery.ReliableSequenced)
         {
+            ValidateMessageSize(messageStream, networkDelivery, isNamed: true);
+
             ulong hash = 0;
             switch (m_NetworkManager.NetworkConfig.RpcHashSize)
             {
@@ -367,6 +372,8 @@ public void SendNamedMessage(string messageName, IReadOnlyList<ulong> clientIds,
                 return;
             }
 
+            ValidateMessageSize(messageStream, networkDelivery, isNamed: true);
+
             ulong hash = 0;
             switch (m_NetworkManager.NetworkConfig.RpcHashSize)
             {
@@ -405,5 +412,32 @@ public void SendNamedMessage(string messageName, IReadOnlyList<ulong> clientIds,
                 m_NetworkManager.NetworkMetrics.TrackNamedMessageSent(clientIds, messageName, size);
             }
         }
+
+        /// <summary>
+        /// Validate the size of the message. If it's a non-fragmented delivery type the message must fit within the
+        /// max allowed size with headers also subtracted. Named messages also include the hash
+        /// of the name string. Only validates in editor and development builds.
+        /// </summary>
+        /// <param name="messageStream">The named message payload</param>
+        /// <param name="networkDelivery">Delivery method</param>
+        /// <param name="isNamed">Is the message named (or unnamed)</param>
+        /// <exception cref="OverflowException">Exception thrown in case validation fails</exception>
+        private unsafe void ValidateMessageSize(FastBufferWriter messageStream, NetworkDelivery networkDelivery, bool isNamed)
+        {
+#if DEVELOPMENT_BUILD || UNITY_EDITOR
+            var maxNonFragmentedSize = m_NetworkManager.MessageManager.NonFragmentedMessageMaxSize - FastBufferWriter.GetWriteSize<NetworkMessageHeader>() - sizeof(NetworkBatchHeader);
+            if (isNamed)
+            {
+                maxNonFragmentedSize -= sizeof(ulong); // MessageName hash
+            }
+            if (networkDelivery != NetworkDelivery.ReliableFragmentedSequenced
+                && messageStream.Length > maxNonFragmentedSize)
+            {
+                throw new OverflowException($"Given message size ({messageStream.Length} bytes) is greater than " +
+                    $"the maximum allowed for the selected delivery method ({maxNonFragmentedSize} bytes). Try using " +
+                    $"ReliableFragmentedSequenced delivery method instead.");
+            }
+#endif
+        }
     }
 }

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

@@ -20,7 +20,7 @@ public bool Deserialize(FastBufferReader reader, ref NetworkContext context, int
 
         public void Handle(ref NetworkContext context)
         {
-            ((NetworkManager)context.SystemOwner).CustomMessagingManager.InvokeUnnamedMessage(context.SenderId, m_ReceivedData, context.SerializedHeaderSize);
+            ((NetworkManager)context.SystemOwner).CustomMessagingManager?.InvokeUnnamedMessage(context.SenderId, m_ReceivedData, context.SerializedHeaderSize);
         }
     }
 }

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

@@ -733,7 +733,11 @@ internal unsafe int SendPreSerializedMessage<TMessageType>(in FastBufferWriter t
                 }
 
                 ref var writeQueueItem = ref sendQueueItem.ElementAt(sendQueueItem.Length - 1);
-                writeQueueItem.Writer.TryBeginWrite(tmpSerializer.Length + headerSerializer.Length);
+                if (!writeQueueItem.Writer.TryBeginWrite(tmpSerializer.Length + headerSerializer.Length))
+                {
+                    Debug.LogError($"Not enough space to write message, size={tmpSerializer.Length + headerSerializer.Length} space used={writeQueueItem.Writer.Position} total size={writeQueueItem.Writer.Capacity}");
+                    continue;
+                }
 
                 writeQueueItem.Writer.WriteBytes(headerSerializer.GetUnsafePtr(), headerSerializer.Length);
                 writeQueueItem.Writer.WriteBytes(tmpSerializer.GetUnsafePtr(), tmpSerializer.Length);

com.unity.netcode.gameobjects/Runtime/Serialization/FastBufferWriter.cs

@@ -700,7 +700,7 @@ public unsafe void WriteBytes(byte* value, int size, int offset = 0)
             }
             if (Handle->Position + size > Handle->AllowedWriteMark)
             {
-                throw new OverflowException($"Attempted to write without first calling {nameof(TryBeginWrite)}()");
+                throw new OverflowException($"Attempted to write without first calling {nameof(TryBeginWrite)}(), Position+Size={Handle->Position + size} > AllowedWriteMark={Handle->AllowedWriteMark}");
             }
 #endif
             UnsafeUtility.MemCpy((Handle->BufferPointer + Handle->Position), value + offset, size);
@@ -729,7 +729,7 @@ public unsafe void WriteBytesSafe(byte* value, int size, int offset = 0)
 
             if (!TryBeginWriteInternal(size))
             {
-                throw new OverflowException("Writing past the end of the buffer");
+                throw new OverflowException($"Writing past the end of the buffer, size is {size} bytes but remaining capacity is {Handle->Capacity - Handle->Position} bytes");
             }
             UnsafeUtility.MemCpy((Handle->BufferPointer + Handle->Position), value + offset, size);
             Handle->Position += size;

com.unity.netcode.gameobjects/Tests/Runtime/Messaging/NamedMessageTests.cs

@@ -239,5 +239,45 @@ public void WhenSendingNamedMessageToNullClientList_ArgumentNullExceptionIsThrow
                     });
             }
         }
+
+        [Test]
+        public unsafe void ErrorMessageIsPrintedWhenAttemptingToSendNamedMessageWithTooBigBuffer()
+        {
+            // First try a valid send with the maximum allowed size (this is atm 1264)
+            var msgSize = m_ServerNetworkManager.MessageManager.NonFragmentedMessageMaxSize - FastBufferWriter.GetWriteSize<NetworkMessageHeader>() - sizeof(ulong)/*MessageName hash*/ - sizeof(NetworkBatchHeader);
+            var bufferSize = m_ServerNetworkManager.MessageManager.NonFragmentedMessageMaxSize;
+            var messageName = Guid.NewGuid().ToString();
+            var messageContent = new byte[msgSize];
+            var writer = new FastBufferWriter(bufferSize, Allocator.Temp, bufferSize * 2);
+            using (writer)
+            {
+                writer.TryBeginWrite(msgSize);
+                writer.WriteBytes(messageContent, msgSize, 0);
+                m_ServerNetworkManager.CustomMessagingManager.SendNamedMessage(messageName, new List<ulong> { FirstClient.LocalClientId }, writer);
+                m_ServerNetworkManager.CustomMessagingManager.SendNamedMessage(messageName, FirstClient.LocalClientId, writer);
+            }
+
+            msgSize++;
+            messageContent = new byte[msgSize];
+            writer = new FastBufferWriter(bufferSize, Allocator.Temp, bufferSize * 2);
+            using (writer)
+            {
+                writer.TryBeginWrite(msgSize);
+                writer.WriteBytes(messageContent, msgSize, 0);
+                var message = Assert.Throws<OverflowException>(
+                    () =>
+                    {
+                        m_ServerNetworkManager.CustomMessagingManager.SendNamedMessage(messageName, new List<ulong> { FirstClient.LocalClientId }, writer);
+                    }).Message;
+                Assert.IsTrue(message.Contains($"Given message size ({msgSize} bytes) is greater than the maximum"), $"Unexpected exception: {message}");
+
+                message = Assert.Throws<OverflowException>(
+                    () =>
+                    {
+                        m_ServerNetworkManager.CustomMessagingManager.SendNamedMessage(messageName, FirstClient.LocalClientId, writer);
+                    }).Message;
+                Assert.IsTrue(message.Contains($"Given message size ({msgSize} bytes) is greater than the maximum"), $"Unexpected exception: {message}");
+            }
+        }
     }
 }

com.unity.netcode.gameobjects/Tests/Runtime/Messaging/UnnamedMessageTests.cs

@@ -194,5 +194,44 @@ public void WhenSendingNamedMessageToNullClientList_ArgumentNullExceptionIsThrow
                     });
             }
         }
+
+        [Test]
+        public unsafe void ErrorMessageIsPrintedWhenAttemptingToSendUnnamedMessageWithTooBigBuffer()
+        {
+            // First try a valid send with the maximum allowed size (this is atm 1272)
+            var msgSize = m_ServerNetworkManager.MessageManager.NonFragmentedMessageMaxSize - FastBufferWriter.GetWriteSize<NetworkMessageHeader>() - sizeof(NetworkBatchHeader);
+            var bufferSize = m_ServerNetworkManager.MessageManager.NonFragmentedMessageMaxSize;
+            var messageContent = new byte[msgSize];
+            var writer = new FastBufferWriter(bufferSize, Allocator.Temp, bufferSize * 2);
+            using (writer)
+            {
+                writer.TryBeginWrite(msgSize);
+                writer.WriteBytes(messageContent, msgSize, 0);
+                m_ServerNetworkManager.CustomMessagingManager.SendUnnamedMessage(new List<ulong> { FirstClient.LocalClientId }, writer);
+                m_ServerNetworkManager.CustomMessagingManager.SendUnnamedMessage(FirstClient.LocalClientId, writer);
+            }
+
+            msgSize++;
+            messageContent = new byte[msgSize];
+            writer = new FastBufferWriter(bufferSize, Allocator.Temp, bufferSize * 2);
+            using (writer)
+            {
+                writer.TryBeginWrite(msgSize);
+                writer.WriteBytes(messageContent, msgSize, 0);
+                var message = Assert.Throws<OverflowException>(
+                    () =>
+                    {
+                        m_ServerNetworkManager.CustomMessagingManager.SendUnnamedMessage(new List<ulong> { FirstClient.LocalClientId }, writer);
+                    }).Message;
+                Assert.IsTrue(message.Contains($"Given message size ({msgSize} bytes) is greater than the maximum"), $"Unexpected exception: {message}");
+
+                message = Assert.Throws<OverflowException>(
+                    () =>
+                    {
+                        m_ServerNetworkManager.CustomMessagingManager.SendUnnamedMessage(FirstClient.LocalClientId, writer);
+                    }).Message;
+                Assert.IsTrue(message.Contains($"Given message size ({msgSize} bytes) is greater than the maximum"), $"Unexpected exception: {message}");
+            }
+        }
     }
 }