refactor: Part 2 of INetworkMessage ILPP refactor (#1282)

Moves all aspects of message registration go ILPP and replaces attribute inspection with an IMessageProvider class that allows tests to change the discovery from the automatic ILPP discovery to returning a defined set of messages for the test.

Author: ShadauxCat

com.unity.netcode.gameobjects/Editor/CodeGen/CodeGenHelpers.cs

@@ -23,8 +23,6 @@ internal static class CodeGenHelpers
         public static readonly string ServerRpcParams_FullName = typeof(ServerRpcParams).FullName;
         public static readonly string ClientRpcParams_FullName = typeof(ClientRpcParams).FullName;
         public static readonly string INetworkSerializable_FullName = typeof(INetworkSerializable).FullName;
-        public static readonly string INetworkSerializable_NetworkSerialize_Name = nameof(INetworkSerializable.NetworkSerialize);
-        public static readonly string IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute_FullName = typeof(IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute).FullName;
         public static readonly string UnityColor_FullName = typeof(Color).FullName;
         public static readonly string UnityColor32_FullName = typeof(Color32).FullName;
         public static readonly string UnityVector2_FullName = typeof(Vector2).FullName;

com.unity.netcode.gameobjects/Editor/CodeGen/INetworkMessageILPP.cs

@@ -59,7 +59,6 @@ public override ILPostProcessResult Process(ICompiledAssembly compiledAssembly)
 
                     try
                     {
-                        types.ForEach(b => ProcessINetworkMessage(b));
                         CreateModuleInitializer(assemblyDefinition, types);
                     }
                     catch (Exception e)
@@ -98,7 +97,11 @@ public override ILPostProcessResult Process(ICompiledAssembly compiledAssembly)
 
         private TypeReference m_FastBufferReader_TypeRef;
         private TypeReference m_NetworkContext_TypeRef;
-        private FieldReference m_MessagingSystem___network_message_types_FieldRef;
+        private TypeReference m_MessagingSystem_MessageWithHandler_TypeRef;
+        private MethodReference m_MessagingSystem_MessageHandler_Constructor_TypeRef;
+        private FieldReference m_ILPPMessageProvider___network_message_types_FieldRef;
+        private FieldReference m_MessagingSystem_MessageWithHandler_MessageType_FieldRef;
+        private FieldReference m_MessagingSystem_MessageWithHandler_Handler_FieldRef;
         private MethodReference m_Type_GetTypeFromHandle_MethodRef;
 
         private MethodReference m_List_Add_MethodRef;
@@ -107,6 +110,24 @@ private bool ImportReferences(ModuleDefinition moduleDefinition)
         {
             m_FastBufferReader_TypeRef = moduleDefinition.ImportReference(typeof(FastBufferReader));
             m_NetworkContext_TypeRef = moduleDefinition.ImportReference(typeof(NetworkContext));
+            m_MessagingSystem_MessageHandler_Constructor_TypeRef =
+                moduleDefinition.ImportReference(typeof(MessagingSystem.MessageHandler).GetConstructors()[0]);
+
+            var messageWithHandlerType = typeof(MessagingSystem.MessageWithHandler);
+            m_MessagingSystem_MessageWithHandler_TypeRef =
+                moduleDefinition.ImportReference(messageWithHandlerType);
+            foreach (var fieldInfo in messageWithHandlerType.GetFields())
+            {
+                switch (fieldInfo.Name)
+                {
+                    case nameof(MessagingSystem.MessageWithHandler.MessageType):
+                        m_MessagingSystem_MessageWithHandler_MessageType_FieldRef = moduleDefinition.ImportReference(fieldInfo);
+                        break;
+                    case nameof(MessagingSystem.MessageWithHandler.Handler):
+                        m_MessagingSystem_MessageWithHandler_Handler_FieldRef = moduleDefinition.ImportReference(fieldInfo);
+                        break;
+                }
+            }
 
             var typeType = typeof(Type);
             foreach (var methodInfo in typeType.GetMethods())
@@ -119,23 +140,23 @@ private bool ImportReferences(ModuleDefinition moduleDefinition)
                 }
             }
 
-            var messagingSystemType = typeof(MessagingSystem);
-            foreach (var fieldInfo in messagingSystemType.GetFields(BindingFlags.Static | BindingFlags.NonPublic))
+            var ilppMessageProviderType = typeof(ILPPMessageProvider);
+            foreach (var fieldInfo in ilppMessageProviderType.GetFields(BindingFlags.Static | BindingFlags.NonPublic))
             {
                 switch (fieldInfo.Name)
                 {
-                    case nameof(MessagingSystem.__network_message_types):
-                        m_MessagingSystem___network_message_types_FieldRef = moduleDefinition.ImportReference(fieldInfo);
+                    case nameof(ILPPMessageProvider.__network_message_types):
+                        m_ILPPMessageProvider___network_message_types_FieldRef = moduleDefinition.ImportReference(fieldInfo);
                         break;
                 }
             }
 
-            var listType = typeof(List<Type>);
+            var listType = typeof(List<MessagingSystem.MessageWithHandler>);
             foreach (var methodInfo in listType.GetMethods())
             {
                 switch (methodInfo.Name)
                 {
-                    case nameof(List<Type>.Add):
+                    case nameof(List<MessagingSystem.MessageWithHandler>.Add):
                         m_List_Add_MethodRef = moduleDefinition.ImportReference(methodInfo);
                         break;
                 }
@@ -145,9 +166,8 @@ private bool ImportReferences(ModuleDefinition moduleDefinition)
             return true;
         }
 
-        private void ProcessINetworkMessage(TypeDefinition typeDefinition)
+        private MethodReference GetNetworkMessageRecieveHandler(TypeDefinition typeDefinition)
         {
-            var foundAValidMethod = false;
             SequencePoint typeSequence = null;
             foreach (var method in typeDefinition.Methods)
             {
@@ -168,15 +188,12 @@ private void ProcessINetworkMessage(TypeDefinition typeDefinition)
                     && resolved.Parameters[1].ParameterType.GetElementType().Resolve() == m_NetworkContext_TypeRef.Resolve()
                     && resolved.ReturnType == resolved.Module.TypeSystem.Void)
                 {
-                    foundAValidMethod = true;
-                    break;
+                    return method;
                 }
             }
 
-            if (!foundAValidMethod)
-            {
-                m_Diagnostics.AddError(typeSequence, $"Class {typeDefinition.FullName} does not implement required function: `public static void Receive(FastBufferReader, in NetworkContext)`");
-            }
+            m_Diagnostics.AddError(typeSequence, $"Class {typeDefinition.FullName} does not implement required method: `public static void Receive(FastBufferReader, in NetworkContext)`");
+            return null;
         }
 
         private MethodDefinition GetOrCreateStaticConstructor(TypeDefinition typeDefinition)
@@ -198,12 +215,32 @@ private MethodDefinition GetOrCreateStaticConstructor(TypeDefinition typeDefinit
             return staticCtorMethodDef;
         }
 
-        private void CreateInstructionsToRegisterType(ILProcessor processor, List<Instruction> instructions, TypeReference type)
+        private void CreateInstructionsToRegisterType(ILProcessor processor, List<Instruction> instructions, TypeReference type, MethodReference receiveMethod)
         {
-            // MessagingSystem.__network_message_types.Add(typeof(type));
-            instructions.Add(processor.Create(OpCodes.Ldsfld, m_MessagingSystem___network_message_types_FieldRef));
+            // MessagingSystem.__network_message_types.Add(new MessagingSystem.MessageWithHandler{MessageType=typeof(type), Handler=type.Receive});
+            processor.Body.Variables.Add(new VariableDefinition(m_MessagingSystem_MessageWithHandler_TypeRef));
+            int messageWithHandlerLocIdx = processor.Body.Variables.Count - 1;
+
+            instructions.Add(processor.Create(OpCodes.Ldsfld, m_ILPPMessageProvider___network_message_types_FieldRef));
+            instructions.Add(processor.Create(OpCodes.Ldloca, messageWithHandlerLocIdx));
+            instructions.Add(processor.Create(OpCodes.Initobj, m_MessagingSystem_MessageWithHandler_TypeRef));
+
+            // tmp.MessageType = typeof(type);
+            instructions.Add(processor.Create(OpCodes.Ldloca, messageWithHandlerLocIdx));
             instructions.Add(processor.Create(OpCodes.Ldtoken, type));
             instructions.Add(processor.Create(OpCodes.Call, m_Type_GetTypeFromHandle_MethodRef));
+            instructions.Add(processor.Create(OpCodes.Stfld, m_MessagingSystem_MessageWithHandler_MessageType_FieldRef));
+
+            // tmp.Handler = type.Receive
+            instructions.Add(processor.Create(OpCodes.Ldloca, messageWithHandlerLocIdx));
+            instructions.Add(processor.Create(OpCodes.Ldnull));
+
+            instructions.Add(processor.Create(OpCodes.Ldftn, receiveMethod));
+            instructions.Add(processor.Create(OpCodes.Newobj, m_MessagingSystem_MessageHandler_Constructor_TypeRef));
+            instructions.Add(processor.Create(OpCodes.Stfld, m_MessagingSystem_MessageWithHandler_Handler_FieldRef));
+
+            // ILPPMessageProvider.__network_message_types.Add(tmp);
+            instructions.Add(processor.Create(OpCodes.Ldloc, messageWithHandlerLocIdx));
             instructions.Add(processor.Create(OpCodes.Callvirt, m_List_Add_MethodRef));
         }
 
@@ -227,7 +264,12 @@ private void CreateModuleInitializer(AssemblyDefinition assembly, List<TypeDefin
 
                     foreach (var type in networkMessageTypes)
                     {
-                        CreateInstructionsToRegisterType(processor, instructions, type);
+                        var receiveMethod = GetNetworkMessageRecieveHandler(type);
+                        if (receiveMethod == null)
+                        {
+                            continue;
+                        }
+                        CreateInstructionsToRegisterType(processor, instructions, type, receiveMethod);
                     }
 
                     instructions.ForEach(instruction => processor.Body.Instructions.Insert(processor.Body.Instructions.Count - 1, instruction));

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

@@ -241,11 +241,7 @@ public GameObject GetNetworkPrefabOverride(GameObject gameObject)
         public ulong LocalClientId
         {
             get => IsServer ? NetworkConfig.NetworkTransport.ServerClientId : m_LocalClientId;
-            internal set
-            {
-                m_LocalClientId = value;
-                MessagingSystem.SetLocalClientId(value);
-            }
+            internal set => m_LocalClientId = value;
         }
 
         private ulong m_LocalClientId;
@@ -497,7 +493,7 @@ private void Initialize(bool server)
             this.RegisterNetworkUpdate(NetworkUpdateStage.EarlyUpdate);
             this.RegisterNetworkUpdate(NetworkUpdateStage.PostLateUpdate);
 
-            MessagingSystem = new MessagingSystem(new NetworkManagerMessageSender(this), this, ulong.MaxValue);
+            MessagingSystem = new MessagingSystem(new NetworkManagerMessageSender(this), this);
 
             MessagingSystem.Hook(new NetworkManagerHooks(this));
 #if DEVELOPMENT_BUILD || UNITY_EDITOR

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

@@ -0,0 +1,17 @@
+using System.Collections.Generic;
+
+namespace Unity.Netcode
+{
+    internal struct ILPPMessageProvider : IMessageProvider
+    {
+#pragma warning disable IDE1006 // disable naming rule violation check
+        // This is NOT modified by RuntimeAccessModifiersILPP right now, but is populated by ILPP.
+        internal static readonly List<MessagingSystem.MessageWithHandler> __network_message_types = new List<MessagingSystem.MessageWithHandler>();
+#pragma warning restore IDE1006 // restore naming rule violation check
+
+        public List<MessagingSystem.MessageWithHandler> GetMessages()
+        {
+            return __network_message_types;
+        }
+    }
+}

com.unity.netcode.gameobjects/Runtime/Messaging/IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute.cs.meta renamed to com.unity.netcode.gameobjects/Runtime/Messaging/ILPPMessageProvider.cs.meta

@@ -0,0 +1,9 @@
+using System.Collections.Generic;
+
+namespace Unity.Netcode
+{
+    internal interface IMessageProvider
+    {
+        List<MessagingSystem.MessageWithHandler> GetMessages();
+    }
+}

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

@@ -17,13 +17,6 @@ public InvalidMessageStructureException(string issue) : base(issue) { }
 
     internal class MessagingSystem : IDisposable
     {
-
-
-#pragma warning disable IDE1006 // disable naming rule violation check
-        // This is NOT modified by RuntimeAccessModifiersILPP right now, but is populated by ILPP.
-        internal static readonly List<Type> __network_message_types = new List<Type>();
-#pragma warning restore IDE1006 // restore naming rule violation check
-
         private struct ReceiveQueueItem
         {
             public FastBufferReader Reader;
@@ -61,7 +54,6 @@ public SendQueueItem(NetworkDelivery delivery, int writerSize, Allocator writerA
         private byte m_HighMessageType;
         private object m_Owner;
         private IMessageSender m_MessageSender;
-        private ulong m_LocalClientId;
         private bool m_Disposed;
 
         internal Type[] MessageTypes => m_ReverseTypeMap;
@@ -76,50 +68,26 @@ internal byte GetMessageType(Type t)
         public const int NON_FRAGMENTED_MESSAGE_MAX_SIZE = 1300;
         public const int FRAGMENTED_MESSAGE_MAX_SIZE = 64000;
 
-        public MessagingSystem(IMessageSender messageSender, object owner, ulong localClientId = long.MaxValue)
+        internal struct MessageWithHandler
+        {
+            public Type MessageType;
+            public MessageHandler Handler;
+        }
+
+        public MessagingSystem(IMessageSender messageSender, object owner, IMessageProvider provider = null)
         {
             try
             {
-                m_LocalClientId = localClientId;
                 m_MessageSender = messageSender;
                 m_Owner = owner;
 
-                var allowedTypes = new List<Type>();
-                foreach (var type in __network_message_types)
+                if (provider == null)
                 {
-                    var attributes = type.GetCustomAttributes(typeof(IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute), false);
-                    // If [IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute(ownerType)] isn't provided, it defaults
-                    // to being bound to NetworkManager. This is technically a breach of domain by having
-                    // MessagingSystem know about the existence of NetworkManager... but ultimately,
-                    // IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute is provided to support testing, not to support
-                    // general use of MessagingSystem outside of Netcode for GameObjects, so having MessagingSystem
-                    // know about NetworkManager isn't so bad. Especially since it's just a default value.
-                    // This is just a convenience to keep us and our users from having to use
-                    // [Bind(typeof(NetworkManager))] on every message - only tests that don't want to use
-                    // the full NetworkManager need to worry about it.
-                    var shouldSkip = attributes.Length != 0 || !(m_Owner is NetworkManager);
-                    for (var i = 0; i < attributes.Length; ++i)
-                    {
-                        var bindAttribute = (IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute)attributes[i];
-                        if (
-                            (bindAttribute.BoundType != null &&
-                             bindAttribute.BoundType.IsInstanceOfType(m_Owner)) ||
-                            (m_Owner == null && bindAttribute.BoundType == null))
-                        {
-                            shouldSkip = false;
-                            break;
-                        }
-                    }
-
-                    if (shouldSkip)
-                    {
-                        continue;
-                    }
-
-                    allowedTypes.Add(type);
+                    provider = new ILPPMessageProvider();
                 }
+                var allowedTypes = provider.GetMessages();
 
-                allowedTypes.Sort((a, b) => string.CompareOrdinal(a.FullName, b.FullName));
+                allowedTypes.Sort((a, b) => string.CompareOrdinal(a.MessageType.FullName, b.MessageType.FullName));
                 foreach (var type in allowedTypes)
                 {
                     RegisterMessageType(type);
@@ -154,40 +122,16 @@ public void Dispose()
             Dispose();
         }
 
-        public void SetLocalClientId(ulong localClientId)
-        {
-            m_LocalClientId = localClientId;
-        }
-
         public void Hook(INetworkHooks hooks)
         {
             m_Hooks.Add(hooks);
         }
 
-        private void RegisterMessageType(Type messageType)
+        private void RegisterMessageType(MessageWithHandler messageWithHandler)
         {
-            if (!typeof(INetworkMessage).IsAssignableFrom(messageType))
-            {
-                throw new ArgumentException("RegisterMessageType types must be INetworkMessage types.");
-            }
-
-            var method = messageType.GetMethod("Receive");
-            if (method == null)
-            {
-                throw new InvalidMessageStructureException(
-                    $"{messageType.FullName}: All INetworkMessage types must implement public static void Receive(FastBufferReader reader, in NetworkContext context)");
-            }
-
-            var asDelegate = Delegate.CreateDelegate(typeof(MessageHandler), method, false);
-            if (asDelegate == null)
-            {
-                throw new InvalidMessageStructureException(
-                    $"{messageType.FullName}: All INetworkMessage types must implement public static void Receive(FastBufferReader reader, in NetworkContext context)");
-            }
-
-            m_MessageHandlers[m_HighMessageType] = (MessageHandler)asDelegate;
-            m_ReverseTypeMap[m_HighMessageType] = messageType;
-            m_MessageTypes[messageType] = m_HighMessageType++;
+            m_MessageHandlers[m_HighMessageType] = messageWithHandler.Handler;
+            m_ReverseTypeMap[m_HighMessageType] = messageWithHandler.MessageType;
+            m_MessageTypes[messageWithHandler.MessageType] = m_HighMessageType++;
         }
 
         internal void HandleIncomingData(ulong clientId, ArraySegment<byte> data, float receiveTime)