whitemagic/WhiteMagicTest/Injection/CodeInjectorTests.cs

using System.ComponentModel;
using WhiteMagic;
using WhiteMagic.Injection;
using WhiteMagic.Memory;
using WhiteMagic.Native;

namespace WhiteMagicTest.Injection;

/// <summary>
/// Tests for <see cref="CodeInjector"/>.
/// </summary>
public class CodeInjectorTests
{
    private static InProcessReader CreateReader()
    {
        return new InProcessReader();
    }

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

        // Allocate a buffer to write to
        byte[] buffer = new byte[32];
        var handle = System.Runtime.InteropServices.GCHandle.Alloc(
            buffer,
            System.Runtime.InteropServices.GCHandleType.Pinned);
        try
        {
            IntPtr addr = handle.AddrOfPinnedObject();

            // Simple x64 payload: mov eax, 42; ret
            // B8 2A 00 00 00 C3
            byte[] code = { 0xB8, 0x2A, 0x00, 0x00, 0x00, 0xC3 };

            if (Environment.Is64BitProcess)
            {
                // 64-bit: mov eax, 42 (B8 2A 00 00 00) + ret (C3)
                code = new byte[] { 0xB8, 0x2A, 0x00, 0x00, 0x00, 0xC3 };
            }
            else
            {
                // 32-bit: mov eax, 42 (B8 2A 00 00 00) + ret (C3) - same encoding
                code = new byte[] { 0xB8, 0x2A, 0x00, 0x00, 0x00, 0xC3 };
            }

            IntPtr result = CodeInjector.InjectAtAddress(reader, addr, code);

            Assert.Equal(addr, result);
            Assert.Equal(code, buffer.Take(code.Length).ToArray());
        }
        finally
        {
            handle.Free();
        }
    }

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

        byte[] code = Array.Empty<byte>();

        var ex = Assert.Throws<ArgumentException>(() =>
            CodeInjector.InjectAtAddress(reader, IntPtr.Zero, code));

        Assert.Contains("Code cannot be empty", ex.Message);
    }

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

        byte[] code = { 0x90, 0x90, 0xC3 }; // nop; nop; ret

        var ex = Assert.Throws<ArgumentException>(() =>
            CodeInjector.InjectAtAddress(reader, IntPtr.Zero, code));

        Assert.Contains("Address cannot be zero", ex.Message);
    }

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

        // Simple x64 payload: ret (C3)
        byte[] code = { 0xC3 };

        using var allocated = CodeInjector.Inject(reader, code);

        Assert.NotEqual(IntPtr.Zero, allocated.BaseAddress);
        Assert.Equal(code.Length, allocated.Size);

        // Verify the code was written
        byte[] readBack = reader.ReadBytes(allocated.BaseAddress, code.Length);
        Assert.Equal(code, readBack);
    }

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

        byte[] code = { 0xC3 }; // ret

        using var allocated = CodeInjector.Inject(
            reader,
            code,
            MemoryProtectionType.ExecuteReadWrite);

        Assert.NotEqual(IntPtr.Zero, allocated.BaseAddress);
    }

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

        byte[] code = { 0xC3 }; // ret

        using var allocated = CodeInjector.Inject(
            reader,
            code,
            MemoryProtectionType.ExecuteRead);

        Assert.NotEqual(IntPtr.Zero, allocated.BaseAddress);
    }

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

        byte[] code = Array.Empty<byte>();

        var ex = Assert.Throws<ArgumentException>(() =>
            CodeInjector.Inject(reader, code));

        Assert.Contains("Code cannot be empty", ex.Message);
    }

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

        byte[] code = { 0xC3 }; // ret

        var allocated = CodeInjector.Inject(reader, code);

        Assert.NotNull(allocated);

        // Dispose should free the memory
        allocated.Dispose();

        // No exception should be thrown during disposal
    }

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

        // Allocate a buffer
        byte[] buffer = new byte[128];
        var handle = System.Runtime.InteropServices.GCHandle.Alloc(
            buffer,
            System.Runtime.InteropServices.GCHandleType.Pinned);
        try
        {
            IntPtr addr = handle.AddrOfPinnedObject();

            // Create a larger payload
            byte[] code = new byte[64];
            for (int i = 0; i < code.Length; i++)
                code[i] = (byte)(i & 0xFF);

            IntPtr result = CodeInjector.InjectAtAddress(reader, addr, code);

            Assert.Equal(addr, result);

            // Verify all bytes were written
            byte[] readBack = reader.ReadBytes(addr, code.Length);
            Assert.Equal(code, readBack);
        }
        finally
        {
            handle.Free();
        }
    }

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

        // mov eax, 12345678h; ret
        // x64: B8 78 56 34 12 C3
        // x86: B8 78 56 34 12 C3 (same)
        byte[] code = { 0xB8, 0x78, 0x56, 0x34, 0x12, 0xC3 };

        using var allocated = CodeInjector.Inject(reader, code);

        Assert.NotEqual(IntPtr.Zero, allocated.BaseAddress);

        // Verify the exact payload was written
        byte[] readBack = reader.ReadBytes(allocated.BaseAddress, code.Length);
        Assert.Equal(code, readBack);
    }
}