Fix reference issues with IMonoExtension

This commit is contained in:
D1GQ committed 2026-07-10 23:40:19 -05:00
1 parent 96a204dc8e
commit 61e355f103
4 files changed
+116 -121

No files matched your search

+66 -32
View File
@@ -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
/// <param name="baseMono">The base MonoBehaviour this extension is attached to.</param>
void OnExtensionAwake(MonoBehaviour baseMono);
/// <summary>
/// Called every frame if the extension is enabled.
/// </summary>
void Update();
/// <summary>
/// Called when the extension is being destroyed.
/// </summary>
@@ -39,11 +36,10 @@ internal interface IMonoExtension
/// <summary>
/// Static dictionary mapping MonoBehaviour instances to their active IMonoExtension instances.
/// </summary>
private static readonly Dictionary<MonoBehaviour, IMonoExtension> MonoToMonoExtensionLookup = [];
private static readonly ConditionalWeakTable<MonoBehaviour, List<IMonoExtension>> MonoToMonoExtensionLookup = [];
/// <summary>
/// Registers all auto-extension types from the assembly.
/// Scans for types implementing IAutoMonoExtension&lt;T&gt; and registers them.
/// </summary>
internal static void RegisterAll()
{
@@ -84,9 +80,16 @@ internal interface IMonoExtension
/// <returns>The extension instance if found, otherwise default(T).</returns>
internal static T? GetExtension<T>(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
/// <returns>The newly created extension instance, or null if creation failed.</returns>
internal static T? AddExtension<T>(MonoBehaviour monoBehaviour) where T : MonoBehaviour, IMonoExtension
{
if (monoBehaviour == null)
return null;
// Check if extension already exists
var existing = GetExtension<T>(monoBehaviour);
if (existing != null)
return existing;
T? monoExtension = monoBehaviour.gameObject.AddComponent<T>();
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
/// <param name="monoBehaviour">The MonoBehaviour to add an auto-extension to.</param>
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}");
}
}
/// <summary>
/// Removes and destroys an extension from its attached MonoBehaviour.
/// Tries to removes an extension from its attached MonoBehaviour.
/// </summary>
/// <param name="monoExtension">The extension instance to remove.</param>
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);
}
}
}
@@ -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<PlayerInfoDisplay>().Init(playerControl);
}
public void Update()
{
}
public void OnDestroy()
{
}
/// <summary>
/// Coroutine to add extended player data.
/// </summary>
/// <returns>IEnumerator for coroutine.</returns>
[HideFromIl2Cpp]
private IEnumerator CoAddBetterData()
{
while (BaseMono == null || BaseMono.Data == null)
{
yield return null;
}
TryCreateExtendedData(BaseMono.Data);
}
/// <summary>
/// Attempts to create extended data for a player.
/// </summary>
/// <param name="data">The player data to extend.</param>
internal static void TryCreateExtendedData(NetworkedPlayerInfo data)
{
if (data.ExtendedData() == null)
{
ExtendedPlayerInfo newBetterData = data.gameObject.AddComponent<ExtendedPlayerInfo>();
newBetterData.SetInfo(data);
}
}
}
/// <summary>
+2 -25
View File
@@ -36,24 +36,6 @@ internal sealed class ExtendedPlayerInfo : MonoBehaviour, IAutoMonoExtension<Net
public NetworkedPlayerInfo? BaseMono { get; set; }
private bool hasSet = false;
/// <summary>
/// Initializes the extended player info.
/// </summary>
/// <param name="data">The NetworkedPlayerInfo to extend.</param>
[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<Net
}
}
public void Update()
private float timeAccumulator = 0f;
internal void Update()
{
var time = Time.deltaTime;
@@ -128,12 +111,6 @@ internal sealed class ExtendedPlayerInfo : MonoBehaviour, IAutoMonoExtension<Net
[HideFromIl2Cpp]
internal HandshakeHandler? HandshakeHandler { get; }
/// <summary>
/// Gets the player ID.
/// </summary>
[HideFromIl2Cpp]
internal byte _PlayerId { get; private set; }
/// <summary>
/// Gets the player's real name.
/// </summary>
+48 -30
View File
@@ -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<Type> 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<T>
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<Type> FindAllExtensionTypes()
{
var types = new List<Type>();
@@ -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);
}
}