From 7c31daa60e644de14bdfb8e2977b3681686a3097 Mon Sep 17 00:00:00 2001 From: D1GQ Date: Mon, 13 Jul 2026 13:50:22 -0500 Subject: [PATCH] Rework IMonoExtension auto patching --- src/BAUPlugin.cs | 5 +- src/Interfaces/IMonoExtension.cs | 275 ++++++++++++------ .../Extended/ExtendedPassiveButton.cs | 1 + .../Extended/ExtendedPlayerControl.cs | 18 +- .../Extended/ExtendedPlayerInfo.cs | 22 +- src/Patches/Hooks/MonoExtensionHook.cs | 140 --------- 6 files changed, 226 insertions(+), 235 deletions(-) delete mode 100644 src/Patches/Hooks/MonoExtensionHook.cs diff --git a/src/BAUPlugin.cs b/src/BAUPlugin.cs index 7146cd1..95b784c 100644 --- a/src/BAUPlugin.cs +++ b/src/BAUPlugin.cs @@ -16,7 +16,6 @@ using BetterAmongUs.Modules.Support; using BetterAmongUs.Network; using BetterAmongUs.Patches.Client; using BetterAmongUs.Patches.Gameplay.UI.Settings; -using BetterAmongUs.Patches.Hooks; using BetterAmongUs.Utilities; using HarmonyLib; using UnityEngine; @@ -159,11 +158,10 @@ internal class BAUPlugin : BasePlugin AudioOverrideManager.Initialize(); Translator.Initialize(); Harmony.PatchAll(); - MonoExtensionHook.Install(); + IMonoExtensionPatcher.PatchAll(); GameSettingsPatch.SetupSettings(true); BAUModdedSupportEvents.OnBAUOptionsLoadedEvent.InvokeAll([.. OptionItem.AllOptions.Cast()]); AutoRegisterAttribute.Initialize(); - IMonoExtension.RegisterAll(); OutfitData.Initialize(); SceneManager.add_sceneLoaded((Action)OnSceneLoaded); @@ -186,7 +184,6 @@ internal class BAUPlugin : BasePlugin BetterNotificationManager.Detach(); ClientPatch.Unpatch(); Harmony.UnpatchAll(); - MonoExtensionHook.Uninstall(); ModManager.Instance.ModStamp.gameObject.SetActive(false); SceneChanger.ChangeScene("MainMenu"); } diff --git a/src/Interfaces/IMonoExtension.cs b/src/Interfaces/IMonoExtension.cs index fe1753b..374e1db 100644 --- a/src/Interfaces/IMonoExtension.cs +++ b/src/Interfaces/IMonoExtension.cs @@ -1,6 +1,7 @@ -using BetterAmongUs.Modules; -using Il2CppInterop.Runtime; +using HarmonyLib; +using System.Reflection; using System.Runtime.CompilerServices; +using System.Runtime.Serialization; using UnityEngine; namespace BetterAmongUs.Interfaces; @@ -28,48 +29,32 @@ internal interface IMonoExtension /// void OnDestroy(); - /// - /// Static dictionary mapping MonoBehaviour types to their auto-extension types. - /// - private static readonly Dictionary AutoExtensionTypeLookup = []; - /// /// Static dictionary mapping MonoBehaviour instances to their active IMonoExtension instances. /// private static readonly ConditionalWeakTable> MonoToMonoExtensionLookup = []; /// - /// Registers all auto-extension types from the assembly. + /// Determines whether a MonoBehaviour has an extension of the specified type attached to it. /// - internal static void RegisterAll() + /// The type of IMonoExtension to check for. + /// The MonoBehaviour to check for the extension. + /// true if an extension of type T exists on the MonoBehaviour; otherwise, false. + internal static bool HasExtension(MonoBehaviour monoBehaviour) where T : IMonoExtension { - var assembly = ModInfo.Assembly; + if (monoBehaviour == null) + return false; - foreach (var type in assembly.GetTypes()) + if (MonoToMonoExtensionLookup.TryGetValue(monoBehaviour, out var extensions)) { - if (type.IsInterface || type.IsAbstract) - continue; - - if (!typeof(IMonoExtension).IsAssignableFrom(type)) - continue; - - if (type.IsGenericTypeDefinition) - continue; - - var genericInterface = type.GetInterfaces() - .FirstOrDefault(i => i.IsGenericType && - i.GetGenericTypeDefinition() == typeof(IAutoMonoExtension<>)); - - if (genericInterface == null) - continue; - - var monoType = genericInterface.GetGenericArguments()[0]; - - if (!monoType.IsSubclassOf(typeof(MonoBehaviour))) - continue; - - AutoExtensionTypeLookup[monoType] = type; + foreach (var extension in extensions) + { + if (extension is T) + return true; + } } + + return false; } /// @@ -130,54 +115,6 @@ internal interface IMonoExtension return null; } - /// - /// Attempts to add an auto-extension to a MonoBehaviour if one is registered. - /// - /// The MonoBehaviour to add an auto-extension to. - internal static void TryAddAutoExtension(MonoBehaviour monoBehaviour) - { - 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) - { - // 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}"); - } - } - /// /// Tries to removes an extension from its attached MonoBehaviour. /// @@ -237,8 +174,176 @@ internal interface IMonoExtension : IMonoExtension where T : MonoBehaviour } /// -/// Marker interface for auto-registering extensions. -/// When a type implements this interface with a specific MonoBehaviour type T, it will be automatically registered and attached to matching MonoBehaviour instances. +/// Provides functionality to patch MonoBehaviour methods and apply extension patches. /// -/// The type of MonoBehaviour this extension automatically attaches to. -internal interface IAutoMonoExtension : IMonoExtension where T : MonoBehaviour; \ No newline at end of file +internal interface IMonoExtensionPatcher +{ + /// + /// Gets the target patch information. + /// + TargetPatch Target { get; } + + /// + /// Adds an extension patch to the specified MonoBehaviour. + /// + /// The MonoBehaviour to patch. + void AddExtensionPatch(MonoBehaviour monoBehaviour); + + /// + /// Represents a target method to be patched. + /// + /// The type containing the method. + /// The name of the method to patch. + internal sealed record TargetPatch(Type Type, string MethodName); + + /// + /// Patches all IMonoExtensionPatcher implementations found in the mod assembly. + /// + internal static void PatchAll() + { + var assembly = ModInfo.Assembly; + var patcherTypes = assembly.GetTypes() + .Where(t => typeof(IMonoExtensionPatcher).IsAssignableFrom(t) && !t.IsInterface && !t.IsAbstract) + .ToList(); + + // Group patchers by their target method to avoid duplicate patches + var patchersByTarget = patcherTypes + .Select(t => (Type: t, Instance: (IMonoExtensionPatcher)FormatterServices.GetUninitializedObject(t))) + .GroupBy(x => x.Instance.Target) + .ToList(); + + foreach (var group in patchersByTarget) + { + var target = group.Key; + var patcherTypesForTarget = group.Select(x => x.Type).ToList(); + + var originalMethod = target.Type.GetMethod(target.MethodName, + BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic); + + if (originalMethod == null) + continue; + + _monoBehaviourToPatcherMap[target.Type] = patcherTypesForTarget; + + var postfixMethod = typeof(IMonoExtensionPatcher).GetMethod(nameof(Postfix), + BindingFlags.Static | BindingFlags.NonPublic); + + BAUPlugin.Harmony.Patch(originalMethod, postfix: new HarmonyMethod(postfixMethod)); + } + } + + /// + /// Patches a specific IMonoExtensionPatcher implementation. + /// + /// The Harmony instance to use for patching. + /// The patcher type to apply. + private static void Patch(Harmony harmony, Type monoExtensionPatcherType) + { + var instance = (IMonoExtensionPatcher)FormatterServices.GetUninitializedObject(monoExtensionPatcherType); + var target = instance.Target; + + var originalMethod = target.Type.GetMethod(target.MethodName, + BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic); + + if (originalMethod == null) + return; + + if (!_monoBehaviourToPatcherMap.TryGetValue(target.Type, out var patcherTypes)) + { + patcherTypes = []; + _monoBehaviourToPatcherMap[target.Type] = patcherTypes; + } + + if (!patcherTypes.Contains(monoExtensionPatcherType)) + { + patcherTypes.Add(monoExtensionPatcherType); + } + + var postfixMethod = typeof(IMonoExtensionPatcher).GetMethod(nameof(Postfix), + BindingFlags.Static | BindingFlags.NonPublic); + + harmony.Patch(originalMethod, postfix: new HarmonyMethod(postfixMethod)); + } + + /// + /// Postfix method that applies extension patches to MonoBehaviour instances. + /// + /// The MonoBehaviour instance being patched. + private static void Postfix(MonoBehaviour __instance) + { + var type = __instance.GetType(); + + // Try exact match first + if (!_monoBehaviourToPatcherMap.TryGetValue(type, out var patcherTypes)) + { + // Check for assignable types (handles inheritance) + foreach (var kvp in _monoBehaviourToPatcherMap) + { + if (kvp.Key.IsAssignableFrom(type)) + { + patcherTypes = kvp.Value; + _monoBehaviourToPatcherMap[type] = patcherTypes; + break; + } + } + + if (patcherTypes == null) + return; + } + + foreach (var patcherType in patcherTypes) + { + var patcher = GetUninitializedMonoExtensionPatcher(patcherType, __instance); + patcher.AddExtensionPatch(__instance); + } + } + + /// + /// Cache for uninitialized patcher instances. + /// + private static readonly Dictionary _uninitializedMonoExtensionPatcherLookup = []; + + /// + /// Maps MonoBehaviour types to their corresponding patcher types. + /// + private static readonly Dictionary> _monoBehaviourToPatcherMap = []; + + /// + /// Gets or creates an uninitialized patcher instance for the specified type. + /// + /// The patcher type. + /// The MonoBehaviour instance. + /// An uninitialized patcher instance. + private static IMonoExtensionPatcher GetUninitializedMonoExtensionPatcher(Type monoExtensionPatcherType, MonoBehaviour monoBehaviour) + { + if (!_uninitializedMonoExtensionPatcherLookup.TryGetValue(monoExtensionPatcherType, out var uninitialized)) + { + uninitialized = (IMonoExtensionPatcher)FormatterServices.GetUninitializedObject(monoExtensionPatcherType); + _uninitializedMonoExtensionPatcherLookup[monoExtensionPatcherType] = uninitialized; + } + + return uninitialized; + } +} + +/// +/// Provides generic functionality to patch MonoBehaviour methods and apply extension patches. +/// +/// The specific MonoBehaviour type to patch. +internal interface IMonoExtensionPatcher : IMonoExtensionPatcher where T : MonoBehaviour +{ + /// + /// Adds an extension patch to the specified MonoBehaviour. + /// + /// The MonoBehaviour to patch. + void IMonoExtensionPatcher.AddExtensionPatch(MonoBehaviour monoBehaviour) + { + AddExtensionPatch((T)monoBehaviour); + } + + /// + /// Adds an extension patch to the specified MonoBehaviour of type T. + /// + /// The MonoBehaviour instance to patch. + void AddExtensionPatch(T monoBehaviour); +} \ No newline at end of file diff --git a/src/MonoScripts/Extended/ExtendedPassiveButton.cs b/src/MonoScripts/Extended/ExtendedPassiveButton.cs index 735830d..f530aee 100644 --- a/src/MonoScripts/Extended/ExtendedPassiveButton.cs +++ b/src/MonoScripts/Extended/ExtendedPassiveButton.cs @@ -59,6 +59,7 @@ internal class ExtendedPassiveButton : MonoBehaviour, IMonoExtension diff --git a/src/MonoScripts/Extended/ExtendedPlayerControl.cs b/src/MonoScripts/Extended/ExtendedPlayerControl.cs index ccdb67a..5ab6a83 100644 --- a/src/MonoScripts/Extended/ExtendedPlayerControl.cs +++ b/src/MonoScripts/Extended/ExtendedPlayerControl.cs @@ -9,20 +9,28 @@ namespace BetterAmongUs.MonoScripts.Extended; /// Extends PlayerControl with additional functionality. /// [RegisterInIl2Cpp] -internal sealed class ExtendedPlayerControl : MonoBehaviour, IAutoMonoExtension +internal sealed class ExtendedPlayerControl : MonoBehaviour, IMonoExtension, IMonoExtensionPatcher { - /// - /// Gets or sets the base PlayerControl instance. - /// - public PlayerControl? BaseMono { get; set; } + public IMonoExtensionPatcher.TargetPatch Target => new(typeof(PlayerControl), nameof(PlayerControl.Awake)); + + public void AddExtensionPatch(PlayerControl playerControl) + { + IMonoExtension.AddExtension(playerControl); + } public void OnExtensionAwake(PlayerControl playerControl) { playerControl.gameObject.AddComponent().Init(playerControl); } + /// + /// Gets or sets the base PlayerControl instance. + /// + public PlayerControl? BaseMono { get; set; } + public void OnDestroy() { + IMonoExtension.TryRemoveExtension(this); } } diff --git a/src/MonoScripts/Extended/ExtendedPlayerInfo.cs b/src/MonoScripts/Extended/ExtendedPlayerInfo.cs index 41c4675..f5bf624 100644 --- a/src/MonoScripts/Extended/ExtendedPlayerInfo.cs +++ b/src/MonoScripts/Extended/ExtendedPlayerInfo.cs @@ -1,4 +1,5 @@ using AmongUs.GameOptions; +using BepInEx.Unity.IL2CPP.Utils; using BetterAmongUs.Attributes; using BetterAmongUs.Generated; using BetterAmongUs.Interfaces; @@ -11,6 +12,7 @@ using HarmonyLib; using Hazel; using Il2CppInterop.Runtime.Attributes; using InnerNet; +using System.Collections; using UnityEngine; namespace BetterAmongUs.MonoScripts.Extended; @@ -19,7 +21,7 @@ namespace BetterAmongUs.MonoScripts.Extended; /// Extended player information with additional data and anti-cheat features. /// [RegisterInIl2Cpp] -internal sealed class ExtendedPlayerInfo : MonoBehaviour, IAutoMonoExtension +internal sealed class ExtendedPlayerInfo : MonoBehaviour, IMonoExtension, IMonoExtensionPatcher { internal ExtendedPlayerInfo() { @@ -34,6 +36,23 @@ internal sealed class ExtendedPlayerInfo : MonoBehaviour, IAutoMonoExtension new(typeof(PlayerControl), nameof(PlayerControl.Awake)); + + public void AddExtensionPatch(PlayerControl playerControl) + { + playerControl.StartCoroutine(CoAddExtensionPatch(playerControl)); + } + + private static IEnumerator CoAddExtensionPatch(PlayerControl playerControl) + { + while (playerControl.Data == null) + { + yield return null; + } + + IMonoExtension.AddExtension(playerControl.Data); + } + public NetworkedPlayerInfo? BaseMono { get; set; } public void OnExtensionAwake(NetworkedPlayerInfo networkedPlayerInfo) @@ -81,6 +100,7 @@ internal sealed class ExtendedPlayerInfo : MonoBehaviour, IAutoMonoExtension PatchedTypes = []; - private static readonly List _detours = []; - - internal static void Install() - { - if (_patched) return; - _patched = true; - - try - { - var extensionTypes = FindAllExtensionTypes(); - - foreach (var extType in extensionTypes) - { - var targetType = GetTargetMonoBehaviourType(extType); - if (targetType == null) continue; - - if (!PatchedTypes.Add(targetType)) continue; - - // Hook the IntPtr constructor - var constructor = AccessTools.Constructor(targetType, [typeof(IntPtr)]); - if (constructor != null) - { - var hook = new MonoMod.RuntimeDetour.Hook( - constructor, - new Action, MonoBehaviour, IntPtr>(MonoBehaviour_Constructor_Hook) - ); - _detours.Add(hook); - } - - // Hook OnDestroy - var destroyMethod = AccessTools.Method(extType, nameof(IMonoExtension.OnDestroy)); - if (destroyMethod != null) - { - var hook = new MonoMod.RuntimeDetour.Hook( - destroyMethod, - new Action, IMonoExtension>(IMonoExtension_OnDestroy_Hook) - ); - _detours.Add(hook); - } - } - } - catch (Exception ex) - { - _patched = false; - Logger_.Error($"Failed to hook MonoExtension methods: {ex.Message}"); - throw; - } - } - - internal static void Uninstall() - { - if (!_patched) return; - - try - { - foreach (var detour in _detours) - { - detour.Dispose(); - } - _detours.Clear(); - _patched = false; - } - catch (Exception ex) - { - Logger_.Error($"Failed to unpatch MonoExtension methods: {ex.Message}"); - throw; - } - } - - 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(); - var assembly = ModInfo.Assembly; - - foreach (var type in assembly.GetTypes()) - { - if (type.IsInterface || type.IsAbstract) - continue; - - if (!typeof(IMonoExtension).IsAssignableFrom(type)) - continue; - - if (type.IsGenericTypeDefinition) - continue; - - types.Add(type); - } - - return types; - } - - private static void MonoBehaviour_Constructor_Hook(Action orig, MonoBehaviour instance, IntPtr ptr) - { - orig(instance, ptr); - - try - { - IMonoExtension.TryAddAutoExtension(instance); - } - catch (Exception ex) - { - Logger_.Error($"Error in constructor hook: {ex.Message}"); - } - } - - private static void IMonoExtension_OnDestroy_Hook(Action orig, IMonoExtension instance) - { - orig(instance); - - try - { - IMonoExtension.TryRemoveExtension(instance); - } - catch (Exception ex) - { - Logger_.Error($"Error in OnDestroy hook: {ex.Message}"); - } - } -} \ No newline at end of file