whitemagic/WhiteMagicTest/Injection/DllInjectorTests.cs

using System.Runtime.InteropServices;
using WhiteMagic;
using WhiteMagic.Injection;
using WhiteMagic.Native;

namespace WhiteMagicTest.Injection;

/// <summary>
/// Tests for <see cref="DllInjector"/>.
/// </summary>
/// <remarks>
/// Real injection tests exercise the current process, because it is always available
/// and the injected DLLs are ordinary system modules that are already loaded.
/// </remarks>
public class DllInjectorTests
{
    private static string GetExistingSystemDll()
    {
        // user32.dll exists on every Windows system and matches the host bitness.
        string path = Path.Combine(Environment.SystemDirectory, "user32.dll");
        Assert.True(File.Exists(path), $"{path} must exist for the test.");
        return path;
    }

    [Fact(Skip = "Integration injection test - run against a dedicated target process")]
    public void InjectWithRemoteThread_loads_system_dll_in_current_process()
    {
        using var reader = new InProcessReader();
        var injector = new DllInjector(reader);

        IntPtr moduleBase = injector.InjectWithRemoteThread(GetExistingSystemDll());

        Assert.NotEqual(IntPtr.Zero, moduleBase);
    }

    [Fact]
    public void InjectWithRemoteThread_throws_for_missing_dll()
    {
        using var reader = new InProcessReader();
        var injector = new DllInjector(reader);
        string missingPath = Path.Combine(Path.GetTempPath(), $"wm-missing-{Guid.NewGuid()}.dll");

        Assert.False(File.Exists(missingPath));
        Assert.Throws<FileNotFoundException>(() => injector.InjectWithRemoteThread(missingPath));
    }

    [Fact]
    public void InjectWithRemoteThread_rejects_bitness_mismatch()
    {
        // A fake reader that reports the opposite bitness from the current process.
        using var reader = new FakeBitnessMemoryBase(!Environment.Is64BitProcess);
        var injector = new DllInjector(reader);

        var ex = Assert.Throws<InvalidOperationException>(
            () => injector.InjectWithRemoteThread(GetExistingSystemDll()));

        Assert.Contains("bitness", ex.Message, StringComparison.OrdinalIgnoreCase);
    }

    [Fact(Skip = "Integration injection test - run against a dedicated target process")]
    public void InjectWithThreadHijack_loads_system_dll_and_restores_context()
    {
        using var reader = new InProcessReader();
        var injector = new DllInjector(reader);

        using var stopEvent = new ManualResetEventSlim(false);
        using var startedEvent = new ManualResetEventSlim(false);

        int osThreadId = 0;
        Exception? threadError = null;

        var helper = new Thread(() =>
        {
            try
            {
                osThreadId = (int)NativeMethods.GetCurrentThreadId();
                startedEvent.Set();

                // Loop with short sleeps so the thread can be hijacked safely and can
                // also be stopped once its original context is restored.
                while (!stopEvent.IsSet)
                {
                    Thread.Sleep(10);
                }
            }
            catch (Exception ex)
            {
                threadError = ex;
            }
        });
        helper.IsBackground = true;
        helper.Start();

        try
        {
            Assert.True(startedEvent.Wait(TimeSpan.FromSeconds(5)), "Helper thread did not start.");
            Assert.NotEqual(0, osThreadId);

            IntPtr moduleBase = injector.InjectWithThreadHijack(osThreadId, GetExistingSystemDll());
            Assert.NotEqual(IntPtr.Zero, moduleBase);

            // Tell the helper thread to exit. If the original context was restored correctly,
            // the thread will return to its loop and observe the stop event.
            stopEvent.Set();
            Assert.True(helper.Join(TimeSpan.FromSeconds(5)), "Helper thread did not exit after context restore.");
            Assert.Null(threadError);
        }
        finally
        {
            stopEvent.Set();
            helper.Join(TimeSpan.FromSeconds(5));
        }
    }

    /// <summary>
    /// A minimal <see cref="MemoryBase"/> whose only job is to report a chosen bitness.
    /// Reads and writes are not expected to be called by the rejection path.
    /// </summary>
    private sealed class FakeBitnessMemoryBase : MemoryBase
    {
        public FakeBitnessMemoryBase(bool is64Bit)
        {
            Is64Bit = is64Bit;
        }

        public override IntPtr ImageBase => IntPtr.Zero;

        public override SafeMemoryHandle Handle => new(IntPtr.Zero);

        public override bool Is64Bit { get; }

        public override int ProcessId => Environment.ProcessId;

        public override byte[] ReadBytes(IntPtr address, int count, bool isRelative = false)
            => throw new NotSupportedException();

        public override int WriteBytes(IntPtr address, ReadOnlySpan<byte> bytes, bool isRelative = false)
            => throw new NotSupportedException();
    }
}