From fa0b5b5013f9057335522596c21dfba8c0d43ce7 Mon Sep 17 00:00:00 2001 From: Kevin Bataille Date: Wed, 22 Jul 2026 00:51:39 +0200 Subject: [PATCH] Add observable alloc/free seams to BackgroundTaskExecutor and fix StubAllocator ownership - Add RemoteAllocator / RemoteReleaser internal test seams for string/struct scratch memory. - Route all scratch allocation/freeing through the seams so tests can observe leaks. - Fix inverted StubAllocator ownership: a caller-provided stub is now never freed by the executor. - Rewrite Execute_releases_allocated_remote_memory_on_write_failure to fail pre-fix by tracking fake allocations through the seams. - Update WriteFailingMemoryBase to carry a valid self-handle so the executor reaches the marshal/write path. Tests: 198 passing, 4 integration/interactive skipped. --- WhiteMagic/Execution/RemoteThreadExecutor.cs | 53 +++++++++++++------ .../Execution/RemoteThreadExecutorTests.cs | 30 +++++++++-- 2 files changed, 63 insertions(+), 20 deletions(-) diff --git a/WhiteMagic/Execution/RemoteThreadExecutor.cs b/WhiteMagic/Execution/RemoteThreadExecutor.cs index e843439..d43c933 100644 --- a/WhiteMagic/Execution/RemoteThreadExecutor.cs +++ b/WhiteMagic/Execution/RemoteThreadExecutor.cs @@ -43,6 +43,38 @@ public sealed class RemoteThreadExecutor /// internal Func? StubAllocator { get; set; } + /// Test seam: overrides remote scratch allocation for string/struct args. + /// Defaults to . + internal Func? RemoteAllocator { get; set; } + + /// Test seam: overrides remote scratch release. Defaults to + /// . + internal Action? RemoteReleaser { get; set; } + + private IntPtr AllocateScratch(int size) + { + if (RemoteAllocator is not null) + return RemoteAllocator(size); + + return NativeMethods.VirtualAllocEx( + _reader.Handle, + IntPtr.Zero, + size, + MemoryAllocationType.Commit | MemoryAllocationType.Reserve, + MemoryProtectionType.ReadWrite); + } + + private void ReleaseScratch(IntPtr address) + { + if (RemoteReleaser is not null) + { + RemoteReleaser(address); + return; + } + + NativeMethods.VirtualFreeEx(_reader.Handle, address, 0, MemoryFreeType.Release); + } + /// /// Initializes a new for the process exposed by /// . @@ -108,7 +140,9 @@ public sealed class RemoteThreadExecutor if (StubAllocator != null) { stubAddress = StubAllocator(address, stubBytes.Length); - stubOwnedByExecutor = stubAddress != IntPtr.Zero; + // A caller-provided stub region is owned by the caller; never free it. + if (stubAddress != IntPtr.Zero) + stubOwnedByExecutor = false; } if (stubAddress == IntPtr.Zero) @@ -193,8 +227,7 @@ public sealed class RemoteThreadExecutor foreach (IntPtr alloc in allocations) { - NativeMethods.VirtualFreeEx( - _reader.Handle, alloc, 0, MemoryFreeType.Release); + ReleaseScratch(alloc); } } } @@ -284,12 +317,7 @@ public sealed class RemoteThreadExecutor bytes.CopyTo(buffer, 0); buffer[^1] = 0; - IntPtr remote = NativeMethods.VirtualAllocEx( - _reader.Handle, - IntPtr.Zero, - buffer.Length, - MemoryAllocationType.Commit | MemoryAllocationType.Reserve, - MemoryProtectionType.ReadWrite); + IntPtr remote = AllocateScratch(buffer.Length); if (remote == IntPtr.Zero) { @@ -338,12 +366,7 @@ public sealed class RemoteThreadExecutor pin.Free(); } - IntPtr remote = NativeMethods.VirtualAllocEx( - _reader.Handle, - IntPtr.Zero, - size, - MemoryAllocationType.Commit | MemoryAllocationType.Reserve, - MemoryProtectionType.ReadWrite); + IntPtr remote = AllocateScratch(size); if (remote == IntPtr.Zero) { diff --git a/WhiteMagicTest/Execution/RemoteThreadExecutorTests.cs b/WhiteMagicTest/Execution/RemoteThreadExecutorTests.cs index 43f34ec..0f456e7 100644 --- a/WhiteMagicTest/Execution/RemoteThreadExecutorTests.cs +++ b/WhiteMagicTest/Execution/RemoteThreadExecutorTests.cs @@ -1,4 +1,6 @@ +using System.Collections.Generic; using System.Diagnostics; +using System.Linq; using System.Runtime.InteropServices; using WhiteMagic; using WhiteMagic.Assembly; @@ -156,11 +158,21 @@ public sealed class RemoteThreadExecutorTests 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. + var allocated = new List(); + var freed = new List(); + nint next = 0x4000_0000; + + executor.RemoteAllocator = size => { var p = (IntPtr)(next += 0x1000); allocated.Add(p); return p; }; + executor.RemoteReleaser = p => freed.Add(p); + + // The string arg is marshalled to remote scratch FIRST, then its write fails. + // Pre-fix the scratch was tracked only AFTER the write, so it escaped the finally + // free and leaked. Post-fix every allocation is released on the failure path. Assert.Throws(() => - executor.Execute(new IntPtr(0x123456789ABCDEF0L), CallConvention.Stdcall)); + executor.Execute(new IntPtr(0x123456789ABCDEF0L), CallConvention.Stdcall, "leakme")); + + Assert.NotEmpty(allocated); // the arg scratch was allocated + Assert.Equal(allocated.OrderBy(x => x), freed.OrderBy(x => x)); // and every alloc freed } private static int RunPayload(byte[] payload, CallConvention convention, params object?[] args) @@ -223,14 +235,21 @@ 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. + /// Holds a valid handle to the current process so the executor passes its + /// handle-validity check without performing real memory operations. /// private sealed class WriteFailingMemoryBase : MemoryBase { public override IntPtr ImageBase => IntPtr.Zero; - public override SafeMemoryHandle Handle { get; } = new SafeMemoryHandle(new IntPtr(-1)); + public override SafeMemoryHandle Handle { get; } public override bool Is64Bit => Environment.Is64BitProcess; public override int ProcessId => Environment.ProcessId; + public WriteFailingMemoryBase() + { + Handle = NativeMethods.OpenProcess(ProcessAccess.AllAccess, false, Environment.ProcessId); + } + public override byte[] ReadBytes(IntPtr address, int count, bool isRelative = false) => throw new NotSupportedException(); @@ -239,6 +258,7 @@ public sealed class RemoteThreadExecutorTests public override void Dispose() { + Handle?.Dispose(); } } }