whitemagic/WhiteMagic/Thread/RemoteThread.cs

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

namespace WhiteMagic.Thread;

/// <summary>
/// A handle to an existing thread in the target process. Provides suspend/resume,
/// context read/write, and TEB query.
/// </summary>
public sealed class RemoteThread : IDisposable
{
    private readonly MemoryBase _memory;
    private readonly SafeMemoryHandle _handle;
    private readonly int _id;
    private bool _disposed;

    /// <summary>The operating-system identifier of this thread.</summary>
    public int Id => _id;

    /// <summary>The native thread handle.</summary>
    internal SafeMemoryHandle Handle => _handle;

    internal RemoteThread(MemoryBase memory, int threadId, SafeMemoryHandle handle)
    {
        _memory = memory ?? throw new ArgumentNullException(nameof(memory));
        _id = threadId;
        _handle = handle ?? throw new ArgumentNullException(nameof(handle));
    }

    /// <summary>
    /// Opens the thread specified by <paramref name="threadId"/> in the target process
    /// represented by <paramref name="memory"/>.
    /// </summary>
    public RemoteThread(MemoryBase memory, int threadId)
        : this(memory, threadId, OpenHandle(threadId))
    {
    }

    private static SafeMemoryHandle OpenHandle(int threadId)
    {
        if (threadId <= 0)
            throw new ArgumentException("Thread ID must be positive.", nameof(threadId));

        const ThreadAccess requiredAccess =
            ThreadAccess.SuspendResume |
            ThreadAccess.GetContext |
            ThreadAccess.SetContext |
            ThreadAccess.QueryInformation;

        SafeMemoryHandle handle = NativeMethods.OpenThread(requiredAccess, false, threadId);
        if (handle.IsInvalid)
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException($"OpenThread failed for thread {threadId}: error {error}.");
        }

        return handle;
    }

    /// <summary>
    /// Suspends the thread and returns its previous suspend count.
    /// </summary>
    public uint Suspend()
    {
        uint result = NativeMethods.SuspendThread(_handle);
        if (result == 0xFFFFFFFF)
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException($"SuspendThread failed for thread {_id}: error {error}.");
        }

        return result;
    }

    /// <summary>
    /// Resumes the thread and returns its previous suspend count.
    /// </summary>
    public uint Resume()
    {
        uint result = NativeMethods.ResumeThread(_handle);
        if (result == 0xFFFFFFFF)
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException($"ResumeThread failed for thread {_id}: error {error}.");
        }

        return result;
    }

    /// <summary>
    /// Reads the 64-bit native context of the thread. Valid only for 64-bit targets.
    /// </summary>
    public unsafe void GetContext64(out Context64 context)
    {
        nint size = Marshal.SizeOf<Context64>();
        void* ptr = NativeMemory.AlignedAlloc((nuint)size, 16);
        try
        {
            Unsafe.InitBlock(ptr, 0, (uint)size);
            ((Context64*)ptr)->ContextFlags = ContextFlags.Amd64Full;

            if (!NativeMethods.GetThreadContext(_handle, ref *(Context64*)ptr))
            {
                int error = Marshal.GetLastPInvokeError();
                throw new InvalidOperationException($"GetThreadContext failed for thread {_id}: error {error}.");
            }

            context = *(Context64*)ptr;
        }
        finally
        {
            NativeMemory.AlignedFree(ptr);
        }
    }

    /// <summary>
    /// Writes the 64-bit native context of the thread. Valid only for 64-bit targets.
    /// </summary>
    public unsafe void SetContext64(ref Context64 context)
    {
        nint size = Marshal.SizeOf<Context64>();
        void* ptr = NativeMemory.AlignedAlloc((nuint)size, 16);
        try
        {
            *(Context64*)ptr = context;
            if (!NativeMethods.SetThreadContext(_handle, ref *(Context64*)ptr))
            {
                int error = Marshal.GetLastPInvokeError();
                throw new InvalidOperationException($"SetThreadContext failed for thread {_id}: error {error}.");
            }
        }
        finally
        {
            NativeMemory.AlignedFree(ptr);
        }
    }

    /// <summary>
    /// Reads the 32-bit native context of the thread. Valid only for 32-bit targets.
    /// </summary>
    public void GetContext32(out Context32 context)
    {
        if (_memory.Is64Bit)
        {
            context = default;
            throw new InvalidOperationException(
                "Use GetContext64 for 64-bit targets; GetContext32 is valid for 32-bit targets only.");
        }

        context = new Context32 { ContextFlags = ContextFlags.X86Full };
        if (!NativeMethods.GetThreadContext(_handle, ref context))
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException($"GetThreadContext failed for thread {_id}: error {error}.");
        }
    }

    /// <summary>
    /// Writes the 32-bit native context of the thread. Valid only for 32-bit targets.
    /// </summary>
    public void SetContext32(ref Context32 context)
    {
        if (_memory.Is64Bit)
            throw new InvalidOperationException(
                "Use SetContext64 for 64-bit targets; SetContext32 is valid for 32-bit targets only.");

        if (!NativeMethods.SetThreadContext(_handle, ref context))
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException($"SetThreadContext failed for thread {_id}: error {error}.");
        }
    }

    /// <summary>
    /// Returns a managed reader for this thread's Thread Environment Block.
    /// </summary>
    public ManagedTeb GetTeb()
    {
        return new ManagedTeb(_memory, _id);
    }

    /// <inheritdoc />
    public void Dispose()
    {
        if (!_disposed)
        {
            _disposed = true;
            _handle.Dispose();
        }
    }
}