whitemagic/WhiteMagic/Memory/ProtectionScope.cs

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

namespace WhiteMagic.Memory;

/// <summary>
/// A scope that temporarily changes page protection via <c>VirtualProtectEx</c> and
/// restores the original protection when disposed, including when the guarded body throws.
/// </summary>
public sealed class ProtectionScope : IDisposable
{
    private readonly MemoryBase _memory;
    private readonly IntPtr _address;
    private readonly nint _size;
    private readonly MemoryProtectionType _originalProtection;
    private bool _disposed;

    /// <summary>
    /// Creates a new protection scope, applying <paramref name="newProtection"/> to the
    /// specified range immediately.
    /// </summary>
    internal ProtectionScope(MemoryBase memory, IntPtr address, nint size, MemoryProtectionType newProtection)
    {
        _memory = memory ?? throw new ArgumentNullException(nameof(memory));

        if (address == IntPtr.Zero)
            throw new ArgumentException("Address cannot be zero.", nameof(address));

        if (size <= 0)
            throw new ArgumentOutOfRangeException(nameof(size), "Size must be positive.");

        _address = address;
        _size = size;

        if (!NativeMethods.VirtualProtectEx(
                memory.Handle,
                address,
                size,
                newProtection,
                out _originalProtection))
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException(
                $"VirtualProtectEx failed to change protection: error {error}.");
        }
    }

    /// <summary>Restores the original page protection if it has not already been restored.</summary>
    public void Dispose()
    {
        if (!_disposed)
        {
            _disposed = true;
            NativeMethods.VirtualProtectEx(_memory.Handle, _address, _size, _originalProtection, out _);
        }
    }
}