whitemagic/WhiteMagicTest/Hooking/HookingTests.cs

using System;
using System.Runtime.InteropServices;
using System.Threading.Tasks;
using WhiteMagic;
using WhiteMagic.Hooking;
using WhiteMagic.Native;
using Xunit;

namespace WhiteMagicTest.Hooking;

/// <summary>
/// Tests for <see cref="PatchManager"/>, <see cref="DetourManager"/> and
/// <see cref="Execution.MainThreadPump"/> operating in-process.
/// </summary>
public class HookingTests
{
    [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
    private delegate int FrameFunc();

    private const int FrameResult = 42;

    private static InProcessReader CreateReader()
    {
        return new InProcessReader();
    }

    /// <summary>
    /// Allocates a tiny executable function whose prologue is made entirely of
    /// covered instruction shapes, so detours apply cleanly in tests.
    /// </summary>
    private static IntPtr AllocateFrameStub(MemoryBase reader, out IntPtr allocationBase)
    {
        // x64: push rbp; push rdi; push rsi; push rbx; sub rsp, 0x28; sub rsp, 0x12345678;
        // mov eax, 42; add rsp, 0x12345678; add rsp, 0x28; pop rbx; pop rsi; pop rdi; pop rbp; ret
        byte[] code =
        [
            0x55, // push rbp
            0x57, // push rdi
            0x56, // push rsi
            0x53, // push rbx
            0x48, 0x83, 0xEC, 0x28, // sub rsp, 0x28
            0x48, 0x81, 0xEC, 0x78, 0x56, 0x34, 0x12, // sub rsp, 0x12345678
            0xB8, 0x2A, 0x00, 0x00, 0x00, // mov eax, 42
            0x48, 0x81, 0xC4, 0x78, 0x56, 0x34, 0x12, // add rsp, 0x12345678
            0x48, 0x83, 0xC4, 0x28, // add rsp, 0x28
            0x5B, // pop rbx
            0x5E, // pop rsi
            0x5F, // pop rdi
            0x5D, // pop rbp
            0xC3  // ret
        ];

        allocationBase = NativeMethods.VirtualAllocEx(
            reader.Handle,
            IntPtr.Zero,
            code.Length,
            MemoryAllocationType.Commit | MemoryAllocationType.Reserve,
            MemoryProtectionType.ExecuteReadWrite);
        Assert.NotEqual(IntPtr.Zero, allocationBase);

        reader.WriteBytes(allocationBase, code);
        return allocationBase;
    }

    [Fact]
    public void Patch_apply_writes_bytes_and_remove_restores_original()
    {
        using var reader = CreateReader();
        byte[] slot = new byte[8];
        GCHandle pin = GCHandle.Alloc(slot, GCHandleType.Pinned);
        try
        {
            IntPtr addr = pin.AddrOfPinnedObject();
            byte[] original = reader.ReadBytes(addr, 4);
            byte[] patchBytes = [0x90, 0x90, 0x90, 0x90];

            Patch patch = reader.PatchManager.Create("nop", addr, patchBytes);
            Assert.False(patch.IsApplied);

            patch.Apply();
            Assert.True(patch.IsApplied);
            Assert.Equal(patchBytes, reader.ReadBytes(addr, 4));

            patch.Remove();
            Assert.False(patch.IsApplied);
            Assert.Equal(original, reader.ReadBytes(addr, 4));
        }
        finally
        {
            pin.Free();
        }
    }

    [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()
    {
        using var reader = CreateReader();
        IntPtr targetPtr = AllocateFrameStub(reader, out IntPtr allocation);

        try
        {
            int hookCalls = 0;
            Detour? detour = null;
            FrameFunc hook = () =>
            {
                hookCalls++;
                return (int?)detour?.CallOriginal() ?? 0;
            };

            detour = reader.DetourManager.Create("frame", targetPtr, hook);
            detour.Apply();

            FrameFunc routed = Marshal.GetDelegateForFunctionPointer<FrameFunc>(targetPtr);
            int result = routed();
            Assert.True(hookCalls > 0);
            Assert.Equal(FrameResult, result);

            detour.Remove();
            hookCalls = 0;
            result = routed();
            Assert.Equal(0, hookCalls);
            Assert.Equal(FrameResult, result);

            GC.KeepAlive(hook);
        }
        finally
        {
            NativeMethods.VirtualFreeEx(reader.Handle, allocation, 0, MemoryFreeType.Release);
        }
    }

    [Fact]
    public void Detour_named_lookup_returns_existing_detour()
    {
        using var reader = CreateReader();
        IntPtr targetPtr = AllocateFrameStub(reader, out IntPtr allocation);

        try
        {
            Detour detour = reader.DetourManager.Create("lookup", targetPtr, (FrameFunc)(() => FrameResult));
            Assert.Same(detour, reader.DetourManager["lookup"]);
        }
        finally
        {
            NativeMethods.VirtualFreeEx(reader.Handle, allocation, 0, MemoryFreeType.Release);
        }
    }

    [Fact]
    public void Detour_aligned_prologue_applies_and_unknown_prologue_rejects()
    {
        using var reader = CreateReader();

        // A normal JIT-compiled function has a covered prologue shape.
        IntPtr targetPtr = AllocateFrameStub(reader, out IntPtr goodAllocation);
        try
        {
            Detour good = reader.DetourManager.Create("good", targetPtr, (FrameFunc)(() => FrameResult));
            good.Apply();
            good.Remove();
        }
        finally
        {
            NativeMethods.VirtualFreeEx(reader.Handle, goodAllocation, 0, MemoryFreeType.Release);
        }

        // Allocate a small executable region whose first instruction is outside the
        // covered set. The decoder must refuse to splice it.
        IntPtr code = NativeMethods.VirtualAllocEx(
            reader.Handle,
            IntPtr.Zero,
            32,
            MemoryAllocationType.Commit | MemoryAllocationType.Reserve,
            MemoryProtectionType.ExecuteReadWrite);
        Assert.NotEqual(IntPtr.Zero, code);

        try
        {
            // 0x0F 0x05 = syscall (not covered), followed by padding and a ret.
            byte[] unknown = [0x0F, 0x05, 0xC3, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC];
            reader.WriteBytes(code, unknown);

            Detour bad = reader.DetourManager.Create("bad", code, (FrameFunc)(() => FrameResult));
            Assert.Throws<InvalidOperationException>(() => bad.Apply());
        }
        finally
        {
            NativeMethods.VirtualFreeEx(reader.Handle, code, 0, MemoryFreeType.Release);
        }
    }

    [Fact]
    public void Dispose_restores_active_patches_and_detours()
    {
        InProcessReader reader = CreateReader();
        byte[] slot = new byte[8];
        GCHandle pin = GCHandle.Alloc(slot, GCHandleType.Pinned);
        try
        {
            IntPtr addr = pin.AddrOfPinnedObject();
            byte[] original = reader.ReadBytes(addr, 2);

            Patch patch = reader.PatchManager.Create("dispose-patch", addr, [0x90, 0x90]);
            patch.Apply();

            reader.Dispose();

            // Verify with a fresh reader; the original handle was closed by Dispose.
            using var verify = CreateReader();
            Assert.Equal(original, verify.ReadBytes(addr, 2));
        }
        finally
        {
            pin.Free();
        }
    }

    [Fact]
    public async Task MainThreadPump_drains_work_on_frame_call_and_uninstalls_on_dispose()
    {
        using var reader = CreateReader();

        IntPtr targetPtr = AllocateFrameStub(reader, out IntPtr allocation);
        try
        {
            var pump = new WhiteMagic.Execution.MainThreadPump(reader.DetourManager, targetPtr);
            pump.Install();

            Task<int> work = pump.ExecuteAsync(() => 123);

            // Drive the frame function manually. The detoured frame runs the pump hook on
            // this thread, drains the work queue, then calls the original frame function.
            FrameFunc routed = Marshal.GetDelegateForFunctionPointer<FrameFunc>(targetPtr);
            int frameResult = routed();

            Assert.Equal(FrameResult, frameResult);
            Assert.Equal(123, await work);

            pump.Dispose();

            // After uninstall, calling the frame function should behave like the original.
            routed = Marshal.GetDelegateForFunctionPointer<FrameFunc>(targetPtr);
            Assert.Equal(FrameResult, routed());
        }
        finally
        {
            NativeMethods.VirtualFreeEx(reader.Handle, allocation, 0, MemoryFreeType.Release);
        }
    }

    [Fact]
    public async Task MainThreadPump_exception_survives_and_does_not_kill_pump()
    {
        using var reader = CreateReader();

        IntPtr targetPtr = AllocateFrameStub(reader, out IntPtr allocation);
        try
        {
            var pump = new WhiteMagic.Execution.MainThreadPump(reader.DetourManager, targetPtr);
            pump.Install();

            Task<int> bad = pump.ExecuteAsync<int>(() => throw new InvalidOperationException("boom"));
            Task<int> good = pump.ExecuteAsync(() => 7);

            FrameFunc routed = Marshal.GetDelegateForFunctionPointer<FrameFunc>(targetPtr);
            routed();

            await Assert.ThrowsAsync<InvalidOperationException>(() => bad);
            Assert.Equal(7, await good);

            pump.Dispose();
        }
        finally
        {
            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<int> pending = pump.ExecuteAsync(() => 42);

            pump.Dispose();

            var ex = await Assert.ThrowsAsync<ObjectDisposedException>(() => 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<int>(() => 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);
        }
    }

    [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);
        }
    }
}