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