mirror of
https://github.com/D1GQ/BetterAmongUs.git
synced 2026-10-19 15:51:52 +00:00
Fix reference issues with IMonoExtension
This commit is contained in:
1 parent
96a204dc8e
commit
61e355f103
4 files changed
+116
-121
No files matched your search
@@ -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<T> 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>
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user