From 1911514120ce9ba8f79356642a356f4eb6616246 Mon Sep 17 00:00:00 2001 From: Kevin Bataille Date: Wed, 22 Jul 2026 02:16:23 +0200 Subject: [PATCH] Fix bounds, memory protection, and completion race in core helpers - AllocatedMemory.Read/Write/ReadBytes/WriteBytes now validate that the requested byte range stays within the allocated block before calling into the memory accessor. - Patch.Apply/Remove temporarily changes the target page to read-write and restores the original protection, mirroring the Detour behavior. - MainThreadPump.WorkItem uses TrySetResult/TrySetException and swallows the InvalidOperationException raised when a completion source is already completed, preventing Dispose from failing during concurrent pump drainage. Regression tests added for all three fixes. Tests: 206 passing, 4 skipped. --- WhiteMagic/Execution/MainThreadPump.cs | 24 ++++- WhiteMagic/Hooking/Patch.cs | 54 ++++++++++- WhiteMagic/Memory/AllocatedMemory.cs | 39 ++++++-- WhiteMagicTest/Hooking/HookingTests.cs | 91 +++++++++++++++++++ WhiteMagicTest/Memory/AllocatedMemoryTests.cs | 69 ++++++++++++++ 5 files changed, 261 insertions(+), 16 deletions(-) diff --git a/WhiteMagic/Execution/MainThreadPump.cs b/WhiteMagic/Execution/MainThreadPump.cs index 4d3e613..ce3f824 100644 --- a/WhiteMagic/Execution/MainThreadPump.cs +++ b/WhiteMagic/Execution/MainThreadPump.cs @@ -83,7 +83,7 @@ public sealed class MainThreadPump : IDisposable 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))); + _queue.Enqueue(new WorkItem(Box, r => tcs.TrySetResult((TResult)r!), ex => tcs.TrySetException(ex))); } return tcs.Task; @@ -158,20 +158,34 @@ public sealed class MainThreadPump : IDisposable public WorkItem(Func work, TaskCompletionSource tcs) { Work = work; - _setResult = r => tcs.SetResult(r); - _setException = ex => tcs.SetException(ex); + _setResult = r => tcs.TrySetResult(r); + _setException = ex => tcs.TrySetException(ex); } public Func Work { get; } public void SetResult(object? result) { - _setResult?.Invoke(result); + try + { + _setResult?.Invoke(result); + } + catch (InvalidOperationException) + { + // Already completed, e.g. concurrent Dispose/PumpHook race. + } } public void SetException(Exception exception) { - _setException?.Invoke(exception); + try + { + _setException?.Invoke(exception); + } + catch (InvalidOperationException) + { + // Already completed, e.g. concurrent Dispose/PumpHook race. + } } } } diff --git a/WhiteMagic/Hooking/Patch.cs b/WhiteMagic/Hooking/Patch.cs index 19a6184..8c58f0e 100644 --- a/WhiteMagic/Hooking/Patch.cs +++ b/WhiteMagic/Hooking/Patch.cs @@ -1,5 +1,7 @@ using System; using System.Linq; +using System.Runtime.InteropServices; +using WhiteMagic.Native; namespace WhiteMagic.Hooking; @@ -53,7 +55,31 @@ public sealed class Patch : IDisposable return; OriginalBytes = _memory.ReadBytes(Address, PatchBytes.Length); - _memory.WriteBytes(Address, PatchBytes); + + if (!NativeMethods.VirtualProtectEx( + _memory.Handle, + Address, + PatchBytes.Length, + MemoryProtectionType.ExecuteReadWrite, + out MemoryProtectionType oldProtect)) + { + int error = Marshal.GetLastPInvokeError(); + throw new InvalidOperationException($"Failed to change target memory protection: error {error}"); + } + + try + { + _memory.WriteBytes(Address, PatchBytes); + } + finally + { + NativeMethods.VirtualProtectEx( + _memory.Handle, + Address, + PatchBytes.Length, + oldProtect, + out _); + } } /// Restores the original bytes if they were captured. @@ -62,7 +88,31 @@ public sealed class Patch : IDisposable if (OriginalBytes is null) return; - _memory.WriteBytes(Address, OriginalBytes); + if (!NativeMethods.VirtualProtectEx( + _memory.Handle, + Address, + OriginalBytes.Length, + MemoryProtectionType.ExecuteReadWrite, + out MemoryProtectionType oldProtect)) + { + int error = Marshal.GetLastPInvokeError(); + throw new InvalidOperationException($"Failed to change target memory protection: error {error}"); + } + + try + { + _memory.WriteBytes(Address, OriginalBytes); + } + finally + { + NativeMethods.VirtualProtectEx( + _memory.Handle, + Address, + OriginalBytes.Length, + oldProtect, + out _); + } + OriginalBytes = null; } diff --git a/WhiteMagic/Memory/AllocatedMemory.cs b/WhiteMagic/Memory/AllocatedMemory.cs index 1d31ba8..2d3a5b1 100644 --- a/WhiteMagic/Memory/AllocatedMemory.cs +++ b/WhiteMagic/Memory/AllocatedMemory.cs @@ -86,12 +86,18 @@ public sealed class AllocatedMemory : IDisposable public IntPtr AddressOf(string name) { ObjectDisposedException.ThrowIf(_disposed, this); + int offset = GetRegionOffset(name); + return _baseAddress + offset; + } + + private int GetRegionOffset(string name) + { ArgumentNullException.ThrowIfNull(name); if (!_regions.TryGetValue(name, out int offset)) throw new ArgumentException($"Region '{name}' does not exist.", nameof(name)); - return _baseAddress + offset; + return offset; } /// @@ -105,8 +111,12 @@ public sealed class AllocatedMemory : IDisposable { ObjectDisposedException.ThrowIf(_disposed, this); - IntPtr address = AddressOf(name); - return _memory.Read(address); + int offset = GetRegionOffset(name); + int size = Marshal.SizeOf(); + if (offset > _size - size) + throw new ArgumentOutOfRangeException(nameof(name), $"Region '{name}' read of {size} bytes exceeds allocation size {_size}."); + + return _memory.Read(_baseAddress + offset); } /// @@ -121,8 +131,12 @@ public sealed class AllocatedMemory : IDisposable { ObjectDisposedException.ThrowIf(_disposed, this); - IntPtr address = AddressOf(name); - return _memory.Write(address, value); + int offset = GetRegionOffset(name); + int size = Marshal.SizeOf(); + if (offset > _size - size) + throw new ArgumentOutOfRangeException(nameof(name), $"Region '{name}' write of {size} bytes exceeds allocation size {_size}."); + + return _memory.Write(_baseAddress + offset, value); } /// @@ -136,8 +150,12 @@ public sealed class AllocatedMemory : IDisposable { ObjectDisposedException.ThrowIf(_disposed, this); - IntPtr address = AddressOf(name); - return _memory.ReadBytes(address, count); + int offset = GetRegionOffset(name); + ArgumentOutOfRangeException.ThrowIfNegative(count); + if (offset > _size - count) + throw new ArgumentOutOfRangeException(nameof(count), $"Region '{name}' read of {count} bytes exceeds allocation size {_size}."); + + return _memory.ReadBytes(_baseAddress + offset, count); } /// @@ -151,8 +169,11 @@ public sealed class AllocatedMemory : IDisposable { ObjectDisposedException.ThrowIf(_disposed, this); - IntPtr address = AddressOf(name); - return _memory.WriteBytes(address, bytes); + int offset = GetRegionOffset(name); + if (offset > _size - bytes.Length) + throw new ArgumentOutOfRangeException(nameof(bytes), $"Region '{name}' write of {bytes.Length} bytes exceeds allocation size {_size}."); + + return _memory.WriteBytes(_baseAddress + offset, bytes); } /// diff --git a/WhiteMagicTest/Hooking/HookingTests.cs b/WhiteMagicTest/Hooking/HookingTests.cs index 23c6f46..44df20c 100644 --- a/WhiteMagicTest/Hooking/HookingTests.cs +++ b/WhiteMagicTest/Hooking/HookingTests.cs @@ -91,6 +91,46 @@ public class HookingTests } } + [Fact] + public void Patch_apply_on_execute_only_memory_succeeds() + { + using var reader = CreateReader(); + + byte[] code = [0xB8, 0x2A, 0x00, 0x00, 0x00, 0xC3]; // mov eax, 42; ret + IntPtr alloc = NativeMethods.VirtualAllocEx( + reader.Handle, + IntPtr.Zero, + code.Length, + MemoryAllocationType.Commit | MemoryAllocationType.Reserve, + MemoryProtectionType.ExecuteReadWrite); + Assert.NotEqual(IntPtr.Zero, alloc); + + try + { + reader.WriteBytes(alloc, code); + + // Remove write access. Without VirtualProtectEx in Patch.Apply, + // applying a patch would fail because the page is read-only for writes. + Assert.True(NativeMethods.VirtualProtectEx( + reader.Handle, + alloc, + code.Length, + MemoryProtectionType.ExecuteRead, + out MemoryProtectionType _)); + + Patch patch = reader.PatchManager.Create("nop-ret", alloc, [0x90, 0x90]); + patch.Apply(); + Assert.True(patch.IsApplied); + + patch.Remove(); + Assert.False(patch.IsApplied); + } + finally + { + NativeMethods.VirtualFreeEx(reader.Handle, alloc, 0, MemoryFreeType.Release); + } + } + [Fact] public void Detour_apply_redirects_callOriginal_remove_restores() { @@ -335,4 +375,55 @@ public class HookingTests NativeMethods.VirtualFreeEx(reader.Handle, allocation, 0, MemoryFreeType.Release); } } + + [Fact] + public async Task MainThreadPump_concurrent_dispose_and_pump_does_not_throw() + { + using var reader = CreateReader(); + + IntPtr targetPtr = AllocateFrameStub(reader, out IntPtr allocation); + try + { + var pump = new WhiteMagic.Execution.MainThreadPump(reader.DetourManager, targetPtr); + pump.Install(); + + FrameFunc routed = Marshal.GetDelegateForFunctionPointer(targetPtr); + + var tasks = new List>(); + for (int i = 0; i < 50; i++) + { + int value = i; + tasks.Add(pump.ExecuteAsync(() => value)); + } + + // Drive the detoured frame while disposing from another thread. + // This stresses the race between PumpHook completing work and + // Dispose faulting still-queued work. + Task drive = Task.Run(() => + { + for (int i = 0; i < 10; i++) + { + try { routed(); } catch { } + } + }); + + await Task.Delay(10); + pump.Dispose(); + await drive; + + // Any faulted task must have been cancelled by Dispose; no + // unhandled exceptions should escape from the pump itself. + foreach (Task task in tasks) + { + if (task.IsFaulted) + { + Assert.IsType(task.Exception!.InnerException); + } + } + } + finally + { + NativeMethods.VirtualFreeEx(reader.Handle, allocation, 0, MemoryFreeType.Release); + } + } } diff --git a/WhiteMagicTest/Memory/AllocatedMemoryTests.cs b/WhiteMagicTest/Memory/AllocatedMemoryTests.cs index ac3c759..f07d7b0 100644 --- a/WhiteMagicTest/Memory/AllocatedMemoryTests.cs +++ b/WhiteMagicTest/Memory/AllocatedMemoryTests.cs @@ -251,4 +251,73 @@ public class AllocatedMemoryTests Assert.Contains("does not exist", ex.Message); } + + [Fact] + public void Read_T_throws_when_value_exceeds_allocation() + { + using var reader = CreateReader(); + + using var allocated = new AllocatedMemory(reader, 100); + allocated.AddRegion("boundary", 99); + + var ex = Assert.Throws(() => + allocated.Read("boundary")); + + Assert.Equal("name", ex.ParamName); + } + + [Fact] + public void Write_T_throws_when_value_exceeds_allocation() + { + using var reader = CreateReader(); + + using var allocated = new AllocatedMemory(reader, 100); + allocated.AddRegion("boundary", 99); + + var ex = Assert.Throws(() => + allocated.Write("boundary", 42)); + + Assert.Equal("name", ex.ParamName); + } + + [Fact] + public void ReadBytes_throws_when_count_exceeds_allocation() + { + using var reader = CreateReader(); + + using var allocated = new AllocatedMemory(reader, 100); + allocated.AddRegion("boundary", 99); + + var ex = Assert.Throws(() => + allocated.ReadBytes("boundary", 2)); + + Assert.Equal("count", ex.ParamName); + } + + [Fact] + public void WriteBytes_throws_when_span_exceeds_allocation() + { + using var reader = CreateReader(); + + using var allocated = new AllocatedMemory(reader, 100); + allocated.AddRegion("boundary", 99); + + var ex = Assert.Throws(() => + allocated.WriteBytes("boundary", new byte[2])); + + Assert.Equal("bytes", ex.ParamName); + } + + [Fact] + public void Read_and_Write_at_exact_allocation_boundary_succeed() + { + using var reader = CreateReader(); + + using var allocated = new AllocatedMemory(reader, 100); + allocated.AddRegion("boundary", 96); + + const int expected = unchecked((int)0xDEADBEEF); + Assert.True(allocated.Write("boundary", expected)); + Assert.Equal(expected, allocated.Read("boundary")); + } }