whitemagic/WhiteMagicTest/Memory/MemoryRegionTests.cs

using System;
using System.Linq;
using System.Runtime.InteropServices;
using WhiteMagic;
using WhiteMagic.Memory;
using WhiteMagic.Native;
using Xunit;

namespace WhiteMagicTest.Memory;

/// <summary>
/// Tests for memory-region query, enumeration and scoped protection (tasks 1.2, 1.4, 1.6, 1.8).
/// </summary>
public sealed class MemoryRegionTests
{
    [Fact]
    public void Contains_returns_true_for_addresses_inside_half_open_range()
    {
        var region = new MemoryRegion(
            new IntPtr(0x10000),
            0x1000,
            MemoryProtectionType.ReadWrite,
            MemoryState.Commit,
            MemoryType.Private,
            new IntPtr(0x10000),
            MemoryProtectionType.ReadWrite);

        Assert.True(region.Contains(new IntPtr(0x10000)));
        Assert.True(region.Contains(new IntPtr(0x10FFF)));
        Assert.False(region.Contains(new IntPtr(0x11000)));
        Assert.False(region.Contains(new IntPtr(0x0FFF)));
    }

    [Fact]
    public void QueryRegion_returns_region_containing_committed_address()
    {
        using var reader = new InProcessReader();
        nint pageSize = Environment.SystemPageSize;

        IntPtr block = NativeMethods.VirtualAllocEx(
            reader.Handle,
            IntPtr.Zero,
            pageSize,
            MemoryAllocationType.Commit | MemoryAllocationType.Reserve,
            MemoryProtectionType.ReadWrite);

        Assert.NotEqual(IntPtr.Zero, block);

        try
        {
            MemoryRegion region = reader.QueryRegion(block);

            Assert.Equal(block, region.BaseAddress);
            Assert.True(region.Contains(block));
            Assert.True(region.Contains(block + (int)pageSize - 1));
            Assert.Equal(MemoryState.Commit, region.State);
            Assert.Equal(MemoryType.Private, region.Type);
            Assert.Equal(MemoryProtectionType.ReadWrite, region.Protection);
            Assert.Equal(MemoryProtectionType.ReadWrite, region.AllocationProtect);
            Assert.Equal(block, region.AllocationBase);
        }
        finally
        {
            NativeMethods.VirtualFreeEx(reader.Handle, block, 0, MemoryFreeType.Release);
        }
    }

    [Fact]
    public void EnumerateRegions_yields_ascending_non_overlapping_regions()
    {
        using var reader = new InProcessReader();

        MemoryRegion[] regions = reader.EnumerateRegions().Take(5).ToArray();
        Assert.True(regions.Length > 0);

        for (int i = 1; i < regions.Length; i++)
        {
            Assert.True(
                (nuint)regions[i].BaseAddress >=
                (nuint)regions[i - 1].BaseAddress + regions[i - 1].Size);
        }
    }

    [Fact]
    public void EnumerateRegions_is_lazy_and_stops_early()
    {
        using var reader = new InProcessReader();

        // Taking a single item must not force a full address-space walk.
        MemoryRegion first = reader.EnumerateRegions().First();
        Assert.True(first.Size > 0);
    }

    [Fact]
    public void ChangeProtection_applies_new_protection_inside_scope_and_restores_on_dispose()
    {
        using var reader = new InProcessReader();
        nint pageSize = Environment.SystemPageSize;

        IntPtr block = NativeMethods.VirtualAllocEx(
            reader.Handle,
            IntPtr.Zero,
            pageSize,
            MemoryAllocationType.Commit | MemoryAllocationType.Reserve,
            MemoryProtectionType.ReadWrite);

        Assert.NotEqual(IntPtr.Zero, block);

        try
        {
            Assert.Equal(MemoryProtectionType.ReadWrite, reader.QueryRegion(block).Protection);

            using (reader.ChangeProtection(block, pageSize, MemoryProtectionType.ExecuteReadWrite))
            {
                Assert.Equal(MemoryProtectionType.ExecuteReadWrite, reader.QueryRegion(block).Protection);
            }

            Assert.Equal(MemoryProtectionType.ReadWrite, reader.QueryRegion(block).Protection);
        }
        finally
        {
            NativeMethods.VirtualFreeEx(reader.Handle, block, 0, MemoryFreeType.Release);
        }
    }

    [Fact]
    public void ChangeProtection_restores_original_protection_when_body_throws()
    {
        using var reader = new InProcessReader();
        nint pageSize = Environment.SystemPageSize;

        IntPtr block = NativeMethods.VirtualAllocEx(
            reader.Handle,
            IntPtr.Zero,
            pageSize,
            MemoryAllocationType.Commit | MemoryAllocationType.Reserve,
            MemoryProtectionType.ReadWrite);

        Assert.NotEqual(IntPtr.Zero, block);

        try
        {
            Assert.Throws<InvalidOperationException>(new Action(() =>
            {
                using (reader.ChangeProtection(block, pageSize, MemoryProtectionType.ExecuteReadWrite))
                {
                    Assert.Equal(MemoryProtectionType.ExecuteReadWrite, reader.QueryRegion(block).Protection);
                    throw new InvalidOperationException("Intentional failure inside scope.");
                }
            }));

            Assert.Equal(MemoryProtectionType.ReadWrite, reader.QueryRegion(block).Protection);
        }
        finally
        {
            NativeMethods.VirtualFreeEx(reader.Handle, block, 0, MemoryFreeType.Release);
        }
    }
}