whitemagic/WhiteMagicTest/Thread/RemoteThreadContextTests.cs

using System.Threading;
using Thread = System.Threading.Thread;
using WhiteMagic;
using WhiteMagic.Native;
using WhiteMagic.Thread;
using Xunit;

namespace WhiteMagicTest.Thread;

/// <summary>
/// Tests for <see cref="RemoteThread.GetContext64"/> / <see cref="RemoteThread.SetContext64"/>.
/// 32-bit/WOW64 context is tested on a 32-bit host run.
/// </summary>
public sealed class RemoteThreadContextTests
{
    [Fact]
    public void GetContext64_SetContext64_round_trip_on_suspended_self_thread()
    {
        if (!Environment.Is64BitProcess)
            return;

        using var magic = Magic.OpenInProcess();
        using var cts = new CancellationTokenSource();
        var started = new ManualResetEventSlim(false);
        int osThreadId = 0;

        var worker = new System.Threading.Thread(() =>
        {
            osThreadId = (int)NativeMethods.GetCurrentThreadId();
            started.Set();
            while (!cts.IsCancellationRequested)
                System.Threading.Thread.Sleep(10);
        });

        worker.Start();
        started.Wait();

        try
        {
            using var thread = new RemoteThread(magic.Memory, osThreadId);
            thread.Suspend();
            System.Threading.Thread.Sleep(100);

            thread.GetContext64(out Context64 context);
            Assert.NotEqual(0uL, context.Rip);

            const ulong sentinel = 0x123456789ABCDEF0uL;
            ulong originalRax = context.Rax;
            context.Rax = sentinel;
            thread.SetContext64(ref context);

            thread.GetContext64(out context);
            Assert.Equal(sentinel, context.Rax);

            // Restore the original register before resuming so the worker keeps running.
            context.Rax = originalRax;
            thread.SetContext64(ref context);

            thread.Resume();
            cts.Cancel();
            Assert.True(worker.Join(1000));
        }
        finally
        {
            if (worker.IsAlive)
            {
                cts.Cancel();
                using var thread = new RemoteThread(magic.Memory, osThreadId);
                thread.Resume();
                worker.Join(1000);
            }
        }
    }

    [Fact]
    public void GetContext32_SetContext32_round_trip_on_suspended_self_thread()
    {
        if (Environment.Is64BitProcess)
            return;

        using var magic = Magic.OpenInProcess();
        using var cts = new CancellationTokenSource();
        var started = new ManualResetEventSlim(false);
        int osThreadId = 0;

        var worker = new System.Threading.Thread(() =>
        {
            osThreadId = (int)NativeMethods.GetCurrentThreadId();
            started.Set();
            while (!cts.IsCancellationRequested)
                System.Threading.Thread.Sleep(10);
        });

        worker.Start();
        started.Wait();

        try
        {
            using var thread = new RemoteThread(magic.Memory, osThreadId);
            thread.Suspend();

            thread.GetContext32(out Context32 context);
            Assert.NotEqual(0u, context.Eip);

            const uint sentinel = 0x89ABCDEFu;
            uint originalEax = context.Eax;
            context.Eax = sentinel;
            thread.SetContext32(ref context);

            thread.GetContext32(out context);
            Assert.Equal(sentinel, context.Eax);

            context.Eax = originalEax;
            thread.SetContext32(ref context);

            thread.Resume();
            cts.Cancel();
            Assert.True(worker.Join(1000));
        }
        finally
        {
            if (worker.IsAlive)
            {
                cts.Cancel();
                using var thread = new RemoteThread(magic.Memory, osThreadId);
                thread.Resume();
                worker.Join(1000);
            }
        }
    }
}