From 614806f5f152f779d6b0eb64e95968bde898855e Mon Sep 17 00:00:00 2001
From: Kevin Bataille <kbe@kev1.net>
Date: Wed, 22 Jul 2026 01:24:33 +0200
Subject: [PATCH] Fix LibraryLoader crash on error during context-transfer
 restoration

- Track whether transferred thread's original context was successfully restored.
- In the catch block, leak the remote allocation instead of freeing it if the thread was not restored; this prevents the target process from executing freed memory.
- Remove redundant 'op == 0x55' check in InstructionAnalyzer (already matched by (op & 0xF8) == 0x50).
- Simplify MemoryBase ReadString align-down expression to previousLen - (previousLen % nullLen).
- Change StubAllocator size parameter from nint to int for clarity (internal test seam).

Tests: 199 passing, 4 integration/interactive skipped.
---
 WhiteMagic/Execution/RemoteThreadExecutor.cs |  2 +-
 WhiteMagic/Hooking/PrologueDecoder.cs        |  4 ++--
 WhiteMagic/Injection/DllInjector.cs          | 17 +++++++++++++++--
 WhiteMagic/MemoryBase.cs                     |  4 ++--
 4 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/WhiteMagic/Execution/RemoteThreadExecutor.cs b/WhiteMagic/Execution/RemoteThreadExecutor.cs
index d43c933..b584c1a 100644
--- a/WhiteMagic/Execution/RemoteThreadExecutor.cs
+++ b/WhiteMagic/Execution/RemoteThreadExecutor.cs
@@ -41,7 +41,7 @@ public sealed class RemoteThreadExecutor
     /// When this delegate returns a non-zero pointer, the executor does not take
     /// ownership of that memory and will not free it.
     /// </remarks>
-    internal Func<IntPtr, nint, IntPtr>? StubAllocator { get; set; }
+    internal Func<IntPtr, int, IntPtr>? StubAllocator { get; set; }
 
     /// <summary>Test seam: overrides remote scratch allocation for string/struct args.
     /// Defaults to <see cref="NativeMethods.VirtualAllocEx"/>.</summary>
diff --git a/WhiteMagic/Hooking/PrologueDecoder.cs b/WhiteMagic/Hooking/PrologueDecoder.cs
index 5c70e84..92fd661 100644
--- a/WhiteMagic/Hooking/PrologueDecoder.cs
+++ b/WhiteMagic/Hooking/PrologueDecoder.cs
@@ -41,8 +41,8 @@ internal static class PrologueDecoder
 
         byte op = bytes[i];
 
-        // push reg / push rbp.
-        if ((op & 0xF8) == 0x50 || op == 0x55)
+        // push reg (0x50-0x57), including rbp (0x55).
+        if ((op & 0xF8) == 0x50)
             return i + 1;
 
         // mov r32/64, r/m32/64. Recognize only the specific forms listed above.
diff --git a/WhiteMagic/Injection/DllInjector.cs b/WhiteMagic/Injection/DllInjector.cs
index ae26145..2d9f199 100644
--- a/WhiteMagic/Injection/DllInjector.cs
+++ b/WhiteMagic/Injection/DllInjector.cs
@@ -209,6 +209,8 @@ public sealed class DllInjector
                 throw new InvalidOperationException($"SuspendThread failed (error {error}).");
             }
 
+            bool restored = false;
+
             try
             {
                 IntPtr result;
@@ -257,6 +259,8 @@ public sealed class DllInjector
                         int error = Marshal.GetLastPInvokeError();
                         throw new InvalidOperationException($"ResumeThread restore failed (error {error}).");
                     }
+
+                    restored = true;
                 }
                 else
                 {
@@ -302,6 +306,8 @@ public sealed class DllInjector
                         int error = Marshal.GetLastPInvokeError();
                         throw new InvalidOperationException($"ResumeThread restore failed (error {error}).");
                     }
+
+                    restored = true;
                 }
 
                 if (result == IntPtr.Zero)
@@ -311,8 +317,15 @@ public sealed class DllInjector
             }
             catch
             {
-                // Best effort: resume the thread if we left it suspended.
-                _ = NativeMethods.ResumeThread(thread);
+                // If we never successfully restored the thread's original context, the
+                // thread may still be executing (or about to execute) code inside the
+                // injected allocation. Freeing that memory now would crash the target
+                // process, so leak the block and leave the thread suspended.
+                if (!restored)
+                {
+                    remoteBase = IntPtr.Zero;
+                }
+
                 throw;
             }
         }
diff --git a/WhiteMagic/MemoryBase.cs b/WhiteMagic/MemoryBase.cs
index 9e59822..f7c864f 100644
--- a/WhiteMagic/MemoryBase.cs
+++ b/WhiteMagic/MemoryBase.cs
@@ -211,8 +211,8 @@ public abstract class MemoryBase : IDisposable
             // Search the newly extended buffer at code-unit-aligned positions. A terminator
             // can start as far back as (nullLen - 1) bytes before the new bytes, so start
             // the search just before the previous end, rounded up to the next code-unit.
-            int firstAligned = ((previousLen - nullLen + 1 + nullLen - 1) / nullLen) * nullLen;
-            firstAligned = Math.Max(0, firstAligned);
+            int firstAligned = previousLen - (previousLen % nullLen);
+            if (firstAligned < 0) firstAligned = 0;
 
             int limit = accumulated.Count - nullLen;
             for (int i = firstAligned; i <= limit; i += nullLen)