Fix bounds, memory protection, and completion race in core helpers

- AllocatedMemory.Read<T>/Write<T>/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.
2026-07-22 02:16:23 +02:00
parent ffd72b37ed
commit 1911514120
5 changed files with 261 additions and 16 deletions
@@ -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<object?> work, TaskCompletionSource<object?> 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<object?> 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.
+            }
        }
    }
 }
@@ -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 _);
+        }
    }

    /// <summary>Restores the original bytes if they were captured.</summary>
@@ -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;
    }

@@ -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;
    }

    /// <summary>
@@ -105,8 +111,12 @@ public sealed class AllocatedMemory : IDisposable
    {
        ObjectDisposedException.ThrowIf(_disposed, this);

-        IntPtr address = AddressOf(name);
-        return _memory.Read<T>(address);
+        int offset = GetRegionOffset(name);
+        int size = Marshal.SizeOf<T>();
+        if (offset > _size - size)
+            throw new ArgumentOutOfRangeException(nameof(name), $"Region '{name}' read of {size} bytes exceeds allocation size {_size}.");
+
+        return _memory.Read<T>(_baseAddress + offset);
    }

    /// <summary>
@@ -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<T>();
+        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);
    }

    /// <summary>
@@ -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);
    }

    /// <summary>
@@ -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);
    }

    /// <summary>
@@ -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<FrameFunc>(targetPtr);
+
+            var tasks = new List<Task<int>>();
+            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<int> task in tasks)
+            {
+                if (task.IsFaulted)
+                {
+                    Assert.IsType<ObjectDisposedException>(task.Exception!.InnerException);
+                }
+            }
+        }
+        finally
+        {
+            NativeMethods.VirtualFreeEx(reader.Handle, allocation, 0, MemoryFreeType.Release);
+        }
+    }
 }
@@ -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<ArgumentOutOfRangeException>(() =>
+            allocated.Read<int>("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<ArgumentOutOfRangeException>(() =>
+            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<ArgumentOutOfRangeException>(() =>
+            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<ArgumentOutOfRangeException>(() =>
+            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<int>("boundary"));
+    }
 }