Add thread control surfaces

Adds RemoteThread, ThreadFactory (enumeration, main-thread selection, get-by-id), and FrozenThread scoped freeze. Supports suspend/resume, 32/64-bit context round-trip, TEB query, and reverse-order resume on dispose. Closes section 2 of add-thread-region-finder.
2026-07-22 16:04:30 +02:00
parent f0faca3112
commit 9aef9c21e3
7 changed files with 977 additions and 0 deletions
@@ -0,0 +1,192 @@
+using System.Linq;
+using System.Threading;
+using SysThread = System.Threading.Thread;
+using WhiteMagic;
+using WhiteMagic.Native;
+using WhiteMagic.Thread;
+using Xunit;
+
+namespace WhiteMagicTest.Thread;
+
+/// <summary>
+/// Tests for scoped thread freeze via <see cref="FrozenThread"/> and <see cref="ThreadFactory.Freeze"/>.
+/// </summary>
+public sealed class FrozenThreadTests
+{
+    [Fact]
+    public void Freeze_suspends_selected_workers_until_disposed()
+    {
+        using var magic = Magic.OpenInProcess();
+        var factory = new ThreadFactory(magic.Memory);
+
+        using var cts1 = new CancellationTokenSource();
+        using var cts2 = new CancellationTokenSource();
+        var started1 = new ManualResetEventSlim(false);
+        var started2 = new ManualResetEventSlim(false);
+        int osThreadId1 = 0;
+        int osThreadId2 = 0;
+
+        var worker1 = new SysThread(() =>
+        {
+            osThreadId1 = (int)NativeMethods.GetCurrentThreadId();
+            started1.Set();
+            while (!cts1.IsCancellationRequested)
+                SysThread.Sleep(10);
+        });
+
+        var worker2 = new SysThread(() =>
+        {
+            osThreadId2 = (int)NativeMethods.GetCurrentThreadId();
+            started2.Set();
+            while (!cts2.IsCancellationRequested)
+                SysThread.Sleep(10);
+        });
+
+        worker1.Start();
+        worker2.Start();
+        started1.Wait();
+        started2.Wait();
+
+        int[] targetIds = [osThreadId1, osThreadId2];
+
+        try
+        {
+            var selected = factory.Enumerate().Where(t => targetIds.Contains(t.Id)).ToList();
+            Assert.Equal(2, selected.Count);
+
+            using (factory.Freeze(selected))
+            {
+                cts1.Cancel();
+                cts2.Cancel();
+
+                Assert.False(worker1.Join(100));
+                Assert.False(worker2.Join(100));
+            }
+
+            Assert.True(worker1.Join(1000));
+            Assert.True(worker2.Join(1000));
+        }
+        finally
+        {
+            if (worker1.IsAlive)
+            {
+                cts1.Cancel();
+                using var t = new RemoteThread(magic.Memory, osThreadId1);
+                t.Resume();
+                worker1.Join(1000);
+            }
+
+            if (worker2.IsAlive)
+            {
+                cts2.Cancel();
+                using var t = new RemoteThread(magic.Memory, osThreadId2);
+                t.Resume();
+                worker2.Join(1000);
+            }
+        }
+    }
+
+    [Fact]
+    public void Dispose_resumes_only_frozen_threads_leaving_external_suspends_intact()
+    {
+        using var magic = Magic.OpenInProcess();
+        var factory = new ThreadFactory(magic.Memory);
+
+        using var cts = new CancellationTokenSource();
+        var started = new ManualResetEventSlim(false);
+        int osThreadId = 0;
+
+        var worker = new SysThread(() =>
+        {
+            osThreadId = (int)NativeMethods.GetCurrentThreadId();
+            started.Set();
+            while (!cts.IsCancellationRequested)
+                SysThread.Sleep(10);
+        });
+
+        worker.Start();
+        started.Wait();
+
+        try
+        {
+            // Suspend the worker externally first.
+            using (var external = new RemoteThread(magic.Memory, osThreadId))
+            {
+                external.Suspend();
+
+                var selected = factory.Enumerate().Where(t => t.Id == osThreadId).ToList();
+                using (factory.Freeze(selected))
+                {
+                    // Frozen scope adds one more suspend count.
+                }
+
+                // After the freeze scope disposes, the worker was resumed once.
+                // Because it was already externally suspended, it should still be suspended.
+                cts.Cancel();
+                Assert.False(worker.Join(100));
+
+                external.Resume();
+            }
+
+            Assert.True(worker.Join(1000));
+        }
+        finally
+        {
+            if (worker.IsAlive)
+            {
+                cts.Cancel();
+                using var t = new RemoteThread(magic.Memory, osThreadId);
+                t.Resume();
+                worker.Join(1000);
+            }
+        }
+    }
+
+    [Fact]
+    public void Exception_in_body_still_resumes_frozen_threads()
+    {
+        using var magic = Magic.OpenInProcess();
+        var factory = new ThreadFactory(magic.Memory);
+
+        using var cts = new CancellationTokenSource();
+        var started = new ManualResetEventSlim(false);
+        int osThreadId = 0;
+
+        var worker = new SysThread(() =>
+        {
+            osThreadId = (int)NativeMethods.GetCurrentThreadId();
+            started.Set();
+            while (!cts.IsCancellationRequested)
+                SysThread.Sleep(10);
+        });
+
+        worker.Start();
+        started.Wait();
+
+        try
+        {
+            var selected = factory.Enumerate().Where(t => t.Id == osThreadId).ToList();
+
+            Assert.Throws<InvalidOperationException>(new Action(() =>
+            {
+                using (factory.Freeze(selected))
+                {
+                    throw new InvalidOperationException("Intentional failure inside freeze scope.");
+                }
+            }));
+
+            cts.Cancel();
+            Assert.True(worker.Join(1000));
+        }
+        finally
+        {
+            if (worker.IsAlive)
+            {
+                cts.Cancel();
+                using var t = new RemoteThread(magic.Memory, osThreadId);
+                t.Resume();
+                worker.Join(1000);
+            }
+        }
+    }
+}
@@ -0,0 +1,131 @@
+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);
+            }
+        }
+    }
+}
@@ -0,0 +1,130 @@
+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"/> open/suspend/resume and context round-trip.
+/// </summary>
+public sealed class RemoteThreadTests
+{
+    [Fact]
+    public void Open_by_id_succeeds_for_current_thread()
+    {
+        using var magic = Magic.OpenInProcess();
+        int currentId = (int)NativeMethods.GetCurrentThreadId();
+
+        using var thread = new RemoteThread(magic.Memory, currentId);
+        Assert.Equal(currentId, thread.Id);
+    }
+
+    [Fact]
+    public void Suspend_returns_prior_count_and_stops_worker()
+    {
+        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);
+
+            uint prior = thread.Suspend();
+            Assert.True(prior < 0xFFFFFFFF);
+
+            cts.Cancel();
+            // Worker cannot observe cancellation while suspended.
+            Assert.False(worker.Join(100));
+
+            thread.Resume();
+            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 Resume_restarts_a_suspended_worker()
+    {
+        using var magic = Magic.OpenInProcess();
+        using var cts = new CancellationTokenSource();
+        var started = new ManualResetEventSlim(false);
+        var resumed = new ManualResetEventSlim(false);
+        int osThreadId = 0;
+
+        var worker = new System.Threading.Thread(() =>
+        {
+            osThreadId = (int)NativeMethods.GetCurrentThreadId();
+            started.Set();
+            while (!cts.IsCancellationRequested)
+            {
+                resumed.Set();
+                System.Threading.Thread.Sleep(10);
+            }
+        });
+
+        worker.Start();
+        started.Wait();
+
+        try
+        {
+            using var thread = new RemoteThread(magic.Memory, osThreadId);
+            thread.Suspend();
+            resumed.Reset();
+
+            uint prior = thread.Resume();
+            Assert.True(prior < 0xFFFFFFFF);
+
+            // Worker must reach the resumed flag again.
+            Assert.True(resumed.Wait(1000));
+            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 GetTeb_returns_managed_teb_for_thread()
+    {
+        using var magic = Magic.OpenInProcess();
+        int currentId = (int)NativeMethods.GetCurrentThreadId();
+
+        using var thread = new RemoteThread(magic.Memory, currentId);
+        using var teb = thread.GetTeb();
+
+        Assert.NotEqual(IntPtr.Zero, teb.ReadTebAddress());
+    }
+}
@@ -0,0 +1,61 @@
+using System.Linq;
+using System.Threading;
+using SysThread = System.Threading.Thread;
+using WhiteMagic;
+using WhiteMagic.Native;
+using WhiteMagic.Thread;
+using Xunit;
+
+namespace WhiteMagicTest.Thread;
+
+/// <summary>
+/// Tests for <see cref="ThreadFactory"/> enumeration and main-thread selection.
+/// </summary>
+public sealed class ThreadFactoryTests
+{
+    [Fact]
+    public void Enumerate_returns_only_target_threads()
+    {
+        using var magic = Magic.OpenInProcess();
+        var factory = new ThreadFactory(magic.Memory);
+
+        int currentOsId = (int)NativeMethods.GetCurrentThreadId();
+        var ids = factory.Enumerate().Select(t => t.Id).ToList();
+
+        Assert.True(ids.Count > 0);
+        Assert.Contains(currentOsId, ids);
+    }
+
+    [Fact]
+    public void GetThreadById_returns_matching_thread()
+    {
+        using var magic = Magic.OpenInProcess();
+        var factory = new ThreadFactory(magic.Memory);
+
+        int currentOsId = (int)NativeMethods.GetCurrentThreadId();
+        using RemoteThread thread = factory.GetThreadById(currentOsId);
+        Assert.Equal(currentOsId, thread.Id);
+    }
+
+    [Fact]
+    public void GetThreadById_throws_for_nonexistent_thread()
+    {
+        using var magic = Magic.OpenInProcess();
+        var factory = new ThreadFactory(magic.Memory);
+
+        Assert.Throws<InvalidOperationException>(() => factory.GetThreadById(0x7FFFFFFF));
+    }
+
+    [Fact]
+    public void MainThread_returns_a_thread_belonging_to_the_target()
+    {
+        using var magic = Magic.OpenInProcess();
+        var factory = new ThreadFactory(magic.Memory);
+
+        using RemoteThread main = factory.MainThread;
+        Assert.NotNull(main);
+
+        var ids = factory.Enumerate().Select(t => t.Id).ToList();
+        Assert.Contains(main.Id, ids);
+    }
+}