Merge remote-tracking branch 'xmanning/develop-2.0.0' into feat/server-enforced-RpcInvokePermission

EmandM · EmandM · commit 00fca38009fb · 2025-10-17T13:51:44.000-04:00
diff --git a/com.unity.netcode.gameobjects/Editor/CodeGen/NetworkBehaviourILPP.cs b/com.unity.netcode.gameobjects/Editor/CodeGen/NetworkBehaviourILPP.cs
@@ -609,6 +609,7 @@ private void CreateNetworkVariableTypeInitializers(AssemblyDefinition assembly,
         private const string k_NetworkVariableBase_Initialize = nameof(NetworkVariableBase.Initialize);
 
         private const string k_RpcAttribute_Delivery = nameof(RpcAttribute.Delivery);
+        private const string k_RpcAttribute_InvokePermission = nameof(RpcAttribute.InvokePermission);
         private const string k_ServerRpcAttribute_RequireOwnership = nameof(ServerRpcAttribute.RequireOwnership);
         private const string k_RpcParams_Server = nameof(__RpcParams.Server);
         private const string k_RpcParams_Client = nameof(__RpcParams.Client);
@@ -1311,7 +1312,7 @@ private void ProcessNetworkBehaviour(TypeDefinition typeDefinition, string[] ass
                     return;
                 }
             }
-            var rpcHandlers = new List<(uint RpcMethodId, MethodDefinition RpcHandler, string RpcMethodName)>();
+            var rpcHandlers = new List<(uint RpcMethodId, MethodDefinition RpcHandler, string RpcMethodName, CustomAttribute rpcAttribute)>();
 
             bool isEditorOrDevelopment = assemblyDefines.Contains("UNITY_EDITOR") || assemblyDefines.Contains("DEVELOPMENT_BUILD");
 
@@ -1342,7 +1343,7 @@ private void ProcessNetworkBehaviour(TypeDefinition typeDefinition, string[] ass
 
                 InjectWriteAndCallBlocks(methodDefinition, rpcAttribute, rpcMethodId);
 
-                rpcHandlers.Add((rpcMethodId, GenerateStaticHandler(methodDefinition, rpcAttribute, rpcMethodId), methodDefinition.Name));
+                rpcHandlers.Add((rpcMethodId, GenerateStaticHandler(methodDefinition, rpcAttribute, rpcMethodId), methodDefinition.Name, rpcAttribute));
             }
 
             GenerateVariableInitialization(typeDefinition);
@@ -1424,7 +1425,7 @@ private void ProcessNetworkBehaviour(TypeDefinition typeDefinition, string[] ass
                 var instructions = new List<Instruction>();
                 var processor = initializeRpcsMethodDef.Body.GetILProcessor();
 
-                foreach (var (rpcMethodId, rpcHandler, rpcMethodName) in rpcHandlers)
+                foreach (var (rpcMethodId, rpcHandler, rpcMethodName, rpcAttribute) in rpcHandlers)
                 {
                     typeDefinition.Methods.Add(rpcHandler);
 
@@ -1435,12 +1436,35 @@ private void ProcessNetworkBehaviour(TypeDefinition typeDefinition, string[] ass
                         callMethod = callMethod.MakeGeneric(genericTypes.ToArray());
                     }
 
-                    // __registerRpc(RpcMethodId, HandleFunc, methodName);
+                    var isClientRpc = rpcAttribute.AttributeType.FullName == CodeGenHelpers.ClientRpcAttribute_FullName;
+
+                    var invokePermission = RpcInvokePermission.Anyone;
+
+                    foreach (var attrField in rpcAttribute.Fields)
+                    {
+                        switch (attrField.Name)
+                        {
+                            case k_ServerRpcAttribute_RequireOwnership:
+                                invokePermission = (attrField.Argument.Type == rpcHandler.Module.TypeSystem.Boolean && (bool)attrField.Argument.Value) ? RpcInvokePermission.Owner : RpcInvokePermission.Anyone;
+                                break;
+                            case k_RpcAttribute_InvokePermission:
+                                invokePermission = (RpcInvokePermission)attrField.Argument.Value;
+                                break;
+                        }
+                    }
+
+                    if (isClientRpc)
+                    {
+                        invokePermission = RpcInvokePermission.Server;
+                    }
+
+                    // __registerRpc(RpcMethodId, HandleFunc, invokePermission, methodName);
                     instructions.Add(processor.Create(OpCodes.Ldarg_0));
                     instructions.Add(processor.Create(OpCodes.Ldc_I4, unchecked((int)rpcMethodId)));
                     instructions.Add(processor.Create(OpCodes.Ldnull));
                     instructions.Add(processor.Create(OpCodes.Ldftn, callMethod));
                     instructions.Add(processor.Create(OpCodes.Newobj, m_NetworkHandlerDelegateCtor_MethodRef));
+                    instructions.Add(processor.Create(OpCodes.Ldc_I4, (int)invokePermission));
                     instructions.Add(processor.Create(OpCodes.Ldstr, rpcMethodName));
                     instructions.Add(processor.Create(OpCodes.Call, m_NetworkBehaviour___registerRpc_MethodRef));
                 }
@@ -1517,6 +1541,7 @@ private void ProcessNetworkBehaviour(TypeDefinition typeDefinition, string[] ass
         private CustomAttribute CheckAndGetRpcAttribute(MethodDefinition methodDefinition)
         {
             CustomAttribute rpcAttribute = null;
+
             foreach (var customAttribute in methodDefinition.CustomAttributes)
             {
                 var customAttributeType_FullName = customAttribute.AttributeType.FullName;
@@ -1600,6 +1625,30 @@ private CustomAttribute CheckAndGetRpcAttribute(MethodDefinition methodDefinitio
 
                 return null;
             }
+
+            bool hasRequireOwnership = false, hasInvokePermission = false;
+
+            foreach (var argument in rpcAttribute.Fields)
+            {
+                switch (argument.Name)
+                {
+                    case k_ServerRpcAttribute_RequireOwnership:
+                        hasRequireOwnership = true;
+                        break;
+                    case k_RpcAttribute_InvokePermission:
+                        hasInvokePermission = true;
+                        break;
+                    default:
+                        break;
+                }
+            }
+
+            if (hasRequireOwnership && hasInvokePermission)
+            {
+                m_Diagnostics.AddError("Rpc attribute cannot declare both RequireOwnership and InvokePermission!");
+                return null;
+            }
+
             // Checks for IsSerializable are moved to later as the check is now done by dynamically seeing if any valid
             // serializer OR extension method exists for it.
             return rpcAttribute;
@@ -2346,6 +2395,7 @@ private void InjectWriteAndCallBlocks(MethodDefinition methodDefinition, CustomA
                             m_Diagnostics.AddError($"{nameof(RpcAttribute)} contains field {field} which is not present in {nameof(RpcAttribute.RpcAttributeParams)}.");
                         }
                     }
+
                     instructions.Add(processor.Create(OpCodes.Ldloc, rpcAttributeParamsIdx));
 
                     // defaultTarget
@@ -2845,19 +2895,28 @@ private MethodDefinition GenerateStaticHandler(MethodDefinition methodDefinition
             var processor = rpcHandler.Body.GetILProcessor();
 
             var isServerRpc = rpcAttribute.AttributeType.FullName == CodeGenHelpers.ServerRpcAttribute_FullName;
-            var isCientRpc = rpcAttribute.AttributeType.FullName == CodeGenHelpers.ClientRpcAttribute_FullName;
+            var isClientRpc = rpcAttribute.AttributeType.FullName == CodeGenHelpers.ClientRpcAttribute_FullName;
             var isGenericRpc = rpcAttribute.AttributeType.FullName == CodeGenHelpers.RpcAttribute_FullName;
-            var requireOwnership = true; // default value MUST be == `ServerRpcAttribute.RequireOwnership`
+            var invokePermission = RpcInvokePermission.Anyone;
             foreach (var attrField in rpcAttribute.Fields)
             {
                 switch (attrField.Name)
                 {
                     case k_ServerRpcAttribute_RequireOwnership:
-                        requireOwnership = attrField.Argument.Type == typeSystem.Boolean && (bool)attrField.Argument.Value;
+                        invokePermission = (attrField.Argument.Type == typeSystem.Boolean && (bool)attrField.Argument.Value) ? RpcInvokePermission.Owner : RpcInvokePermission.Anyone;
+                        break;
+                    case k_RpcAttribute_InvokePermission:
+                        invokePermission = (RpcInvokePermission)attrField.Argument.Value;
                         break;
                 }
             }
 
+            // legacy ClientRpc should always be RpcInvokePermission.Server
+            if (isClientRpc)
+            {
+                invokePermission = RpcInvokePermission.Server;
+            }
+
             rpcHandler.Body.InitLocals = true;
             // NetworkManager networkManager;
             rpcHandler.Body.Variables.Add(new VariableDefinition(m_NetworkManager_TypeRef));
@@ -2883,7 +2942,7 @@ private MethodDefinition GenerateStaticHandler(MethodDefinition methodDefinition
                 processor.Append(lastInstr);
             }
 
-            if (isServerRpc && requireOwnership)
+            if (isServerRpc && invokePermission == RpcInvokePermission.Owner)
             {
                 var roReturnInstr = processor.Create(OpCodes.Ret);
                 var roLastInstr = processor.Create(OpCodes.Nop);
@@ -2917,6 +2976,42 @@ private MethodDefinition GenerateStaticHandler(MethodDefinition methodDefinition
 
                 processor.Append(logNextInstr);
 
+                processor.Append(roReturnInstr);
+                processor.Append(roLastInstr);
+            } else if (invokePermission == RpcInvokePermission.Server)
+            {
+                var roReturnInstr = processor.Create(OpCodes.Ret);
+                var roLastInstr = processor.Create(OpCodes.Nop);
+
+                // if (rpcParams.Server.Receive.SenderClientId != NetworkManager.IsServer) { ... } return;
+                processor.Emit(OpCodes.Ldarg_2);
+                processor.Emit(OpCodes.Ldfld, m_RpcParams_Server_FieldRef);
+                processor.Emit(OpCodes.Ldfld, m_ServerRpcParams_Receive_FieldRef);
+                processor.Emit(OpCodes.Ldfld, m_ServerRpcParams_Receive_SenderClientId_FieldRef);
+                processor.Emit(OpCodes.Ldarg_0);
+                processor.Emit(OpCodes.Call, m_NetworkManager_getIsServer_MethodRef);
+                processor.Emit(OpCodes.Ceq);
+                processor.Emit(OpCodes.Ldc_I4, 0);
+                processor.Emit(OpCodes.Ceq);
+                processor.Emit(OpCodes.Brfalse, roLastInstr);
+
+                var logNextInstr = processor.Create(OpCodes.Nop);
+
+                // if (LogLevel.Normal > networkManager.LogLevel)
+                processor.Emit(OpCodes.Ldloc, netManLocIdx);
+                processor.Emit(OpCodes.Ldfld, m_NetworkManager_LogLevel_FieldRef);
+                processor.Emit(OpCodes.Ldc_I4, (int)LogLevel.Normal);
+                processor.Emit(OpCodes.Cgt);
+                processor.Emit(OpCodes.Ldc_I4, 0);
+                processor.Emit(OpCodes.Ceq);
+                processor.Emit(OpCodes.Brfalse, logNextInstr);
+
+                // Debug.LogError(...);
+                processor.Emit(OpCodes.Ldstr, "Only the server can invoke an Rpc with RpcInvokePermission.Server!");
+                processor.Emit(OpCodes.Call, m_Debug_LogError_MethodRef);
+
+                processor.Append(logNextInstr);
+
                 processor.Append(roReturnInstr);
                 processor.Append(roLastInstr);
             }
diff --git a/com.unity.netcode.gameobjects/Runtime/Core/NetworkBehaviour.cs b/com.unity.netcode.gameobjects/Runtime/Core/NetworkBehaviour.cs
@@ -38,6 +38,7 @@ public abstract class NetworkBehaviour : MonoBehaviour
 
         // RuntimeAccessModifiersILPP will make this `public`
         internal static readonly Dictionary<Type, Dictionary<uint, RpcReceiveHandler>> __rpc_func_table = new Dictionary<Type, Dictionary<uint, RpcReceiveHandler>>();
+        internal static readonly Dictionary<Type, Dictionary<uint, RpcInvokePermission>> __rpc_permission_table = new Dictionary<Type, Dictionary<uint, RpcInvokePermission>>();
 
 #if DEVELOPMENT_BUILD || UNITY_EDITOR || UNITY_MP_TOOLS_NET_STATS_MONITOR_ENABLED_IN_RELEASE
         // RuntimeAccessModifiersILPP will make this `public`
@@ -325,10 +326,14 @@ internal FastBufferWriter __beginSendRpc(uint rpcMethodId, RpcParams rpcParams,
 #pragma warning restore IDE1006 // restore naming rule violation check
         {
             if (m_NetworkObject == null && !IsSpawned)
-            {
+            {   
                 throw new RpcException("The NetworkBehaviour must be spawned before calling this method.");
             }
-            if (attributeParams.RequireOwnership && !IsOwner)
+            else if (attributeParams.InvokePermission == RpcInvokePermission.Server && !IsServer)
+            {
+                throw new RpcException("This RPC can only be sent by the server.");
+            }
+            else if ((attributeParams.RequireOwnership || attributeParams.InvokePermission == RpcInvokePermission.Owner) && !IsOwner)
             {
                 throw new RpcException("This RPC can only be sent by its owner.");
             }
@@ -935,10 +940,11 @@ internal virtual void __initializeRpcs()
 
 #pragma warning disable IDE1006 // disable naming rule violation check
         // RuntimeAccessModifiersILPP will make this `protected`
-        internal void __registerRpc(uint hash, RpcReceiveHandler handler, string rpcMethodName)
+        internal void __registerRpc(uint hash, RpcReceiveHandler handler, RpcInvokePermission permission, string rpcMethodName)
 #pragma warning restore IDE1006 // restore naming rule violation check
         {
             __rpc_func_table[GetType()][hash] = handler;
+            __rpc_permission_table[GetType()][hash] = permission;
 #if DEVELOPMENT_BUILD || UNITY_EDITOR || UNITY_MP_TOOLS_NET_STATS_MONITOR_ENABLED_IN_RELEASE
             __rpc_name_table[GetType()][hash] = rpcMethodName;
 #endif
@@ -985,6 +991,7 @@ internal void InitializeVariables()
             if (!__rpc_func_table.ContainsKey(GetType()))
             {
                 __rpc_func_table[GetType()] = new Dictionary<uint, RpcReceiveHandler>();
+                __rpc_permission_table[GetType()] = new Dictionary<uint, RpcInvokePermission>();
 #if UNITY_EDITOR || DEVELOPMENT_BUILD || UNITY_MP_TOOLS_NET_STATS_MONITOR_ENABLED_IN_RELEASE
                 __rpc_name_table[GetType()] = new Dictionary<uint, string>();
 #endif
diff --git a/com.unity.netcode.gameobjects/Runtime/Messaging/Messages/ProxyMessage.cs b/com.unity.netcode.gameobjects/Runtime/Messaging/Messages/ProxyMessage.cs
@@ -41,9 +41,32 @@ public unsafe void Handle(ref NetworkContext context)
                 }
                 return;
             }
-
             var observers = networkObject.Observers;
 
+            // Validate message if server
+            if (networkManager.IsServer)
+            {
+                var networkBehaviour = networkObject.GetNetworkBehaviourAtOrderIndex(WrappedMessage.Metadata.NetworkBehaviourId);
+
+                RpcInvokePermission permission = NetworkBehaviour.__rpc_permission_table[networkBehaviour.GetType()][WrappedMessage.Metadata.NetworkRpcMethodId];
+                bool hasPermission = permission switch
+                {
+                    RpcInvokePermission.Anyone => true,
+                    RpcInvokePermission.Server => context.SenderId == networkManager.LocalClientId,
+                    RpcInvokePermission.Owner => context.SenderId == networkBehaviour.OwnerClientId,
+                    _ => false,
+                };
+
+                // Do not handle the message if the sender does not have permission to do so.
+                if (!hasPermission)
+                {
+                    return;
+                }
+
+                WrappedMessage.SenderClientId = context.SenderId;
+            }
+
+
             var nonServerIds = new NativeList<ulong>(Allocator.Temp);
             for (var i = 0; i < TargetClientIds.Length; ++i)
             {
diff --git a/com.unity.netcode.gameobjects/Runtime/Messaging/Messages/RpcMessages.cs b/com.unity.netcode.gameobjects/Runtime/Messaging/Messages/RpcMessages.cs
@@ -66,12 +66,29 @@ public static void Handle(ref NetworkContext context, ref RpcMetadata metadata,
                 {
                     NetworkLog.LogWarning($"[{metadata.NetworkObjectId}, {metadata.NetworkBehaviourId}, {metadata.NetworkRpcMethodId}] An RPC called on a {nameof(NetworkObject)} that is not in the spawned objects list. Please make sure the {nameof(NetworkObject)} is spawned before calling RPCs.");
                 }
-                return;
             }
             var networkBehaviour = networkObject.GetNetworkBehaviourAtOrderIndex(metadata.NetworkBehaviourId);
 
             try
             {
+                Type type = networkBehaviour.GetType();
+                if (networkManager.IsServer)
+                {
+                    RpcInvokePermission permission = NetworkBehaviour.__rpc_permission_table[networkBehaviour.GetType()][metadata.NetworkRpcMethodId];
+                    bool hasPermission = permission switch
+                    {
+                        RpcInvokePermission.Anyone => true,
+                        RpcInvokePermission.Server => context.SenderId == networkManager.LocalClientId,
+                        RpcInvokePermission.Owner => context.SenderId == networkBehaviour.OwnerClientId,
+                        _ => false,
+                    };
+
+                    // Do not handle the message if the sender does not have permission to do so.
+                    if (!hasPermission)
+                    {
+                        return;
+                    }
+                }    
                 NetworkBehaviour.__rpc_func_table[networkBehaviour.GetType()][metadata.NetworkRpcMethodId](networkBehaviour, payload, rpcParams);
             }
             catch (Exception ex)
@@ -196,6 +213,12 @@ public unsafe bool Deserialize(FastBufferReader reader, ref NetworkContext conte
 
         public void Handle(ref NetworkContext context)
         {
+            var networkManager = (NetworkManager)context.SystemOwner;
+            if (networkManager.IsServer)
+            {
+                SenderClientId = context.SenderId;
+            }
+            
             var rpcParams = new __RpcParams
             {
                 Ext = new RpcParams
diff --git a/com.unity.netcode.gameobjects/Runtime/Messaging/RpcAttributes.cs b/com.unity.netcode.gameobjects/Runtime/Messaging/RpcAttributes.cs
