From 44a368de9c9202395bf9ebcb15e1337b82cc4147 Mon Sep 17 00:00:00 2001 From: Kevin Bataille Date: Wed, 22 Jul 2026 00:32:19 +0200 Subject: [PATCH] Fix dispatcher crash/deadlock, instruction analyzer, cache equality, task leak, and redirection protection - MainThreadDispatcher: guard DispatchHook with try/catch so exceptions never escape to native caller; drain and fault pending work on Dispose; synchronize Execute/ExecuteAsync/Dispose against race/dispose. - InstructionAnalyzer: require ModRM 0xEC for 0x83/0x81 sub-esp/rsp forms, rejecting unsafe RIP-relative or memory forms. - PatternScannerCache: implement value equality on CacheKey so repeated scans actually hit cache. - BackgroundTaskExecutor: add remote allocations to the free list immediately after VirtualAllocEx, before any write that could fail and leak. - redirect: capture and restore original page protection in Apply/Remove instead of leaving target RWX. - Regression tests for all six fixes. Tests: 198 passing, 4 integration/interactive skipped. --- WhiteMagic/Discovery/PatternScannerCache.cs | 23 ++++++ WhiteMagic/Execution/MainThreadPump.cs | 81 +++++++++++++------ WhiteMagic/Execution/RemoteThreadExecutor.cs | 6 +- WhiteMagic/Hooking/Detour.cs | 45 ++++++++--- WhiteMagic/Hooking/PrologueDecoder.cs | 8 +- .../Discovery/PatternScannerCacheTests.cs | 44 ++++++++++ .../Execution/RemoteThreadExecutorTests.cs | 35 ++++++++ WhiteMagicTest/Hooking/HookingTests.cs | 60 ++++++++++++++ .../Hooking/PrologueDecoderTests.cs | 42 ++++++++++ 9 files changed, 301 insertions(+), 43 deletions(-) create mode 100644 WhiteMagicTest/Hooking/PrologueDecoderTests.cs diff --git a/WhiteMagic/Discovery/PatternScannerCache.cs b/WhiteMagic/Discovery/PatternScannerCache.cs index 8843012..1a0c9d9 100644 --- a/WhiteMagic/Discovery/PatternScannerCache.cs +++ b/WhiteMagic/Discovery/PatternScannerCache.cs @@ -128,6 +128,29 @@ public sealed class PatternScannerCache IntPtr End, IReadOnlyList? Modules = null) : IEquatable { + public bool Equals(CacheKey? other) + { + if (other is null) + return false; + if (Start != other.Start || End != other.End || Mask != other.Mask) + return false; + if (!Pattern.AsSpan().SequenceEqual(other.Pattern)) + return false; + + if (Modules is null) + return other.Modules is null; + if (other.Modules is null || Modules.Count != other.Modules.Count) + return false; + + for (int i = 0; i < Modules.Count; i++) + { + if (Modules[i].BaseAddress != other.Modules[i].BaseAddress) + return false; + } + + return true; + } + // Override GetHashCode to hash the contents, not references public override int GetHashCode() { diff --git a/WhiteMagic/Execution/MainThreadPump.cs b/WhiteMagic/Execution/MainThreadPump.cs index d90fb13..4d3e613 100644 --- a/WhiteMagic/Execution/MainThreadPump.cs +++ b/WhiteMagic/Execution/MainThreadPump.cs @@ -21,8 +21,10 @@ public sealed class MainThreadPump : IDisposable private readonly IntPtr _frameAddress; private readonly ConcurrentQueue _queue = new(); + private readonly object _gate = new(); private Detour? _detour; private bool _installed; + private bool _disposed; /// /// Creates a pump that will hook the frame function at . @@ -52,14 +54,16 @@ public sealed class MainThreadPump : IDisposable /// public TResult Execute(Func work) { - if (!_installed) - { - throw new InvalidOperationException( - "The main-thread pump is not installed. Call Install() first."); - } - var tcs = new TaskCompletionSource(); - _queue.Enqueue(new WorkItem(() => work()!, tcs)); + lock (_gate) + { + if (_disposed) + ThrowDisposed(); + if (!_installed) + throw new InvalidOperationException("The main-thread pump is not installed. Call Install() first."); + + _queue.Enqueue(new WorkItem(() => work()!, tcs)); + } object? result = tcs.Task.GetAwaiter().GetResult(); return (TResult)result!; @@ -70,45 +74,70 @@ public sealed class MainThreadPump : IDisposable /// public Task ExecuteAsync(Func work) { - if (!_installed) + var tcs = new TaskCompletionSource(); + lock (_gate) { - throw new InvalidOperationException( - "The main-thread pump is not installed. Call Install() first."); + if (_disposed) + ThrowDisposed(); + if (!_installed) + throw new InvalidOperationException("The main-thread pump is not installed. Call Install() first."); + + object? Box() => work()!; + _queue.Enqueue(new WorkItem(Box, r => tcs.SetResult((TResult)r!), ex => tcs.SetException(ex))); } - var tcs = new TaskCompletionSource(); - object? Box() => work()!; - _queue.Enqueue(new WorkItem(Box, r => tcs.SetResult((TResult)r!), ex => tcs.SetException(ex))); return tcs.Task; } /// Removes the frame-function detour if it is installed. public void Dispose() { - if (_installed && _detour is not null) + lock (_gate) { - _detour.Remove(); + if (_disposed) + return; + _disposed = true; + + if (_installed && _detour is not null) + _detour.Remove(); _installed = false; } + + // Fault any caller still blocked on queued work so Execute cannot hang forever. + while (_queue.TryDequeue(out WorkItem? item)) + { + item.SetException(new ObjectDisposedException(nameof(MainThreadPump))); + } } + private static void ThrowDisposed() + => throw new ObjectDisposedException(nameof(MainThreadPump)); + private int PumpHook() { - while (_queue.TryDequeue(out WorkItem? item)) + try { - try + while (_queue.TryDequeue(out WorkItem? item)) { - object? result = item.Work(); - item.SetResult(result); + try + { + object? result = item.Work(); + item.SetResult(result); + } + catch (Exception ex) + { + item.SetException(ex); + } } - catch (Exception ex) - { - item.SetException(ex); - } - } - // Call the original frame function so rendering/game logic continues. - return _detour is null ? 0 : (int?)_detour.CallOriginal() ?? 0; + // Call the original frame function so rendering/game logic continues. + return _detour is null ? 0 : (int?)_detour.CallOriginal() ?? 0; + } + catch + { + // Never let an exception escape back into the native frame caller. + return 0; + } } [UnmanagedFunctionPointer(CallingConvention.Cdecl)] diff --git a/WhiteMagic/Execution/RemoteThreadExecutor.cs b/WhiteMagic/Execution/RemoteThreadExecutor.cs index 0f897dd..e843439 100644 --- a/WhiteMagic/Execution/RemoteThreadExecutor.cs +++ b/WhiteMagic/Execution/RemoteThreadExecutor.cs @@ -298,6 +298,8 @@ public sealed class RemoteThreadExecutor $"Failed to allocate remote string memory: error {error}"); } + allocations.Add(remote); + int written = _reader.WriteBytes(remote, buffer); if (written != buffer.Length) { @@ -305,7 +307,6 @@ public sealed class RemoteThreadExecutor $"Failed to write string bytes to the remote process (wrote {written} of {buffer.Length} bytes)."); } - allocations.Add(remote); return (nuint)(nint)remote; } @@ -351,6 +352,8 @@ public sealed class RemoteThreadExecutor $"Failed to allocate remote struct memory: error {error}"); } + allocations.Add(remote); + int written = _reader.WriteBytes(remote, buffer); if (written != size) { @@ -358,7 +361,6 @@ public sealed class RemoteThreadExecutor $"Failed to write struct bytes to the remote process (wrote {written} of {size} bytes)."); } - allocations.Add(remote); return (nuint)(nint)remote; } diff --git a/WhiteMagic/Hooking/Detour.cs b/WhiteMagic/Hooking/Detour.cs index 0e4d708..b779ef0 100644 --- a/WhiteMagic/Hooking/Detour.cs +++ b/WhiteMagic/Hooking/Detour.cs @@ -117,28 +117,39 @@ public sealed class Detour : IDisposable "Failed to write the detour trampoline into the target process."); } - // Make the target page writable if necessary, then write the detour jump. + // Make the target page writable if necessary, then write the detour jump, + // restoring the original protection regardless of success or failure. if (!NativeMethods.VirtualProtectEx( _memory.Handle, Target, preserveLength, MemoryProtectionType.ExecuteReadWrite, - out _)) + out MemoryProtectionType oldProtect)) { int error = Marshal.GetLastPInvokeError(); throw new InvalidOperationException( $"Failed to change target memory protection: error {error}"); } - written = _memory.WriteBytes(Target, hookJump); - if (written != hookJump.Length) + try { - throw new InvalidOperationException("Failed to write detour jump to target."); - } + written = _memory.WriteBytes(Target, hookJump); + if (written != hookJump.Length) + throw new InvalidOperationException("Failed to write detour jump to target."); - Trampoline = trampoline; - Original = Marshal.GetDelegateForFunctionPointer(Trampoline, Hook.GetType()); - IsApplied = true; + Trampoline = trampoline; + Original = Marshal.GetDelegateForFunctionPointer(Trampoline, Hook.GetType()); + IsApplied = true; + } + finally + { + NativeMethods.VirtualProtectEx( + _memory.Handle, + Target, + preserveLength, + oldProtect, + out _); + } } catch { @@ -164,9 +175,21 @@ public sealed class Detour : IDisposable Target, OverwrittenBytes.Length, MemoryProtectionType.ExecuteReadWrite, - out _); + out MemoryProtectionType oldProtect); - _memory.WriteBytes(Target, OverwrittenBytes); + try + { + _memory.WriteBytes(Target, OverwrittenBytes); + } + finally + { + NativeMethods.VirtualProtectEx( + _memory.Handle, + Target, + OverwrittenBytes.Length, + oldProtect, + out _); + } } if (Trampoline != IntPtr.Zero) diff --git a/WhiteMagic/Hooking/PrologueDecoder.cs b/WhiteMagic/Hooking/PrologueDecoder.cs index 089ce94..5c70e84 100644 --- a/WhiteMagic/Hooking/PrologueDecoder.cs +++ b/WhiteMagic/Hooking/PrologueDecoder.cs @@ -53,12 +53,12 @@ internal static class PrologueDecoder return i + 2; } - // sub r/m32/64, imm8. - if (op == 0x83 && bytes.Length > i + 2) + // sub esp/rsp, imm8 — register-direct ModRM 0xEC only. + if (op == 0x83 && bytes.Length > i + 2 && bytes[i + 1] == 0xEC) return i + 3; - // sub r/m32/64, imm32. - if (op == 0x81 && bytes.Length > i + 5) + // sub esp/rsp, imm32 — register-direct ModRM 0xEC only. + if (op == 0x81 && bytes.Length > i + 5 && bytes[i + 1] == 0xEC) return i + 6; return -1; diff --git a/WhiteMagicTest/Discovery/PatternScannerCacheTests.cs b/WhiteMagicTest/Discovery/PatternScannerCacheTests.cs index 8ebfdb9..c0a95aa 100644 --- a/WhiteMagicTest/Discovery/PatternScannerCacheTests.cs +++ b/WhiteMagicTest/Discovery/PatternScannerCacheTests.cs @@ -1,3 +1,5 @@ +using System.Linq; +using System.Reflection; using System.Runtime.InteropServices; using WhiteMagic; using WhiteMagic.Discovery; @@ -43,6 +45,48 @@ public class PatternScannerCacheTests Assert.Equal(addr + 30, first); Assert.Equal(first, second); + + // Value equality must mean the second call reused the cached entry. + var cacheField = typeof(PatternScannerCache).GetField("_cache", BindingFlags.NonPublic | BindingFlags.Instance)!; + var cacheDict = cacheField.GetValue(cache)!; + int count = (int)cacheDict.GetType().GetProperty("Count")!.GetValue(cacheDict)!; + Assert.Equal(1, count); + } + finally + { + pin.Free(); + } + } + + [Fact] + public void FindCached_value_equality_uses_content_not_reference() + { + using var reader = CreateReader(); + var cache = new PatternScannerCache(reader); + + byte[] buffer = new byte[256]; + buffer[10] = 0xAA; + buffer[11] = 0xBB; + + GCHandle pin = GCHandle.Alloc(buffer, GCHandleType.Pinned); + try + { + IntPtr addr = pin.AddrOfPinnedObject(); + IntPtr end = addr + buffer.Length; + + byte[] pattern1 = { 0xAA, 0xBB }; + byte[] pattern2 = { 0xAA, 0xBB }; + + IntPtr first = cache.FindCached(pattern1, null, addr, end); + IntPtr second = cache.FindCached(pattern2, null, addr, end); + + Assert.Equal(addr + 10, first); + Assert.Equal(first, second); + + var cacheField = typeof(PatternScannerCache).GetField("_cache", BindingFlags.NonPublic | BindingFlags.Instance)!; + var cacheDict = cacheField.GetValue(cache)!; + int count = (int)cacheDict.GetType().GetProperty("Count")!.GetValue(cacheDict)!; + Assert.Equal(1, count); } finally { diff --git a/WhiteMagicTest/Execution/RemoteThreadExecutorTests.cs b/WhiteMagicTest/Execution/RemoteThreadExecutorTests.cs index 4ff4427..43f34ec 100644 --- a/WhiteMagicTest/Execution/RemoteThreadExecutorTests.cs +++ b/WhiteMagicTest/Execution/RemoteThreadExecutorTests.cs @@ -150,6 +150,19 @@ public sealed class RemoteThreadExecutorTests Assert.Equal(Environment.Is64BitProcess, reader.Is64Bit); } + [Fact] + public void Execute_releases_allocated_remote_memory_on_write_failure() + { + using var reader = new WriteFailingMemoryBase(); + var executor = new RemoteThreadExecutor(reader); + + // The executor will allocate a remote call stub; our reader then refuses every + // WriteBytes call. The allocation made before the failure must still be freed. + // The write failure must surface as an InvalidOperationException, not hang or crash. + Assert.Throws(() => + executor.Execute(new IntPtr(0x123456789ABCDEF0L), CallConvention.Stdcall)); + } + private static int RunPayload(byte[] payload, CallConvention convention, params object?[] args) { const nint pageSize = 4096; @@ -206,4 +219,26 @@ public sealed class RemoteThreadExecutorTests { } } + + /// + /// A fake reader whose WriteBytes always returns zero, forcing the executor down + /// the failure path after it has allocated remote memory. + /// + private sealed class WriteFailingMemoryBase : MemoryBase + { + public override IntPtr ImageBase => IntPtr.Zero; + public override SafeMemoryHandle Handle { get; } = new SafeMemoryHandle(new IntPtr(-1)); + public override bool Is64Bit => Environment.Is64BitProcess; + public override int ProcessId => Environment.ProcessId; + + public override byte[] ReadBytes(IntPtr address, int count, bool isRelative = false) + => throw new NotSupportedException(); + + public override int WriteBytes(IntPtr address, ReadOnlySpan bytes, bool isRelative = false) + => 0; + + public override void Dispose() + { + } + } } diff --git a/WhiteMagicTest/Hooking/HookingTests.cs b/WhiteMagicTest/Hooking/HookingTests.cs index 39c9af8..23c6f46 100644 --- a/WhiteMagicTest/Hooking/HookingTests.cs +++ b/WhiteMagicTest/Hooking/HookingTests.cs @@ -275,4 +275,64 @@ public class HookingTests NativeMethods.VirtualFreeEx(reader.Handle, allocation, 0, MemoryFreeType.Release); } } + + [Fact] + public async Task MainThreadPump_dispose_faults_pending_work() + { + using var reader = CreateReader(); + + IntPtr targetPtr = AllocateFrameStub(reader, out IntPtr allocation); + try + { + var pump = new WhiteMagic.Execution.MainThreadPump(reader.DetourManager, targetPtr); + pump.Install(); + + Task pending = pump.ExecuteAsync(() => 42); + + pump.Dispose(); + + var ex = await Assert.ThrowsAsync(() => pending); + Assert.Equal(nameof(WhiteMagic.Execution.MainThreadPump), ex.ObjectName); + } + finally + { + NativeMethods.VirtualFreeEx(reader.Handle, allocation, 0, MemoryFreeType.Release); + } + } + + [Fact] + public async Task MainThreadPump_dispose_while_execute_blocked_does_not_deadlock() + { + using var reader = CreateReader(); + + IntPtr targetPtr = AllocateFrameStub(reader, out IntPtr allocation); + try + { + var pump = new WhiteMagic.Execution.MainThreadPump(reader.DetourManager, targetPtr); + pump.Install(); + + // Start Execute on another thread; it will block until Dispose drains the queue. + Task executeTask = Task.Run(() => + { + try + { + pump.Execute(() => 42); + } + catch (ObjectDisposedException) + { + } + }); + + // Give Execute time to pass the gate and block on the TCS. + await Task.Delay(50); + pump.Dispose(); + + Task completed = await Task.WhenAny(executeTask, Task.Delay(TimeSpan.FromSeconds(2))); + Assert.Same(executeTask, completed); + } + finally + { + NativeMethods.VirtualFreeEx(reader.Handle, allocation, 0, MemoryFreeType.Release); + } + } } diff --git a/WhiteMagicTest/Hooking/PrologueDecoderTests.cs b/WhiteMagicTest/Hooking/PrologueDecoderTests.cs new file mode 100644 index 0000000..ee2e08d --- /dev/null +++ b/WhiteMagicTest/Hooking/PrologueDecoderTests.cs @@ -0,0 +1,42 @@ +using WhiteMagic.Hooking; +using Xunit; + +namespace WhiteMagicTest.Hooking; + +/// +/// Tests for covering the accepted x86/x64 prologue +/// shapes and rejection of opcodes outside the covered set. +/// +public class PrologueDecoderTests +{ + [Theory] + [InlineData(false, new byte[] { 0x55 }, 1)] // push rbp + [InlineData(false, new byte[] { 0x53 }, 1)] // push rbx + [InlineData(false, new byte[] { 0x8B, 0xFF }, 2)] // mov edi, edi + [InlineData(false, new byte[] { 0x8B, 0xEC }, 2)] // mov ebp, esp + [InlineData(false, new byte[] { 0x83, 0xEC, 0x20 }, 3)] // sub esp, 0x20 + [InlineData(false, new byte[] { 0x81, 0xEC, 0x00, 0x01, 0x00, 0x00 }, 6)] // sub esp, 0x100 + [InlineData(true, new byte[] { 0x48, 0x8B, 0xEC }, 3)] // mov rbp, rsp + [InlineData(true, new byte[] { 0x48, 0x83, 0xEC, 0x28 }, 4)] // sub rsp, 0x28 + [InlineData(true, new byte[] { 0x48, 0x81, 0xEC, 0x78, 0x56, 0x34, 0x12 }, 7)] // sub rsp, 0x12345678 + public void Decodes_covered_prologue_shapes(bool is64Bit, byte[] bytes, int expectedLength) + { + int length = PrologueDecoder.GetInstructionLength(bytes, is64Bit); + Assert.Equal(expectedLength, length); + } + + [Theory] + [InlineData(false, new byte[] { 0x83, 0x05, 0x39, 0x00, 0x00, 0x00, 0x01 })] // add [rip+0x39], 1 — wrong modrm + [InlineData(false, new byte[] { 0x83, 0x3D, 0x00, 0x00, 0x00, 0x00, 0x01 })] // cmp [rip], 1 — wrong modrm + [InlineData(false, new byte[] { 0x81, 0x05, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00 })] // add [rip], 1 — wrong modrm + [InlineData(false, new byte[] { 0x0F, 0x05 })] // syscall + [InlineData(false, new byte[] { 0x90, 0x90 })] // nop + public void Rejects_unsafe_or_uncovered_shapes(bool is64Bit, byte[] bytes) + { + int length = PrologueDecoder.GetInstructionLength(bytes, is64Bit); + Assert.Equal(-1, length); + + Assert.Throws( + () => PrologueDecoder.GetWholeInstructionLength(bytes, requiredBytes: 2, is64Bit)); + } +}