Improve unhooking of hooked functions

Dalamud/Hooking/AsmHook.cs

@@ -35,12 +35,8 @@ public sealed class AsmHook : IDisposable, IDalamudHook
     {
         address = HookManager.FollowJmp(address);
 
-        var hasOtherHooks = HookManager.Originals.ContainsKey(address);
+        // We cannot call TrimAfterHook here because the hook is activated by the caller.
-        if (!hasOtherHooks)
+        HookManager.RegisterUnhooker(address);
-        {
-            MemoryHelper.ReadRaw(address, 0x32, out var original);
-            HookManager.Originals[address] = original;
-        }
 
         this.address = address;
         this.hookImpl = ReloadedHooks.Instance.CreateAsmHook(assembly, address.ToInt64(), (Reloaded.Hooks.Definitions.Enums.AsmHookBehaviour)asmHookBehaviour);
@@ -65,12 +61,8 @@ public sealed class AsmHook : IDisposable, IDalamudHook
     {
         address = HookManager.FollowJmp(address);
 
-        var hasOtherHooks = HookManager.Originals.ContainsKey(address);
+        // We cannot call TrimAfterHook here because the hook is activated by the caller.
-        if (!hasOtherHooks)
+        HookManager.RegisterUnhooker(address);
-        {
-            MemoryHelper.ReadRaw(address, 0x32, out var original);
-            HookManager.Originals[address] = original;
-        }
 
         this.address = address;
         this.hookImpl = ReloadedHooks.Instance.CreateAsmHook(assembly, address.ToInt64(), (Reloaded.Hooks.Definitions.Enums.AsmHookBehaviour)asmHookBehaviour);

Dalamud/Hooking/Internal/FunctionPointerVariableHook.cs

@@ -41,12 +41,7 @@ internal class FunctionPointerVariableHook<T> : Hook<T>
     {
         lock (HookManager.HookEnableSyncRoot)
         {
-            var hasOtherHooks = HookManager.Originals.ContainsKey(this.Address);
+            var unhooker = HookManager.RegisterUnhooker(this.Address);
-            if (!hasOtherHooks)
-            {
-                MemoryHelper.ReadRaw(this.Address, 0x32, out var original);
-                HookManager.Originals[this.Address] = original;
-            }
 
             if (!HookManager.MultiHookTracker.TryGetValue(this.Address, out var indexList))
             {
@@ -104,6 +99,8 @@ internal class FunctionPointerVariableHook<T> : Hook<T>
             // Add afterwards, so the hookIdent starts at 0.
             indexList.Add(this);
 
+            unhooker.TrimAfterHook();
+
             HookManager.TrackedHooks.TryAdd(Guid.NewGuid(), new HookInfo(this, detour, callingAssembly));
         }
     }

Dalamud/Hooking/Internal/HookManager.cs

@@ -16,7 +16,10 @@ namespace Dalamud.Hooking.Internal;
 [ServiceManager.EarlyLoadedService]
 internal class HookManager : IDisposable, IServiceType
 {
-    private static readonly ModuleLog Log = new("HM");
+    /// <summary>
+    /// Logger shared with <see cref="Unhooker"/>
+    /// </summary>
+    internal static readonly ModuleLog Log = new("HM");
 
     [ServiceManager.ServiceConstructor]
     private HookManager()
@@ -34,9 +37,21 @@ internal class HookManager : IDisposable, IServiceType
     internal static ConcurrentDictionary<Guid, HookInfo> TrackedHooks { get; } = new();
 
     /// <summary>
-    /// Gets a static dictionary of original code for a hooked address.
+    /// Gets a static dictionary of unhookers for a hooked address.
     /// </summary>
-    internal static ConcurrentDictionary<IntPtr, byte[]> Originals { get; } = new();
+    internal static ConcurrentDictionary<IntPtr, Unhooker> Unhookers { get; } = new();
+
+    /// <summary>
+    /// Creates a new Unhooker instance for the provided address if no such unhooker was already registered, or returns
+    /// an existing instance if the address registered previously.
+    /// </summary>
+    /// <param name="address">The address of the instruction.</param>
+    /// <returns>A new Unhooker instance.</returns>
+    public static Unhooker RegisterUnhooker(IntPtr address)
+    {
+        Log.Verbose($"Registering hook at 0x{address.ToInt64():X}");
+        return Unhookers.GetOrAdd(address, adr => new Unhooker(adr));
+    }
 
     /// <summary>
     /// Gets a static dictionary of the number of hooks on a given address.
@@ -48,7 +63,7 @@ internal class HookManager : IDisposable, IServiceType
     {
         RevertHooks();
         TrackedHooks.Clear();
-        Originals.Clear();
+        Unhookers.Clear();
     }
 
     /// <summary>
@@ -60,7 +75,7 @@ internal class HookManager : IDisposable, IServiceType
     {
         while (true)
         {
-            var hasOtherHooks = HookManager.Originals.ContainsKey(address);
+            var hasOtherHooks = HookManager.Unhookers.ContainsKey(address);
             if (hasOtherHooks)
             {
                 // This address has been hooked already. Do not follow a jmp into a trampoline of our own making.
@@ -124,28 +139,11 @@ internal class HookManager : IDisposable, IServiceType
         return address;
     }
 
-    private static unsafe void RevertHooks()
+    private static void RevertHooks()
     {
-        foreach (var (address, originalBytes) in Originals)
+        foreach (var unhooker in Unhookers.Values)
         {
-            var i = 0;
+            unhooker.Unhook();
-            var current = (byte*)address;
-            // Find how many bytes have been modified by comparing to the saved original
-            for (; i < originalBytes.Length; i++)
-            {
-                if (current[i] == originalBytes[i])
-                    break;
-            }
-
-            var snippet = originalBytes[0..i];
-
-            if (i > 0)
-            {
-                Log.Verbose($"Reverting hook at 0x{address.ToInt64():X} ({snippet.Length} bytes)");
-                MemoryHelper.ChangePermission(address, i, MemoryProtection.ExecuteReadWrite, out var oldPermissions);
-                MemoryHelper.WriteRaw(address, snippet);
-                MemoryHelper.ChangePermission(address, i, oldPermissions);
-            }
         }
     }
 }

Dalamud/Hooking/Internal/MinHookHook.cs

@@ -1,8 +1,6 @@
 using System;
 using System.Reflection;
 
-using Dalamud.Memory;
-
 namespace Dalamud.Hooking.Internal;
 
 /// <summary>
@@ -24,12 +22,7 @@ internal class MinHookHook<T> : Hook<T> where T : Delegate
     {
         lock (HookManager.HookEnableSyncRoot)
         {
-            var hasOtherHooks = HookManager.Originals.ContainsKey(this.Address);
+            var unhooker = HookManager.RegisterUnhooker(this.Address);
-            if (!hasOtherHooks)
-            {
-                MemoryHelper.ReadRaw(this.Address, 0x32, out var original);
-                HookManager.Originals[this.Address] = original;
-            }
 
             if (!HookManager.MultiHookTracker.TryGetValue(this.Address, out var indexList))
                 indexList = HookManager.MultiHookTracker[this.Address] = new();
@@ -41,6 +34,8 @@ internal class MinHookHook<T> : Hook<T> where T : Delegate
             // Add afterwards, so the hookIdent starts at 0.
             indexList.Add(this);
 
+            unhooker.TrimAfterHook();
+
             HookManager.TrackedHooks.TryAdd(Guid.NewGuid(), new HookInfo(this, detour, callingAssembly));
         }
     }

Dalamud/Hooking/Internal/ReloadedHook.cs

@@ -1,7 +1,6 @@
 using System;
 using System.Reflection;
 
-using Dalamud.Memory;
 using Reloaded.Hooks;
 
 namespace Dalamud.Hooking.Internal;
@@ -25,17 +24,13 @@ internal class ReloadedHook<T> : Hook<T> where T : Delegate
     {
         lock (HookManager.HookEnableSyncRoot)
         {
-            var hasOtherHooks = HookManager.Originals.ContainsKey(this.Address);
+            var unhooker = HookManager.RegisterUnhooker(address);
-            if (!hasOtherHooks)
-            {
-                MemoryHelper.ReadRaw(this.Address, 0x32, out var original);
-                HookManager.Originals[this.Address] = original;
-            }
-
             this.hookImpl = ReloadedHooks.Instance.CreateHook<T>(detour, address.ToInt64());
             this.hookImpl.Activate();
             this.hookImpl.Disable();
 
+            unhooker.TrimAfterHook();
+
             HookManager.TrackedHooks.TryAdd(Guid.NewGuid(), new HookInfo(this, detour, callingAssembly));
         }
     }

Dalamud/Hooking/Internal/Unhooker.cs

@@ -0,0 +1,91 @@
+using System;
+
+using Dalamud.Memory;
+
+namespace Dalamud.Hooking.Internal;
+
+/// <summary>
+/// A class which stores a copy of the bytes at a location which will be hooked in the future, such that those bytes can
+/// be restored later to "unhook" the function.
+/// </summary>
+public class Unhooker
+{
+    private readonly IntPtr address;
+    private byte[] originalBytes;
+    private bool trimmed;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="Unhooker"/> class. Upon creation, the Unhooker stores a copy of
+    /// the bytes stored at the provided address, and can be used to restore these bytes when the hook should be
+    /// removed. As such this class should be instantiated before the function is actually hooked.
+    /// </summary>
+    /// <param name="address">The address which will be hooked.</param>
+    public Unhooker(IntPtr address)
+    {
+        this.address = address;
+        MemoryHelper.ReadRaw(address, 0x32, out this.originalBytes);
+    }
+
+    /// <summary>
+    /// When called after a hook is created, checks the pre-hook original bytes and post-hook modified bytes, trimming
+    /// the original bytes stored and removing unmodified bytes from the end of the byte sequence. Assuming no
+    /// concurrent actions modified the same address space, this should result in storing only the minimum bytes
+    /// required to unhook the function.
+    /// </summary>
+    public void TrimAfterHook()
+    {
+        if (this.trimmed)
+        {
+            return;
+        }
+
+        this.originalBytes = this.originalBytes[..this.GetFullHookLength()];
+        this.trimmed = true;
+    }
+
+    /// <summary>
+    /// Attempts to unhook the function by replacing the hooked bytes with the original bytes. If
+    /// <see cref="TrimAfterHook"/> was called, the trimmed original bytes stored at that time will be used for
+    /// unhooking. Otherwise, a naive algorithm which only restores bytes until the first unchanged byte will be used in
+    /// order to avoid overwriting adjacent data.
+    /// </summary>
+    public void Unhook()
+    {
+        var len = this.trimmed ? this.originalBytes.Length : this.GetNaiveHookLength();
+        if (len > 0)
+        {
+            HookManager.Log.Verbose($"Reverting hook at 0x{this.address.ToInt64():X} ({len} bytes, trimmed={this.trimmed})");
+            MemoryHelper.ChangePermission(this.address, len, MemoryProtection.ExecuteReadWrite, out var oldPermissions);
+            MemoryHelper.WriteRaw(this.address, this.originalBytes[..len]);
+            MemoryHelper.ChangePermission(this.address, len, oldPermissions);
+        }
+    }
+
+    private unsafe int GetNaiveHookLength()
+    {
+        var current = (byte*)this.address;
+        for (var i = 0; i < this.originalBytes.Length; i++)
+        {
+            if (current[i] == this.originalBytes[i])
+            {
+                return i;
+            }
+        }
+
+        return 0;
+    }
+
+    private unsafe int GetFullHookLength()
+    {
+        var current = (byte*)this.address;
+        for (var i = this.originalBytes.Length - 1; i >= 0; i--)
+        {
+            if (current[i] != this.originalBytes[i])
+            {
+                return i + 1;
+            }
+        }
+
+        return 0;
+    }
+}