whitemagic/WhiteMagic/Hooking/Patch.cs

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

namespace WhiteMagic.Hooking;

/// <summary>
/// A single reversible byte patch. Captures the original bytes when applied,
/// restores them when removed, and reports its state by comparing live memory.
/// </summary>
public sealed class Patch : IDisposable
{
    private readonly MemoryBase _memory;

    /// <summary>The unique name of this patch.</summary>
    public string Name { get; }

    /// <summary>The address the patch overwrites.</summary>
    public IntPtr Address { get; }

    /// <summary>The bytes written by the patch.</summary>
    public byte[] PatchBytes { get; }

    /// <summary>The bytes captured before the patch was applied.</summary>
    public byte[]? OriginalBytes { get; private set; }

    /// <summary>
    /// <see langword="true"/> when the live bytes at <see cref="Address"/> match
    /// <see cref="PatchBytes"/>.
    /// </summary>
    public bool IsApplied
    {
        get
        {
            byte[] current = _memory.ReadBytes(Address, PatchBytes.Length);
            return current.SequenceEqual(PatchBytes);
        }
    }

    internal Patch(MemoryBase memory, string name, IntPtr address, byte[] patchBytes)
    {
        ArgumentNullException.ThrowIfNull(patchBytes);

        _memory = memory;
        Name = name;
        Address = address;
        PatchBytes = patchBytes;
    }

    /// <summary>Captures the original bytes and writes the patch bytes.</summary>
    public void Apply()
    {
        if (IsApplied)
            return;

        OriginalBytes = _memory.ReadBytes(Address, PatchBytes.Length);

        if (!NativeMethods.VirtualProtectEx(
            _memory.Handle,
            Address,
            PatchBytes.Length,
            MemoryProtectionType.ExecuteReadWrite,
            out MemoryProtectionType oldProtect))
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException($"Failed to change target memory protection: error {error}");
        }

        try
        {
            _memory.WriteBytes(Address, PatchBytes);
        }
        finally
        {
            NativeMethods.VirtualProtectEx(
                _memory.Handle,
                Address,
                PatchBytes.Length,
                oldProtect,
                out _);
        }
    }

    /// <summary>Restores the original bytes if they were captured.</summary>
    public void Remove()
    {
        if (OriginalBytes is null)
            return;

        if (!NativeMethods.VirtualProtectEx(
            _memory.Handle,
            Address,
            OriginalBytes.Length,
            MemoryProtectionType.ExecuteReadWrite,
            out MemoryProtectionType oldProtect))
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException($"Failed to change target memory protection: error {error}");
        }

        try
        {
            _memory.WriteBytes(Address, OriginalBytes);
        }
        finally
        {
            NativeMethods.VirtualProtectEx(
                _memory.Handle,
                Address,
                OriginalBytes.Length,
                oldProtect,
                out _);
        }

        OriginalBytes = null;
    }

    /// <inheritdoc />
    public void Dispose()
    {
        Remove();
    }
}