whitemagic/WhiteMagic/Execution/RemoteThreadExecutor.cs

using System.Globalization;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading.Tasks;
using WhiteMagic.Assembly;
using WhiteMagic.Native;

namespace WhiteMagic.Execution;

/// <summary>
/// Executes a function in the target process by creating a remote thread at a
/// calling-convention-aware stub. Waits for the thread to finish and returns the
/// typed exit value read from the thread's exit code.
/// </summary>
/// <remarks>
/// <para>This executor is safe only for thread-agnostic payloads. Calls that touch
/// single-threaded process state should use <see cref="MainThreadPump"/> instead.</para>
/// <para>String arguments are encoded as null-terminated UTF-8 and allocated in the
/// remote process; struct arguments are serialized with the default interop marshaler
/// (<see cref="Marshal.StructureToPtr"/>) and allocated with <see cref="Marshal.SizeOf(Type)"/>
/// bytes. All temporary remote allocations are released after the call, including on failure.</para>
/// </remarks>
public sealed class RemoteThreadExecutor
{
    private const uint WaitObject0 = 0x00000000;
    private const uint WaitTimeout = 0x00000102;
    private const uint WaitFailed = 0xFFFFFFFF;

    private const nuint AllocationGranularity = 0x10000; // 64 KB
    private const int NearAllocationAttempts = 64;

    private readonly MemoryBase _reader;
    private readonly StubAssembler _assembler;

    /// <summary>
    /// Internal hook for tests that need to place the generated call stub inside an
    /// already-allocated executable region (for example, immediately after the target
    /// payload to keep the relative CALL within ±2 GiB).
    /// </summary>
    /// <remarks>
    /// When this delegate returns a non-zero pointer, the executor does not take
    /// ownership of that memory and will not free it.
    /// </remarks>
    internal Func<IntPtr, nint, IntPtr>? StubAllocator { get; set; }

    /// <summary>Test seam: overrides remote scratch allocation for string/struct args.
    /// Defaults to <see cref="NativeMethods.VirtualAllocEx"/>.</summary>
    internal Func<int, IntPtr>? RemoteAllocator { get; set; }

    /// <summary>Test seam: overrides remote scratch release. Defaults to
    /// <see cref="NativeMethods.VirtualFreeEx"/>.</summary>
    internal Action<IntPtr>? RemoteReleaser { get; set; }

    private IntPtr AllocateScratch(int size)
    {
        if (RemoteAllocator is not null)
            return RemoteAllocator(size);

        return NativeMethods.VirtualAllocEx(
            _reader.Handle,
            IntPtr.Zero,
            size,
            MemoryAllocationType.Commit | MemoryAllocationType.Reserve,
            MemoryProtectionType.ReadWrite);
    }

    private void ReleaseScratch(IntPtr address)
    {
        if (RemoteReleaser is not null)
        {
            RemoteReleaser(address);
            return;
        }

        NativeMethods.VirtualFreeEx(_reader.Handle, address, 0, MemoryFreeType.Release);
    }

    /// <summary>
    /// Initializes a new <see cref="RemoteThreadExecutor"/> for the process exposed by
    /// <paramref name="reader"/>.
    /// </summary>
    /// <param name="reader">The memory reader that owns the target process handle.</param>
    public RemoteThreadExecutor(MemoryBase reader)
    {
        _reader = reader ?? throw new ArgumentNullException(nameof(reader));
        _assembler = new StubAssembler();
    }

    /// <summary>
    /// Calls the function at <paramref name="address"/> in the target process using a
    /// remote thread and returns its exit value cast to <typeparamref name="T"/>.
    /// </summary>
    /// <typeparam name="T">The expected return type.</typeparam>
    /// <param name="address">The target function address.</param>
    /// <param name="convention">The calling convention (ignored on x64 targets).</param>
    /// <param name="args">Arguments to pass. Primitives, pointers and enums are packed
    /// into pointer-sized slots. Strings and structs are allocated remotely and passed
    /// by pointer.</param>
    /// <returns>The function's exit value converted to <typeparamref name="T"/>.</returns>
    /// <exception cref="InvalidOperationException">The process handle is not open or a
    /// required native operation failed.</exception>
    /// <exception cref="TimeoutException">The remote thread did not complete in time.</exception>
    public Task<T> ExecuteAsync<T>(IntPtr address, CallConvention convention, params object?[] args)
    {
        return Task.Run(() => Execute<T>(address, convention, args));
    }

    /// <summary>
    /// Synchronous variant of <see cref="ExecuteAsync{T}"/>.
    /// </summary>
    public T Execute<T>(IntPtr address, CallConvention convention, params object?[] args)
    {
        if (_reader.Handle.IsInvalid)
        {
            throw new InvalidOperationException(
                "Cannot execute a remote function: the target process handle is not open.");
        }

        if (address == IntPtr.Zero)
        {
            throw new ArgumentException(
                "Target function address cannot be zero.", nameof(address));
        }

        int pointerSize = _reader.Is64Bit ? 8 : 4;
        var allocations = new List<IntPtr>(args.Length + 1);
        IntPtr stubAddress = IntPtr.Zero;
        SafeMemoryHandle? thread = null;

        try
        {
            nuint[] nativeArgs = MarshalArguments(args, pointerSize, allocations);

            // Compute the exact stub size with a dummy address close to the target;
            // the emitted byte count does not depend on the stub's final address.
            byte[] stubBytes = _assembler.BuildCallStub(
                address, address, nativeArgs, pointerSize, convention);

            bool stubOwnedByExecutor = true;
            if (StubAllocator != null)
            {
                stubAddress = StubAllocator(address, stubBytes.Length);
                // A caller-provided stub region is owned by the caller; never free it.
                if (stubAddress != IntPtr.Zero)
                    stubOwnedByExecutor = false;
            }

            if (stubAddress == IntPtr.Zero)
            {
                stubAddress = AllocateExecutableMemory(_reader.Handle, address, stubBytes.Length);
                stubOwnedByExecutor = true;
            }

            if (stubAddress == IntPtr.Zero)
            {
                int error = Marshal.GetLastPInvokeError();
                throw new InvalidOperationException(
                    $"Failed to allocate remote stub memory: error {error}");
            }

            if (stubOwnedByExecutor)
            {
                allocations.Add(stubAddress);
            }

            // Re-emit with the real stub address so the relative call lands correctly.
            stubBytes = _assembler.BuildCallStub(
                stubAddress, address, nativeArgs, pointerSize, convention);

            int written = _reader.WriteBytes(stubAddress, stubBytes);
            if (written != stubBytes.Length)
            {
                throw new InvalidOperationException(
                    $"Failed to write the call stub to the remote process (wrote {written} of {stubBytes.Length} bytes).");
            }

            thread = NativeMethods.CreateRemoteThread(
                _reader.Handle,
                IntPtr.Zero,
                0,
                stubAddress,
                IntPtr.Zero,
                ThreadCreationFlags.RunImmediately,
                out _);

            if (thread.IsInvalid)
            {
                int error = Marshal.GetLastPInvokeError();
                throw new InvalidOperationException(
                    $"CreateRemoteThread failed: error {error}");
            }

            uint waitResult = NativeMethods.WaitForSingleObject(thread, uint.MaxValue);
            if (waitResult == WaitFailed)
            {
                int error = Marshal.GetLastPInvokeError();
                throw new InvalidOperationException(
                    $"WaitForSingleObject failed: error {error}");
            }

            if (waitResult == WaitTimeout)
            {
                throw new TimeoutException(
                    "The remote thread did not complete within the requested timeout.");
            }

            if (waitResult != WaitObject0)
            {
                throw new InvalidOperationException(
                    $"Unexpected wait status: 0x{waitResult:X}");
            }

            if (!NativeMethods.GetExitCodeThread(thread, out uint exitCode))
            {
                int error = Marshal.GetLastPInvokeError();
                throw new InvalidOperationException(
                    $"GetExitCodeThread failed: error {error}");
            }

            return ConvertExitCode<T>(exitCode);
        }
        finally
        {
            // Dispose the thread handle explicitly so the safe handle releases it
            // before any virtual memory is freed.
            thread?.Dispose();

            foreach (IntPtr alloc in allocations)
            {
                ReleaseScratch(alloc);
            }
        }
    }

    /// <summary>
    /// Converts the raw DWORD exit code into the requested return type.
    /// </summary>
    private static T ConvertExitCode<T>(uint exitCode)
    {
        Type target = typeof(T);

        if (target == typeof(IntPtr) || target == typeof(nint))
        {
            return (T)(object)(IntPtr)(nint)exitCode;
        }

        if (target == typeof(UIntPtr) || target == typeof(nuint))
        {
            return (T)(object)(UIntPtr)(nuint)exitCode;
        }

        if (Nullable.GetUnderlyingType(target) is Type underlying)
        {
            return (T)Convert.ChangeType(exitCode, underlying, CultureInfo.InvariantCulture);
        }

        return (T)Convert.ChangeType(exitCode, target, CultureInfo.InvariantCulture);
    }

    /// <summary>
    /// Marshals managed arguments into pointer-sized native argument slots. Allocates
    /// remote memory for strings and structs and records each allocation in
    /// <paramref name="allocations"/>.
    /// </summary>
    private nuint[] MarshalArguments(object?[] args, int pointerSize, List<IntPtr> allocations)
    {
        var nativeArgs = new nuint[args.Length];

        for (int i = 0; i < args.Length; i++)
        {
            object? arg = args[i];
            nativeArgs[i] = MarshalArgument(arg, pointerSize, allocations);
        }

        return nativeArgs;
    }

    /// <summary>
    /// Marshals a single argument. Strings and structs become remote pointers; primitives,
    /// enums and pointer values are packed directly.
    /// </summary>
    private nuint MarshalArgument(object? arg, int pointerSize, List<IntPtr> allocations)
    {
        if (arg is null)
        {
            return 0;
        }

        if (arg is string s)
        {
            return MarshalString(s, allocations);
        }

        Type type = arg.GetType();
        if (IsPrimitiveOrPointer(type))
        {
            return PackPrimitive(arg, pointerSize);
        }

        if (type.IsValueType)
        {
            return MarshalStruct(arg, type, allocations);
        }

        throw new ArgumentException(
            $"Unsupported argument type: {type.FullName}. Only primitives, pointers, enums, strings and structs are supported.");
    }

    /// <summary>
    /// Allocates the UTF-8 encoding of a string in the target process and returns its
    /// remote address.
    /// </summary>
    private nuint MarshalString(string value, List<IntPtr> allocations)
    {
        byte[] bytes = Encoding.UTF8.GetBytes(value);
        byte[] buffer = new byte[bytes.Length + 1];
        bytes.CopyTo(buffer, 0);
        buffer[^1] = 0;

        IntPtr remote = AllocateScratch(buffer.Length);

        if (remote == IntPtr.Zero)
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException(
                $"Failed to allocate remote string memory: error {error}");
        }

        allocations.Add(remote);

        int written = _reader.WriteBytes(remote, buffer);
        if (written != buffer.Length)
        {
            throw new InvalidOperationException(
                $"Failed to write string bytes to the remote process (wrote {written} of {buffer.Length} bytes).");
        }

        return (nuint)(nint)remote;
    }

    /// <summary>
    /// Allocates unmanaged space for a struct in the target process, writes its bytes with
    /// the default interop marshaler, and returns the remote address.
    /// </summary>
    private nuint MarshalStruct(object value, Type type, List<IntPtr> allocations)
    {
        int size;
        try
        {
            size = Marshal.SizeOf(type);
        }
        catch (ArgumentException ex)
        {
            throw new InvalidOperationException(
                $"Cannot marshal argument of type {type.FullName}: {ex.Message}", ex);
        }

        byte[] buffer = new byte[size];
        GCHandle pin = GCHandle.Alloc(buffer, GCHandleType.Pinned);
        try
        {
            Marshal.StructureToPtr(value, pin.AddrOfPinnedObject(), false);
        }
        finally
        {
            pin.Free();
        }

        IntPtr remote = AllocateScratch(size);

        if (remote == IntPtr.Zero)
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException(
                $"Failed to allocate remote struct memory: error {error}");
        }

        allocations.Add(remote);

        int written = _reader.WriteBytes(remote, buffer);
        if (written != size)
        {
            throw new InvalidOperationException(
                $"Failed to write struct bytes to the remote process (wrote {written} of {size} bytes).");
        }

        return (nuint)(nint)remote;
    }

    /// <summary>
    /// Determines whether a type can be passed directly as a pointer-sized value.
    /// </summary>
    private static bool IsPrimitiveOrPointer(Type type)
    {
        if (type == typeof(IntPtr) || type == typeof(UIntPtr) ||
            type == typeof(nint) || type == typeof(nuint))
        {
            return true;
        }

        TypeCode code = Type.GetTypeCode(type);

        switch (code)
        {
            case TypeCode.Boolean:
            case TypeCode.Char:
            case TypeCode.SByte:
            case TypeCode.Byte:
            case TypeCode.Int16:
            case TypeCode.UInt16:
            case TypeCode.Int32:
            case TypeCode.UInt32:
            case TypeCode.Int64:
            case TypeCode.UInt64:
                return true;

            case TypeCode.Object when type.IsEnum:
            case TypeCode.Object when type == typeof(IntPtr) || type == typeof(UIntPtr):
                return true;

            default:
                return false;
        }
    }

    /// <summary>
    /// Packs a primitive, enum or pointer value into a pointer-sized unsigned integer.
    /// Values are truncated to the target pointer width so x86 arguments receive their
    /// low 32 bits.
    /// </summary>
    private static nuint PackPrimitive(object value, int pointerSize)
    {
        Type type = value.GetType();
        ulong raw;

        if (type == typeof(IntPtr) || type == typeof(nint))
        {
            raw = unchecked((ulong)(nint)value);
        }
        else if (type == typeof(UIntPtr) || type == typeof(nuint))
        {
            raw = (ulong)(UIntPtr)value;
        }
        else if (type.IsEnum)
        {
            raw = Convert.ToUInt64(value);
        }
        else
        {
            raw = Convert.ToUInt64(value);
        }

        if (pointerSize == 4)
        {
            raw &= uint.MaxValue;
        }

        return unchecked((nuint)raw);
    }

    /// <summary>
    /// Attempts to allocate executable memory close to <paramref name="preferredAddress"/>
    /// so that the relative CALL instruction in the generated stub stays within its
    /// ±2 GiB range.
    /// </summary>
    private static IntPtr AllocateExecutableMemory(
        SafeMemoryHandle handle,
        IntPtr preferredAddress,
        nint size)
    {
        nuint preferred = (nuint)(nint)preferredAddress;
        nuint mask = AllocationGranularity - (nuint)1;
        nuint aligned = (preferred + AllocationGranularity - (nuint)1) & ~mask;

        for (int i = 0; i < NearAllocationAttempts; i++)
        {
            nuint candidate;
            if (i == 0)
            {
                candidate = aligned;
            }
            else if ((i & 1) == 1)
            {
                candidate = aligned + (nuint)i * AllocationGranularity;
            }
            else
            {
                nuint offset = (nuint)i * AllocationGranularity;
                if (offset > aligned)
                {
                    continue;
                }

                candidate = aligned - offset;
            }

            IntPtr result = NativeMethods.VirtualAllocEx(
                handle,
                (IntPtr)(nint)candidate,
                size,
                MemoryAllocationType.Commit | MemoryAllocationType.Reserve,
                MemoryProtectionType.ExecuteReadWrite);

            if (result != IntPtr.Zero)
            {
                return result;
            }
        }

        return NativeMethods.VirtualAllocEx(
            handle,
            IntPtr.Zero,
            size,
            MemoryAllocationType.Commit | MemoryAllocationType.Reserve,
            MemoryProtectionType.ExecuteReadWrite);
    }
}