diff --git a/src/Interfaces/IMonoExtension.cs b/src/Interfaces/IMonoExtension.cs index decfab3..fe1753b 100644 --- a/src/Interfaces/IMonoExtension.cs +++ b/src/Interfaces/IMonoExtension.cs @@ -1,4 +1,6 @@ -using Il2CppInterop.Runtime; +using BetterAmongUs.Modules; +using Il2CppInterop.Runtime; +using System.Runtime.CompilerServices; using UnityEngine; namespace BetterAmongUs.Interfaces; @@ -21,11 +23,6 @@ internal interface IMonoExtension /// The base MonoBehaviour this extension is attached to. void OnExtensionAwake(MonoBehaviour baseMono); - /// - /// Called every frame if the extension is enabled. - /// - void Update(); - /// /// Called when the extension is being destroyed. /// @@ -39,11 +36,10 @@ internal interface IMonoExtension /// /// Static dictionary mapping MonoBehaviour instances to their active IMonoExtension instances. /// - private static readonly Dictionary MonoToMonoExtensionLookup = []; + private static readonly ConditionalWeakTable> MonoToMonoExtensionLookup = []; /// /// Registers all auto-extension types from the assembly. - /// Scans for types implementing IAutoMonoExtension<T> and registers them. /// internal static void RegisterAll() { @@ -84,9 +80,16 @@ internal interface IMonoExtension /// The extension instance if found, otherwise default(T). internal static T? GetExtension(MonoBehaviour monoBehaviour) where T : IMonoExtension { - if (MonoToMonoExtensionLookup.TryGetValue(monoBehaviour, out var monoExtension)) + if (monoBehaviour == null) + return default; + + if (MonoToMonoExtensionLookup.TryGetValue(monoBehaviour, out var extensions)) { - return (T)monoExtension; + foreach (var extension in extensions) + { + if (extension is T typedExtension) + return typedExtension; + } } return default; @@ -100,10 +103,25 @@ internal interface IMonoExtension /// The newly created extension instance, or null if creation failed. internal static T? AddExtension(MonoBehaviour monoBehaviour) where T : MonoBehaviour, IMonoExtension { + if (monoBehaviour == null) + return null; + + // Check if extension already exists + var existing = GetExtension(monoBehaviour); + if (existing != null) + return existing; + T? monoExtension = monoBehaviour.gameObject.AddComponent(); if (monoExtension != null) { - MonoToMonoExtensionLookup[monoBehaviour] = monoExtension; + // Get or create the extension list for this MonoBehaviour + if (!MonoToMonoExtensionLookup.TryGetValue(monoBehaviour, out var extensions)) + { + extensions = []; + MonoToMonoExtensionLookup.Add(monoBehaviour, extensions); + } + + extensions.Add(monoExtension); monoExtension.BaseMono = monoBehaviour; monoExtension.OnExtensionAwake(monoBehaviour); return monoExtension; @@ -118,48 +136,64 @@ internal interface IMonoExtension /// The MonoBehaviour to add an auto-extension to. internal static void TryAddAutoExtension(MonoBehaviour monoBehaviour) { - if (AutoExtensionTypeLookup.TryGetValue(monoBehaviour.GetType(), out var extensionType)) + if (monoBehaviour == null) + return; + + var monoType = monoBehaviour.GetType(); + + // Check if this type has an auto-extension registered + if (!AutoExtensionTypeLookup.TryGetValue(monoType, out var extensionType)) + return; + + // Check if extension already exists for this MonoBehaviour instance + if (MonoToMonoExtensionLookup.TryGetValue(monoBehaviour, out var existingExtensions)) + { + foreach (var ext in existingExtensions) + { + if (ext.GetType() == extensionType) + return; + } + } + + try { IMonoExtension monoExtension = (IMonoExtension)monoBehaviour.gameObject.AddComponent(Il2CppType.From(extensionType)); if (monoExtension != null) { - MonoToMonoExtensionLookup[monoBehaviour] = monoExtension; + // Get or create the extension list for this MonoBehaviour + if (!MonoToMonoExtensionLookup.TryGetValue(monoBehaviour, out var extensions)) + { + extensions = []; + MonoToMonoExtensionLookup.Add(monoBehaviour, extensions); + } + + extensions.Add(monoExtension); monoExtension.BaseMono = monoBehaviour; monoExtension.OnExtensionAwake(monoBehaviour); } } + catch (Exception ex) + { + Logger_.Error($"Failed to add auto-extension {extensionType.Name} to {monoType.Name}: {ex.Message}"); + } } /// - /// Removes and destroys an extension from its attached MonoBehaviour. + /// Tries to removes an extension from its attached MonoBehaviour. /// /// The extension instance to remove. - internal static void RemoveExtension(IMonoExtension monoExtension) + internal static void TryRemoveExtension(IMonoExtension monoExtension) { if (monoExtension == null) return; - if (!MonoToMonoExtensionLookup.ContainsValue(monoExtension)) - return; - - MonoBehaviour? keyToRemove = null; - bool found = false; - - foreach (var kvp in MonoToMonoExtensionLookup) + if (monoExtension.BaseMono != null) { - if (kvp.Value == monoExtension) + if (MonoToMonoExtensionLookup.TryGetValue(monoExtension.BaseMono, out var extensions)) { - keyToRemove = kvp.Key; - found = true; - break; + extensions.Remove(monoExtension); } } - - if (found) - { - MonoToMonoExtensionLookup.Remove(keyToRemove); - UnityEngine.Object.Destroy((MonoBehaviour)monoExtension); - } } } diff --git a/src/MonoScripts/Extended/ExtendedPlayerControl.cs b/src/MonoScripts/Extended/ExtendedPlayerControl.cs index 68fa329..ccdb67a 100644 --- a/src/MonoScripts/Extended/ExtendedPlayerControl.cs +++ b/src/MonoScripts/Extended/ExtendedPlayerControl.cs @@ -1,8 +1,6 @@ using BetterAmongUs.Attributes; using BetterAmongUs.Interfaces; using HarmonyLib; -using Il2CppInterop.Runtime.Attributes; -using System.Collections; using UnityEngine; namespace BetterAmongUs.MonoScripts.Extended; @@ -23,41 +21,9 @@ internal sealed class ExtendedPlayerControl : MonoBehaviour, IAutoMonoExtension< playerControl.gameObject.AddComponent().Init(playerControl); } - public void Update() - { - } - public void OnDestroy() { } - - /// - /// Coroutine to add extended player data. - /// - /// IEnumerator for coroutine. - [HideFromIl2Cpp] - private IEnumerator CoAddBetterData() - { - while (BaseMono == null || BaseMono.Data == null) - { - yield return null; - } - - TryCreateExtendedData(BaseMono.Data); - } - - /// - /// Attempts to create extended data for a player. - /// - /// The player data to extend. - internal static void TryCreateExtendedData(NetworkedPlayerInfo data) - { - if (data.ExtendedData() == null) - { - ExtendedPlayerInfo newBetterData = data.gameObject.AddComponent(); - newBetterData.SetInfo(data); - } - } } /// diff --git a/src/MonoScripts/Extended/ExtendedPlayerInfo.cs b/src/MonoScripts/Extended/ExtendedPlayerInfo.cs index 1c7c8b6..41c4675 100644 --- a/src/MonoScripts/Extended/ExtendedPlayerInfo.cs +++ b/src/MonoScripts/Extended/ExtendedPlayerInfo.cs @@ -36,24 +36,6 @@ internal sealed class ExtendedPlayerInfo : MonoBehaviour, IAutoMonoExtension - /// Initializes the extended player info. - /// - /// The NetworkedPlayerInfo to extend. - [HideFromIl2Cpp] - internal void SetInfo(NetworkedPlayerInfo data) - { - if (hasSet) - return; - - hasSet = true; - _PlayerId = data.PlayerId; - } - - private float timeAccumulator = 0f; - public void OnExtensionAwake(NetworkedPlayerInfo networkedPlayerInfo) { if (HandshakeHandler != null) @@ -62,7 +44,8 @@ internal sealed class ExtendedPlayerInfo : MonoBehaviour, IAutoMonoExtension - /// Gets the player ID. - /// - [HideFromIl2Cpp] - internal byte _PlayerId { get; private set; } - /// /// Gets the player's real name. /// diff --git a/src/Patches/Unity/MonoExtensionPatch.cs b/src/Patches/Unity/MonoExtensionPatch.cs index aafbefe..1b5f57a 100644 --- a/src/Patches/Unity/MonoExtensionPatch.cs +++ b/src/Patches/Unity/MonoExtensionPatch.cs @@ -9,13 +9,7 @@ namespace BetterAmongUs.Patches.Unity; internal static class MonoExtensionPatch { private static bool _patched = false; - - [HarmonyPatch(typeof(MonoBehaviour), MethodType.Constructor, [typeof(IntPtr)])] - [HarmonyPostfix] - private static void MonoBehaviour_Constructor_Postfix(MonoBehaviour __instance) - { - IMonoExtension.TryAddAutoExtension(__instance); - } + private static readonly HashSet PatchedTypes = []; internal static void Patch(Harmony harmony) { @@ -27,20 +21,31 @@ internal static class MonoExtensionPatch // Find all types that implement IMonoExtension var extensionTypes = FindAllExtensionTypes(); - foreach (var type in extensionTypes) + foreach (var extType in extensionTypes) { - // Patch Update method if it exists - var updateMethod = AccessTools.Method(type, nameof(IMonoExtension.Update)); - if (updateMethod != null) + // Get the target MonoBehaviour type from IAutoMonoExtension + var targetType = GetTargetMonoBehaviourType(extType); + if (targetType == null) continue; + + // Only patch each target type once + if (!PatchedTypes.Add(targetType)) continue; + + // Patch the constructor of the target type + var constructor = AccessTools.Constructor(targetType, [typeof(IntPtr)]); + if (constructor != null) { - harmony.Patch(updateMethod, prefix: new HarmonyMethod(typeof(MonoExtensionPatch), nameof(IMonoExtension_Update_Prefix))); + harmony.Patch(constructor, + postfix: new HarmonyMethod(typeof(MonoExtensionPatch), + nameof(MonoBehaviour_Constructor_Postfix))); } - // Patch OnDestroy method if it exists - var destroyMethod = AccessTools.Method(type, nameof(IMonoExtension.OnDestroy)); + // Patch OnDestroy method if it exists on the extension type + var destroyMethod = AccessTools.Method(extType, nameof(IMonoExtension.OnDestroy)); if (destroyMethod != null) { - harmony.Patch(destroyMethod, postfix: new HarmonyMethod(typeof(MonoExtensionPatch), nameof(IMonoExtension_OnDestroy_Postfix))); + harmony.Patch(destroyMethod, + postfix: new HarmonyMethod(typeof(MonoExtensionPatch), + nameof(IMonoExtension_OnDestroy_Postfix))); } } } @@ -58,20 +63,22 @@ internal static class MonoExtensionPatch try { - // Find all types that implement IMonoExtension var extensionTypes = FindAllExtensionTypes(); - foreach (var type in extensionTypes) + foreach (var extType in extensionTypes) { - // Unpatch Update method - var updateMethod = AccessTools.Method(type, nameof(IMonoExtension.Update)); - if (updateMethod != null) + var targetType = GetTargetMonoBehaviourType(extType); + if (targetType == null) continue; + + // Unpatch constructor + var constructor = AccessTools.Constructor(targetType, [typeof(IntPtr)]); + if (constructor != null) { - harmony.Unpatch(updateMethod, HarmonyPatchType.Prefix, harmony.Id); + harmony.Unpatch(constructor, HarmonyPatchType.Postfix, harmony.Id); } - // Unpatch OnDestroy method - var destroyMethod = AccessTools.Method(type, nameof(IMonoExtension.OnDestroy)); + // Unpatch OnDestroy + var destroyMethod = AccessTools.Method(extType, nameof(IMonoExtension.OnDestroy)); if (destroyMethod != null) { harmony.Unpatch(destroyMethod, HarmonyPatchType.Postfix, harmony.Id); @@ -85,6 +92,15 @@ internal static class MonoExtensionPatch } } + private static Type? GetTargetMonoBehaviourType(Type extensionType) + { + var genericInterface = extensionType.GetInterfaces() + .FirstOrDefault(i => i.IsGenericType && + i.GetGenericTypeDefinition() == typeof(IAutoMonoExtension<>)); + + return genericInterface?.GetGenericArguments()[0]; + } + private static List FindAllExtensionTypes() { var types = new List(); @@ -107,19 +123,21 @@ internal static class MonoExtensionPatch return types; } - private static bool IMonoExtension_Update_Prefix(IMonoExtension __instance) + private static void MonoBehaviour_Constructor_Postfix(MonoBehaviour __instance) { - if (__instance.BaseMono == null) + try { - IMonoExtension.RemoveExtension(__instance); - return false; + if (__instance == null) return; + IMonoExtension.TryAddAutoExtension(__instance); + } + catch (Exception ex) + { + Logger_.Error($"Error in MonoExtension constructor patch: {ex.Message}"); } - - return true; } private static void IMonoExtension_OnDestroy_Postfix(IMonoExtension __instance) { - IMonoExtension.RemoveExtension(__instance); + IMonoExtension.TryRemoveExtension(__instance); } } \ No newline at end of file