Cleanup NativeStackWalk and Il2CppInterop Support

LavaGang · Aug 23, 2024 · 008054a · 008054a
1 parent 58473a4
commit 008054a
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 39 deletions.
diff --git a/Dependencies/SupportModules/Il2Cpp/Main.cs b/Dependencies/SupportModules/Il2Cpp/Main.cs
@@ -130,74 +130,96 @@ private sealed class MelonDetour : IDetour
             public nint Target => _detourFrom;
 
             public nint Detour => _targetPtr;
-
             public nint OriginalTrampoline => _originalPtr;
-//             {
-//                 get
-//                 {
-//                     MelonLogger.Msg($"Getting original trampoline => 0x{_originalPtr:X}");
-//                     return _originalPtr;
-//                 }
-//             }
 
             public MelonDetour(nint detourFrom, Delegate target)
             {
                 _detourFrom = detourFrom;
                 _target = target;
-
-                Apply(); //We have to apply immediately because we're gonna be asked for a trampoline right away
+
+                // We have to apply immediately because we're gonna be asked for a trampoline right away
+                Apply();
             }
 
             public unsafe void Apply()
             {
-                if(_targetPtr != IntPtr.Zero)
-                    // Already applied
+                if (_targetPtr != IntPtr.Zero)
                     return;
 
                 _targetPtr = Marshal.GetFunctionPointerForDelegate(_target);
 
-//                 MelonLogger.Msg($"About to detour 0x{_detourFrom:X} to 0x{_targetPtr:X} for method {_target.Method.Name}");
-
                 var addr = _detourFrom;
-                MelonUtils.NativeHookAttachDirect((nint) (&addr), _targetPtr);
+                nint addrPtr = (nint)(&addr);
+                MelonUtils.NativeHookAttachDirect(addrPtr, _targetPtr);
+                NativeStackWalk.RegisterHookAddr((ulong)addrPtr, $"Il2CppInterop detour of 0x{addrPtr:X} -> 0x{_targetPtr:X}");
+
                 _originalPtr = addr;
-
-//                 MelonLogger.Msg($"Applied detour from {_detourFrom:X} to {_targetPtr:X} for method {_target.Method.Name}, original is now: {_originalPtr:X})");
-
-                NativeStackWalk.RegisterHookAddr((ulong)_targetPtr, $"Harmony Hook to {_target.Method.Name}");
             }
 
             public unsafe void Dispose()
             {
-//                 MelonLogger.Msg($"Removing detour from 0x{_detourFrom:X} to 0x{_targetPtr:X} for method {_target.Method.Name}");
+                if (_targetPtr == IntPtr.Zero)
+                    return;
+
                 var addr = _detourFrom;
-                MelonUtils.NativeHookDetach((nint) (&addr), _targetPtr);
+                nint addrPtr = (nint)(&addr);
+
+                MelonUtils.NativeHookDetach(addrPtr, _targetPtr);
+                NativeStackWalk.UnregisterHookAddr((ulong)addrPtr);
+
                 _targetPtr = IntPtr.Zero;
-//                 _originalPtr = IntPtr.Zero;
-//                 MelonLogger.Msg($"Address after removing detour {_target.Method.Name}: {addr:X}");
+                _originalPtr = IntPtr.Zero;
             }
 
-            public T GenerateTrampoline<T>() where T : Delegate
+            public T GenerateTrampoline<T>()
+                where T : Delegate
             {
-//                 MelonLogger.Msg($"Getting delegate for original method at 0x{_originalPtr:X}, type: {typeof(T)}, method name {_target.Method.Name}");
+                if (_originalPtr == IntPtr.Zero)
+                    return null;
                 return Marshal.GetDelegateForFunctionPointer<T>(_originalPtr);
             }
         }
     }
 
-    internal class InteropLogger : Microsoft.Extensions.Logging.ILogger {
-        public void Log<TState>(LogLevel logLevel, EventId eventId, TState state, Exception exception, Func<TState, Exception, string> formatter) {
-            if (logLevel is LogLevel.Debug or LogLevel.Trace)
-                MelonDebug.Msg(formatter(state, exception));
-            else
-                MelonLogger.Msg(formatter(state, exception));
+    internal class InteropLogger
+        : Microsoft.Extensions.Logging.ILogger
+    {
+        private MelonLogger.Instance _logger = new("Il2CppInterop");
+
+        public void Log<TState>(LogLevel logLevel, EventId eventId, TState state, Exception exception,
+            Func<TState, Exception, string> formatter)
+        {
+            string formattedTxt = formatter(state, exception);
+            switch (logLevel)
+            {
+                case LogLevel.Debug:
+                case LogLevel.Trace:
+                    MelonDebug.Msg(formattedTxt);
+                    break;
+
+                case LogLevel.Error:
+                    _logger.Error(formattedTxt);
+                    break;
+
+                case LogLevel.Warning:
+                    _logger.Warning(formattedTxt);
+                    break;
+
+                case LogLevel.Information:
+                default:
+                    _logger.Msg(formattedTxt);
+                    break;
+            }
         }
 
-        public bool IsEnabled(LogLevel logLevel) => logLevel switch {
-            LogLevel.Debug or LogLevel.Trace => MelonDebug.IsEnabled(),
-            _ => true
-        };
+        public bool IsEnabled(LogLevel logLevel)
+            => logLevel switch
+            {
+                LogLevel.Debug or LogLevel.Trace => MelonDebug.IsEnabled(),
+                _ => true
+            };
 
-        public IDisposable BeginScope<TState>(TState state) => throw new NotImplementedException();
+        public IDisposable BeginScope<TState>(TState state)
+            => throw new NotImplementedException();
     }
 }
diff --git a/MelonLoader/BootstrapInterop.cs b/MelonLoader/BootstrapInterop.cs
@@ -67,8 +67,7 @@ public static void NativeHookAttach(IntPtr target, IntPtr detour)
                 return;
 
             NativeHookAttachDirect(target, detour);
-
-            NativeStackWalk.RegisterHookAddr((ulong)detour, $"Mod-requested detour of 0x{target:X}");
+            NativeStackWalk.RegisterHookAddr((ulong)target, $"Mod-requested detour of 0x{target:X} -> 0x{detour:X}");
         }
 
         internal static unsafe void NativeHookAttachDirect(IntPtr target, IntPtr detour)
@@ -79,6 +78,7 @@ internal static unsafe void NativeHookAttachDirect(IntPtr target, IntPtr detour)
         public static unsafe void NativeHookDetach(IntPtr target, IntPtr detour)
         {
             HookDetach((void**)target, (void*)detour);
+            NativeStackWalk.UnregisterHookAddr((ulong)target);
         }
 #endif
     }

diff --git a/MelonLoader/CoreClrUtils/NativeStackWalk.cs b/MelonLoader/CoreClrUtils/NativeStackWalk.cs
@@ -455,6 +455,12 @@ internal static void RegisterHookAddr(ulong addr, string hookName)
         _registeredHooks[addr] = hookName;
     }
 
+    internal static void UnregisterHookAddr(ulong addr)
+    {
+        if (_registeredHooks.ContainsKey(addr))
+            _registeredHooks.Remove(addr);
+    }
+
     [SupportedOSPlatform("windows")]
     [Pure]
     public static List<NativeStackFrame> GetNativeStackFrames()

diff --git a/MelonLoader/MelonUtils.cs b/MelonLoader/MelonUtils.cs
@@ -557,11 +557,11 @@ internal static string GetOSVersion()
         [Obsolete("Use NativeUtils.NativeHook instead")]
         public static void NativeHookAttach(IntPtr target, IntPtr detour) => BootstrapInterop.NativeHookAttach(target, detour);
 
+        [Obsolete("Use NativeUtils.NativeHook instead")]
 #if NET6_0
         internal static void NativeHookAttachDirect(IntPtr target, IntPtr detour) => BootstrapInterop.NativeHookAttachDirect(target, detour);
 #else
         //On mono, NativeHookAttach *is* direct.
-        [Obsolete("Use NativeUtils.NativeHook instead")]
         internal static void NativeHookAttachDirect(IntPtr target, IntPtr detour) => BootstrapInterop.NativeHookAttach(target, detour);
 #endif
         [Obsolete("Use NativeUtils.NativeHook instead")]