whitemagic/WhiteMagic/Memory/AllocatedMemory.cs

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

namespace WhiteMagic.Memory;

/// <summary>
/// Represents a chunk of remote memory subdivided into named regions.
/// </summary>
public sealed class AllocatedMemory : IDisposable
{
    private readonly MemoryBase _memory;
    private readonly IntPtr _baseAddress;
    private readonly int _size;
    private readonly Dictionary<string, int> _regions;
    private bool _disposed;

    /// <summary>
    /// Creates a new allocated memory chunk.
    /// </summary>
    /// <param name="memory">The memory accessor.</param>
    /// <param name="size">The size of the allocation in bytes.</param>
    /// <param name="protection">The initial memory protection.</param>
    /// <exception cref="Win32Exception">Allocation fails.</exception>
    public AllocatedMemory(MemoryBase memory, int size, MemoryProtectionType protection = MemoryProtectionType.ExecuteReadWrite)
    {
        ArgumentNullException.ThrowIfNull(memory);
        ArgumentOutOfRangeException.ThrowIfNegativeOrZero(size);

        _memory = memory;
        _size = size;
        _regions = new Dictionary<string, int>();

        // Allocate using VirtualAllocEx
        _baseAddress = NativeMethods.VirtualAllocEx(
            memory.Handle,
            IntPtr.Zero,
            size,
            MemoryAllocationType.Commit | MemoryAllocationType.Reserve,
            protection);

        if (_baseAddress == IntPtr.Zero)
        {
            int error = Marshal.GetLastPInvokeError();
            throw new Win32Exception(error, $"VirtualAllocEx failed (size={size}).");
        }
    }

    /// <summary>
    /// Gets the base address of the allocated memory.
    /// </summary>
    public IntPtr BaseAddress => _baseAddress;

    /// <summary>
    /// Gets the size of the allocation in bytes.
    /// </summary>
    public int Size => _size;

    /// <summary>
    /// Adds a named region at a specific offset within the allocation.
    /// </summary>
    /// <param name="name">The unique name for the region.</param>
    /// <param name="offset">The offset from the base address.</param>
    /// <exception cref="ArgumentException">A region with this name already exists.</exception>
    /// <exception cref="ArgumentOutOfRangeException">Offset is outside the allocation bounds.</exception>
    public void AddRegion(string name, int offset)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);
        ArgumentNullException.ThrowIfNull(name);

        ArgumentOutOfRangeException.ThrowIfNegative(offset);
        ArgumentOutOfRangeException.ThrowIfGreaterThanOrEqual(offset, _size);

        if (_regions.ContainsKey(name))
            throw new ArgumentException($"Region '{name}' already exists.", nameof(name));

        _regions[name] = offset;
    }

    /// <summary>
    /// Gets the absolute address of a named region.
    /// </summary>
    /// <param name="name">The region name.</param>
    /// <returns>The absolute address of the region.</returns>
    /// <exception cref="ArgumentException">No region with this name exists.</exception>
    public IntPtr AddressOf(string name)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);
        int offset = GetRegionOffset(name);
        return _baseAddress + offset;
    }

    private int GetRegionOffset(string name)
    {
        ArgumentNullException.ThrowIfNull(name);

        if (!_regions.TryGetValue(name, out int offset))
            throw new ArgumentException($"Region '{name}' does not exist.", nameof(name));

        return offset;
    }

    /// <summary>
    /// Reads a value of type <typeparamref name="T"/> from a named region.
    /// </summary>
    /// <typeparam name="T">The value type.</typeparam>
    /// <param name="name">The region name.</param>
    /// <returns>The value read from memory.</returns>
    /// <exception cref="ArgumentException">No region with this name exists.</exception>
    public T Read<T>(string name) where T : struct
    {
        ObjectDisposedException.ThrowIf(_disposed, this);

        int offset = GetRegionOffset(name);
        int size = Marshal.SizeOf<T>();
        if (offset > _size - size)
            throw new ArgumentOutOfRangeException(nameof(name), $"Region '{name}' read of {size} bytes exceeds allocation size {_size}.");

        return _memory.Read<T>(_baseAddress + offset);
    }

    /// <summary>
    /// Writes a value of type <typeparamref name="T"/> to a named region.
    /// </summary>
    /// <typeparam name="T">The value type.</typeparam>
    /// <param name="name">The region name.</param>
    /// <param name="value">The value to write.</param>
    /// <returns><see langword="true"/> if all bytes were written.</returns>
    /// <exception cref="ArgumentException">No region with this name exists.</exception>
    public bool Write<T>(string name, T value) where T : struct
    {
        ObjectDisposedException.ThrowIf(_disposed, this);

        int offset = GetRegionOffset(name);
        int size = Marshal.SizeOf<T>();
        if (offset > _size - size)
            throw new ArgumentOutOfRangeException(nameof(name), $"Region '{name}' write of {size} bytes exceeds allocation size {_size}.");

        return _memory.Write(_baseAddress + offset, value);
    }

    /// <summary>
    /// Reads bytes from a named region.
    /// </summary>
    /// <param name="name">The region name.</param>
    /// <param name="count">The number of bytes to read.</param>
    /// <returns>The bytes read from memory.</returns>
    /// <exception cref="ArgumentException">No region with this name exists.</exception>
    public byte[] ReadBytes(string name, int count)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);

        int offset = GetRegionOffset(name);
        ArgumentOutOfRangeException.ThrowIfNegative(count);
        if (offset > _size - count)
            throw new ArgumentOutOfRangeException(nameof(count), $"Region '{name}' read of {count} bytes exceeds allocation size {_size}.");

        return _memory.ReadBytes(_baseAddress + offset, count);
    }

    /// <summary>
    /// Writes bytes to a named region.
    /// </summary>
    /// <param name="name">The region name.</param>
    /// <param name="bytes">The bytes to write.</param>
    /// <returns>The number of bytes written.</returns>
    /// <exception cref="ArgumentException">No region with this name exists.</exception>
    public int WriteBytes(string name, ReadOnlySpan<byte> bytes)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);

        int offset = GetRegionOffset(name);
        if (offset > _size - bytes.Length)
            throw new ArgumentOutOfRangeException(nameof(bytes), $"Region '{name}' write of {bytes.Length} bytes exceeds allocation size {_size}.");

        return _memory.WriteBytes(_baseAddress + offset, bytes);
    }

    /// <summary>
    /// Frees the allocated memory.
    /// </summary>
    public void Dispose()
    {
        if (!_disposed)
        {
            _disposed = true;

            // Free using VirtualFreeEx
            if (_baseAddress != IntPtr.Zero)
            {
                NativeMethods.VirtualFreeEx(
                    _memory.Handle,
                    _baseAddress,
                    0,
                    MemoryFreeType.Release);
            }

            _regions.Clear();
        }
    }
}