Address PR feedback

BepInEx · Jul 5, 2023 · d590fc6 · d590fc6
1 parent 18c9f9b
commit d590fc6
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 16 deletions.
diff --git a/Harmony/Internal/Util/StackTraceFixes.cs b/Harmony/Internal/Util/StackTraceFixes.cs
@@ -3,6 +3,7 @@
 using System.Diagnostics;
 using System.Reflection;
 using HarmonyLib.Tools;
+using MonoMod.Core.Platforms;
 using MonoMod.RuntimeDetour;
 using System.Linq;
 
@@ -19,7 +20,9 @@ internal static class StackTraceFixes
         private static readonly Dictionary<MethodBase, MethodBase> RealMethodMap =
             new Dictionary<MethodBase, MethodBase>();
 
-        private static Hook stackTraceHook;
+        private static Hook getAssemblyHookManaged;
+        private static NativeHook getAssemblyHookNative;
+        private static Hook getMethodHook;
 
         public static void Install()
         {
@@ -31,8 +34,17 @@ public static void Install()
 	            DetourManager.ILHookApplied += OnILChainRefresh;
 	            DetourManager.ILHookUndone += OnILChainRefresh;
 
-	            stackTraceHook = new Hook(AccessTools.Method(typeof(Assembly), nameof(Assembly.GetExecutingAssembly)),
-		            AccessTools.Method(typeof(StackTraceFixes), nameof(GetAssemblyFix)));
+	            var getAssemblyMethod = AccessTools.DeclaredMethod(typeof(Assembly), nameof(Assembly.GetExecutingAssembly), EmptyType.NoArgs);
+	            if (getAssemblyMethod.HasMethodBody())
+	            {
+		            getAssemblyHookManaged = new Hook(getAssemblyMethod, GetAssemblyFix);
+	            }
+	            else
+	            {
+		            getAssemblyHookNative = new NativeHook(PlatformTriple.Current.GetNativeMethodBody(getAssemblyMethod), GetAssemblyFix);
+	            }
+
+	            getMethodHook = new Hook(AccessTools.DeclaredMethod(typeof(StackFrame), nameof(StackFrame.GetMethod), EmptyType.NoArgs), GetMethodFix);
             }
             catch (Exception e)
             {
@@ -41,17 +53,26 @@ public static void Install()
             _applied = true;
         }
 
+        delegate Assembly GetAssemblyDelegate();
+
         // We need to force GetExecutingAssembly make use of stack trace
         // This is to fix cases where calling assembly is actually the patch
         // This solves issues with code where it uses the method to get current filepath etc
-        private static Assembly GetAssemblyFix(Func<Assembly> orig)
+        private static Assembly GetAssemblyFix(GetAssemblyDelegate orig)
+        {
+	        var entry = getAssemblyHookManaged?.DetourInfo.Entry ?? getAssemblyHookNative.DetourInfo.Entry;
+	        var method = new StackTrace().GetFrames()!.Select(f => f.GetMethod()).SkipWhile(method => method != entry).Skip(1).First();
+	        return method.Module.Assembly;
+        }
+
+        private static MethodBase GetMethodFix(Func<StackFrame, MethodBase> orig, StackFrame self)
         {
-	        var method = new StackTrace().GetFrames()!.SkipWhile(frame => frame.GetMethod() != stackTraceHook.DetourInfo.Entry).Skip(1).First().GetMethod();
-	        if (RealMethodMap.TryGetValue(method, out var real))
+	        var method = orig(self);
+	        if (method is not null && RealMethodMap.TryGetValue(PlatformTriple.Current.GetIdentifiable(method), out var real))
 	        {
-		        return real.Module.Assembly;
+		        return real;
 	        }
-	        return orig();
+	        return method;
         }
 
         private static readonly AccessTools.FieldRef<MethodDetourInfo, object> GetDetourState = AccessTools.FieldRefAccess<MethodDetourInfo, object>(AccessTools.DeclaredField(typeof(MethodDetourInfo), "state"));
@@ -64,7 +85,7 @@ private static void OnILChainRefresh(ILHookInfo self)
         {
             lock (RealMethodMap)
             {
-                RealMethodMap[GetEndOfChain(GetDetourState(self.Method))] = self.Method.Method;
+                RealMethodMap[PlatformTriple.Current.GetIdentifiable(GetEndOfChain(GetDetourState(self.Method)))] = PlatformTriple.Current.GetIdentifiable(self.Method.Method);
             }
         }
     }

diff --git a/...ny/Public/Patching/NativeMethodPatcher.cs → ...ublic/Patching/NativeHookMethodPatcher.cs b/...ny/Public/Patching/NativeMethodPatcher.cs → ...ublic/Patching/NativeHookMethodPatcher.cs
@@ -9,7 +9,7 @@
 namespace HarmonyLib.Public.Patching;
 
 /// <inheritdoc/>
-public class NativeMethodPatcher : MethodPatcher
+public class NativeHookMethodPatcher : MethodPatcher
 {
 	private PlatformTriple.NativeDetour? _hook;
 	private Delegate _replacementDelegate;
@@ -21,7 +21,7 @@ public class NativeMethodPatcher : MethodPatcher
 	private readonly MethodInfo _altEntryDelegateInvoke;
 
 	/// <inheritdoc/>
-	public NativeMethodPatcher(MethodBase original) : base(original)
+	public NativeHookMethodPatcher(MethodBase original) : base(original)
 	{
 		var originalParameters = Original.GetParameters();
 
@@ -97,6 +97,6 @@ public override DynamicMethodDefinition CopyOriginal()
 	public static void TryResolve(object _, PatchManager.PatcherResolverEventArgs args)
 	{
 		if (args.Original.GetMethodBody() == null)
-			args.MethodPatcher = new NativeMethodPatcher(args.Original);
+			args.MethodPatcher = new NativeHookMethodPatcher(args.Original);
 	}
 }
diff --git a/Harmony/Public/Patching/PatchManager.cs b/Harmony/Public/Patching/PatchManager.cs
@@ -25,7 +25,7 @@ public static class PatchManager
 		static PatchManager()
 		{
 			ResolvePatcher += ManagedMethodPatcher.TryResolve;
-			ResolvePatcher += NativeMethodPatcher.TryResolve;
+			ResolvePatcher += NativeHookMethodPatcher.TryResolve;
 		}
 
 		/// <summary>
@@ -132,8 +132,8 @@ internal static MethodBase FindReplacement(StackFrame frame)
 			}
 			else
 			{
-				var baseMethod = PlatformTriple.Current.Runtime.GetIdentifiable(frameMethod);
-				methodStart = PlatformTriple.Current.Runtime.GetMethodEntryPoint(baseMethod).ToInt64();
+				var baseMethod = PlatformTriple.Current.GetIdentifiable(frameMethod);
+				methodStart = PlatformTriple.Current.GetNativeMethodBody(baseMethod).ToInt64();
 			}
 
 			// Failed to find any usable method, if `frameMethod` is null, we can not find any
@@ -143,7 +143,7 @@ internal static MethodBase FindReplacement(StackFrame frame)
 
 			lock (ReplacementToOriginals)
 				return ReplacementToOriginals
-					.FirstOrDefault(kv => kv.Key.IsAlive && PlatformTriple.Current.Runtime.GetMethodEntryPoint((MethodBase)kv.Key.Target).ToInt64() == methodStart).Key.Target as MethodBase;
+					.FirstOrDefault(kv => kv.Key.IsAlive && PlatformTriple.Current.GetNativeMethodBody((MethodBase)kv.Key.Target).ToInt64() == methodStart).Key.Target as MethodBase;
 		}
 
 		internal static void AddReplacementOriginal(MethodBase original, MethodInfo replacement)