function call trace is more advanced

Eddio0141 · Nov 25, 2024 · 24f1a7f · 24f1a7f
1 parent 30ea119
commit 24f1a7f
Show file tree

Hide file tree

Showing 5 changed files with 136 additions and 47 deletions.
diff --git a/UniTAS/Patcher/Implementations/Customization/Config.Section.cs b/UniTAS/Patcher/Implementations/Customization/Config.Section.cs
@@ -8,17 +8,18 @@ public static class Debug
         {
             public static class FunctionCallTrace
             {
-                public const string SECTION_NAME = $"{nameof(Debug)}.FunctionCallTrace";
-                public const string ENABLE = "Enable";
-                public const string MATCHING_TYPES = "MatchingTypes";
+                public const string SectionName = $"{nameof(Debug)}.FunctionCallTrace";
+                public const string Enable = "Enable";
+                public const string Methods = "Methods";
             }
         }
+
         public static class Remote
         {
-            public const string SECTION_NAME = $"{nameof(Remote)}";
-            public const string ENABLE = "Enable";
-            public const string ADDRESS = "Address";
-            public const string PORT = "Port";
+            public const string SectionName = $"{nameof(Remote)}";
+            public const string Enable = "Enable";
+            public const string Address = "Address";
+            public const string Port = "Port";
         }
     }
 }
diff --git a/UniTAS/Patcher/Implementations/Customization/Config.cs b/UniTAS/Patcher/Implementations/Customization/Config.cs
@@ -48,12 +48,12 @@ public Config(ILogger logger)
         }
 
         // add entry for patcher config entry
-        BepInExConfigFile.Bind(Sections.Debug.FunctionCallTrace.SECTION_NAME, Sections.Debug.FunctionCallTrace.ENABLE,
+        BepInExConfigFile.Bind(Sections.Debug.FunctionCallTrace.SectionName, Sections.Debug.FunctionCallTrace.Enable,
             false,
             "If enabled, will hook on most functions and log every function call");
-        BepInExConfigFile.Bind(Sections.Debug.FunctionCallTrace.SECTION_NAME,
-            Sections.Debug.FunctionCallTrace.MATCHING_TYPES, "*",
-            "A list of glob pattern of types to hook function call tracing to. You can append to the list by separating each entry with a comma. Example: `UnityEngine.Application, UnityEngine.Time, UnityEngine.InputSystem.*`");
+        BepInExConfigFile.Bind(Sections.Debug.FunctionCallTrace.SectionName,
+            Sections.Debug.FunctionCallTrace.Methods, "",
+            "A list of methods to hook call tracing to. Each entry must consist of `Type:method`. You can append to the list by separating each entry with a comma. Example: `UnityEngine.Time:get_captureDeltaTime, Game.HeadCrab:Update`");
 
         var backendConfigRaw = new List<string>();
         try

diff --git a/UniTAS/Patcher/Implementations/PreloadPatcherProcessor.cs b/UniTAS/Patcher/Implementations/PreloadPatcherProcessor.cs
@@ -11,8 +11,8 @@ public class PreloadPatcherProcessor
         new StaticCtorHeaders(),
         new UnityInitInvoke(),
         new FinalizeSuppressionPatch(),
-        new FunctionCallTrace(),
         new SteamAPIPatch(),
-        new SerializationCallbackPatch()
+        new SerializationCallbackPatch(),
+        new FunctionCallTrace() // this must run last, it hooks logs on start and ret
     ];
 }
diff --git a/UniTAS/Patcher/Implementations/RemoteControl.cs b/UniTAS/Patcher/Implementations/RemoteControl.cs
@@ -32,11 +32,11 @@ public RemoteControl(IConfig config, ILogger logger, ILiveScripting liveScriptin
     {
         var conf = config.BepInExConfigFile;
 
-        var enableConf = conf.Bind(Config.Sections.Remote.SECTION_NAME, Config.Sections.Remote.ENABLE, false,
+        var enableConf = conf.Bind(Config.Sections.Remote.SectionName, Config.Sections.Remote.Enable, false,
             "If enabled, starts a TCP/IP connection for remote controlling UniTAS");
-        var ipAddrConf = conf.Bind(Config.Sections.Remote.SECTION_NAME, Config.Sections.Remote.ADDRESS, "127.0.0.1",
+        var ipAddrConf = conf.Bind(Config.Sections.Remote.SectionName, Config.Sections.Remote.Address, "127.0.0.1",
             "IP address to use for remote connection");
-        var portConf = conf.Bind(Config.Sections.Remote.SECTION_NAME, Config.Sections.Remote.PORT, 8080,
+        var portConf = conf.Bind(Config.Sections.Remote.SectionName, Config.Sections.Remote.Port, 8080,
             "Port to use for remote connection");
 
         if (!enableConf.Value) return;

diff --git a/UniTAS/Patcher/Patches/Preloader/FunctionCallTrace.cs b/UniTAS/Patcher/Patches/Preloader/FunctionCallTrace.cs
@@ -1,12 +1,15 @@
+using System;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
 using System.Linq;
+using System.Threading;
 using BepInEx.Logging;
 using HarmonyLib;
 using Mono.Cecil;
 using Mono.Cecil.Cil;
 using Mono.Cecil.Rocks;
+using MonoMod.Utils;
 using UniTAS.Patcher.Extensions;
 using UniTAS.Patcher.Implementations.Customization;
 using UniTAS.Patcher.Interfaces;
@@ -25,68 +28,153 @@ public override void Patch(ref AssemblyDefinition assembly)
         // check config if loading
         var unitasConfig = File.ReadAllText(UniTASPaths.ConfigBepInEx);
 
-        var funcTraceEnable = ConfigUtils.GetEntryKey(unitasConfig, Config.Sections.Debug.FunctionCallTrace.SECTION_NAME,
-            Config.Sections.Debug.FunctionCallTrace.ENABLE);
+        var funcTraceEnable = ConfigUtils.GetEntryKey(unitasConfig,
+            Config.Sections.Debug.FunctionCallTrace.SectionName,
+            Config.Sections.Debug.FunctionCallTrace.Enable);
 
         if (funcTraceEnable != "true") return;
 
         StaticLogger.LogDebug("Hooking functions for trace logging");
 
-        var methodInvokeTrace = AccessTools.Method(typeof(FunctionCallTrace), nameof(MethodInvokeTrace));
-        var methodInvokeTraceRef = assembly.MainModule.ImportReference(methodInvokeTrace);
+        var methodInvokeTraceStart = AccessTools.Method(typeof(FunctionCallTrace), nameof(MethodInvokeStart));
+        var methodInvokeTraceStartRef = assembly.MainModule.ImportReference(methodInvokeTraceStart);
+        var methodInvokeTraceEnd = AccessTools.Method(typeof(FunctionCallTrace), nameof(MethodInvokeEnd));
+        var methodInvokeTraceEndRef = assembly.MainModule.ImportReference(methodInvokeTraceEnd);
 
         var types = assembly.Modules.SelectMany(module => module.GetAllTypes());
-        var matchingTypesConfig = ConfigUtils.GetEntryKey(unitasConfig, Config.Sections.Debug.FunctionCallTrace.SECTION_NAME, Config.Sections.Debug.FunctionCallTrace.MATCHING_TYPES);
+        var methodsHookConfig = ConfigUtils.GetEntryKey(unitasConfig,
+            Config.Sections.Debug.FunctionCallTrace.SectionName,
+            Config.Sections.Debug.FunctionCallTrace.Methods);
 
-        if (matchingTypesConfig != null && matchingTypesConfig.Trim() != "*")
-        {
-            var matchingTypes = matchingTypesConfig.Split([','], System.StringSplitOptions.RemoveEmptyEntries).Select(x => x.Trim()).ToList();
+        if (methodsHookConfig == null || methodsHookConfig.Trim() == "") return;
 
-            types = types.Where(t => matchingTypes.Any(x => t.FullName.Like(x)));
-        }
+        var hookTypeAndMethods = methodsHookConfig.Split([','], StringSplitOptions.RemoveEmptyEntries)
+            .Select(x => x.Trim().Split([":"], StringSplitOptions.None)).Where(x => x.Length == 2)
+            .Select(x => (x[0], x[1])).GroupBy(x => x.Item1)
+            .ToDictionary(x => x.Key, x => new HashSet<string>(x.Select(y => y.Item2)));
+
+        types = types.Where(t => hookTypeAndMethods.ContainsKey(t.FullName));
 
+        var hookedMethods = new HashSet<MethodDefinition>();
         foreach (var type in types)
         {
-            if (type.IsEnum || type.IsInterface) continue;
+            var typeFullName = type.FullName;
 
-            StaticLogger.LogDebug($"hooking on class {type.FullName}");
+            if (type.IsEnum || type.IsInterface)
+            {
+                StaticLogger.LogWarning(
+                    $"Function call trace: skipping hook on type that is enum, or interface: {typeFullName}");
+                continue;
+            }
+
+            StaticLogger.LogDebug($"hooking on class {typeFullName}");
+
+            var methodsToHook = hookTypeAndMethods[typeFullName];
 
             foreach (var method in type.Methods)
             {
-                if (!method.HasBody) continue;
-
-                StaticLogger.LogDebug($"hooking on method {method.Name}");
+                if (!methodsToHook.Contains(method.Name)) continue;
+                if (!method.HasBody)
+                {
+                    StaticLogger.LogWarning($"Function call trace: cannot hook method {method.Name}, it has no body");
+                    continue;
+                }
 
                 var body = method.Body;
-                body.SimplifyMacros();
-
+                // start hook
                 var il = body.GetILProcessor();
+                il.InsertBefore(body.Instructions.First(), il.Create(OpCodes.Call, methodInvokeTraceStartRef));
 
-                var methodInvokeTraceInstruction = il.Create(OpCodes.Call, methodInvokeTraceRef);
+                // end hooks
+                HookOnMethod(method, methodInvokeTraceEndRef, hookedMethods);
+            }
+        }
+    }
 
-                il.InsertBefore(body.Instructions.First(), methodInvokeTraceInstruction);
+    private static void HookOnMethod(MethodDefinition target, MethodReference methodInvokeTraceEndRef,
+        HashSet<MethodDefinition> alreadyHooked)
+    {
+        if (!target.HasBody || alreadyHooked.Contains(target)) return;
+        alreadyHooked.Add(target);
 
-                body.OptimizeMacros();
-            }
+        StaticLogger.LogDebug($"Function call trace: hooking on method {target.FullName}");
+
+        var body = target.Body;
+        body.SimplifyMacros();
+
+        var il = body.GetILProcessor();
+
+        // calls
+        foreach (var inst in body.Instructions)
+        {
+            if (inst.Operand is not MethodReference method) continue;
+            // TODO: for now, we exclude System namespace
+            var ns = method.DeclaringType.Namespace;
+            if (ns == "System" || ns.StartsWith("System.")) continue;
+
+            var methodDef = method.SafeResolve();
+            if (methodDef == null) continue;
+            HookOnMethod(methodDef, methodInvokeTraceEndRef, alreadyHooked);
         }
+
+        // ends
+        var methodInvokeTraceEndInstruction = il.Create(OpCodes.Call, methodInvokeTraceEndRef);
+        for (var i = 0; i < body.Instructions.Count; i++)
+        {
+            var inst = body.Instructions[i];
+            if (inst.OpCode != OpCodes.Ret) continue;
+            il.InsertBefore(inst, methodInvokeTraceEndInstruction);
+            i++;
+        }
+
+        body.OptimizeMacros();
     }
 
-    public static void MethodInvokeTrace()
+    private static readonly ThreadLocal<(int startDepth, int depth, string path)> MethodTracePath = new();
+
+    public static void MethodInvokeStart()
     {
-        var stack = new StackTrace();
-        var frame = stack.GetFrame(1);
+        if (MethodTracePath.Value.startDepth > 0) return;
+        MethodTracePath.Value = (new StackTrace().FrameCount - 1, default, default);
+    }
 
-        if (frame == null)
+    public static void MethodInvokeEnd()
+    {
+        var stack = new StackTrace();
+        var frames = stack.GetFrames();
+        if (frames == null)
         {
-            LOG.LogWarning("trace data is empty, cannot trace call");
+            StaticLogger.LogWarning("Stack trace data is empty, cannot trace call");
             return;
         }
 
-        var space = stack.FrameCount - 2;
-        var func = frame.GetMethod();
+        var frameCount = frames.Length;
+        var prevDepth = MethodTracePath.Value.depth;
+        foreach (var frame in frames.Skip(1 + prevDepth).Reverse())
+        {
+            var func = frame.GetMethod();
+            var funcNameFull = $"{func.DeclaringType.SaneFullName()}.{func.Name}";
+            if (MethodTracePath.Value.path != null)
+            {
+                funcNameFull = $"{MethodTracePath.Value.path} -> {funcNameFull}";
+            }
+
+            // -2, since 1 for this method, and another since we are handling the return trace
+            MethodTracePath.Value = (MethodTracePath.Value.startDepth, 0, funcNameFull);
+        }
 
-        var funcNameFull = $"{func.DeclaringType.SaneFullName()}.{func.Name}";
+        // add return, if anything to be returned to
+        if (frameCount - MethodTracePath.Value.startDepth > 1)
+        {
+            var retMethod = frames[2].GetMethod();
+            var funcNameFull = $" <- {retMethod.DeclaringType.SaneFullName()}.{retMethod.Name}";
+            MethodTracePath.Value = (MethodTracePath.Value.startDepth, frameCount - 2,
+                MethodTracePath.Value.path + funcNameFull);
+            return;
+        }
 
-        LOG.LogInfo($"{new string(' ', space)}{funcNameFull}");
+        // dump the result
+        LOG.LogInfo(MethodTracePath.Value.path);
+        MethodTracePath.Value = default;
     }
-}
+}