whitemagic/WhiteMagic/Assembly/IcedAssembler.cs

using System.Collections.Generic;
using System.Globalization;
using System.Reflection;
using Iced.Intel;

namespace WhiteMagic.Assembly;

/// <summary>
/// Optional <see cref="IAssembler"/> backend that assembles arbitrary x86/x64 mnemonic
/// text to machine code using the Iced library, and provides full instruction-boundary
/// decoding for detour prologue validation.
/// </summary>
/// <remarks>
/// <para>Iced ships a fluent code assembler (typed method calls) and a decoder, but no
/// text parser. This class bridges Intel-syntax text onto Iced's fluent
/// <see cref="Assembler"/> by reflection: each line's mnemonic selects the matching
/// <see cref="Assembler"/> method and its operands are bound to registers, immediates, or
/// labels. Register and immediate operands and label-relative branches are supported;
/// memory operands (<c>[reg+disp]</c>) are not — a caller needing those should emit bytes
/// directly.</para>
/// <para>This backend is entirely optional. Constructing it is the only thing that pulls
/// Iced into a behavioral path; the default <see cref="StubAssembler"/> never references it.</para>
/// </remarks>
public sealed class IcedAssembler : IAssembler
{
    private const int DefaultBitness = 64;

    private readonly int _bitness;

    // Lowercased register name -> boxed AssemblerRegisterNN value, built once from
    // Iced's AssemblerRegisters. Enables binding a text operand like "esp" to a typed
    // fluent-API register argument.
    private static readonly Dictionary<string, object> Registers = BuildRegisterMap();

    /// <summary>Creates an assembler for the given bitness (32 or 64).</summary>
    /// <param name="bitness">32 for x86, 64 for x64. Defaults to 64.</param>
    public IcedAssembler(int bitness = DefaultBitness)
    {
        if (bitness != 32 && bitness != 64)
            throw new ArgumentOutOfRangeException(nameof(bitness), "Bitness must be 32 or 64.");

        _bitness = bitness;
    }

    /// <inheritdoc />
    public byte[] Assemble(string assemblyText, ulong origin = 0)
    {
        ArgumentNullException.ThrowIfNull(assemblyText);

        var assembler = new Assembler(_bitness);

        List<(string Mnemonic, string[] Operands)> lines = Tokenize(assemblyText, out var labelNames);

        // Pre-create every label so a forward branch can reference it before its definition.
        var labels = new Dictionary<string, Label>(StringComparer.OrdinalIgnoreCase);
        foreach (string name in labelNames)
            labels[name] = assembler.CreateLabel(name);

        foreach ((string mnemonic, string[] operands) in lines)
        {
            // A pure label definition (e.g. "loop:") marks the current position.
            if (mnemonic.EndsWith(':'))
            {
                Label label = labels[mnemonic[..^1]];
                assembler.Label(ref label);
                continue;
            }

            EmitInstruction(assembler, mnemonic, operands, labels);
        }

        var writer = new ByteListCodeWriter();
        assembler.Assemble(writer, origin);
        return writer.Bytes.ToArray();
    }

    /// <summary>
    /// Computes the number of whole prologue-instruction bytes that must be preserved for
    /// a splice of <paramref name="requiredBytes"/> bytes, decoding arbitrary instructions
    /// (not just the common prologue shapes the built-in decoder covers). Matches the
    /// <c>PrologueLengthResolver</c> delegate so it can be assigned to
    /// <see cref="WhiteMagic.Hooking.DetourManager.PrologueLengthResolver"/>.
    /// </summary>
    /// <exception cref="InvalidOperationException">A prologue byte sequence does not decode
    /// to a valid instruction.</exception>
    public int GetPrologueLength(byte[] prologue, int requiredBytes, bool is64Bit)
    {
        ArgumentNullException.ThrowIfNull(prologue);

        var reader = new ByteArrayCodeReader(prologue);
        var decoder = Decoder.Create(is64Bit ? 64 : 32, reader);

        int total = 0;
        while (total < requiredBytes)
        {
            decoder.Decode(out Instruction instruction);
            if (instruction.IsInvalid)
            {
                throw new InvalidOperationException(
                    "The target prologue contains a byte sequence that does not decode to a valid instruction.");
            }

            total += instruction.Length;
        }

        return total;
    }

    private void EmitInstruction(
        Assembler assembler,
        string mnemonic,
        string[] operandText,
        Dictionary<string, Label> labels)
    {
        object?[] operands = new object?[operandText.Length];
        for (int i = 0; i < operandText.Length; i++)
            operands[i] = ParseOperand(operandText[i], labels);

        // Find the fluent Assembler method whose name equals the mnemonic and whose
        // parameters bind to the parsed operands.
        foreach (MethodInfo method in typeof(Assembler).GetMethods(BindingFlags.Public | BindingFlags.Instance))
        {
            if (!string.Equals(method.Name, mnemonic, StringComparison.OrdinalIgnoreCase))
                continue;

            ParameterInfo[] parameters = method.GetParameters();
            if (parameters.Length != operands.Length)
                continue;

            if (TryBind(parameters, operands, out object?[]? boundArgs))
            {
                try
                {
                    method.Invoke(assembler, boundArgs);
                }
                catch (TargetInvocationException ex) when (ex.InnerException is not null)
                {
                    // Surface the real Iced failure rather than the reflection wrapper.
                    throw ex.InnerException;
                }

                return;
            }
        }

        throw new NotSupportedException(
            $"Cannot assemble '{mnemonic}{(operandText.Length > 0 ? " " + string.Join(", ", operandText) : "")}': " +
            "no matching Iced assembler overload for the given operands (registers, immediates and " +
            "labels are supported; memory operands are not).");
    }

    private static bool TryBind(ParameterInfo[] parameters, object?[] operands, out object?[]? boundArgs)
    {
        var args = new object?[parameters.Length];
        for (int i = 0; i < parameters.Length; i++)
        {
            Type paramType = parameters[i].ParameterType;
            object? operand = operands[i];

            switch (operand)
            {
                case Immediate imm when IsNumeric(paramType):
                    // An immediate that overflows this parameter's type means this overload
                    // is the wrong width; return false so a wider overload can be tried
                    // instead of crashing the whole assembly.
                    if (!TryChangeType(imm.Value, paramType, out object? converted))
                    {
                        boundArgs = null;
                        return false;
                    }

                    args[i] = converted;
                    break;

                case not null when paramType.IsInstanceOfType(operand):
                    args[i] = operand;
                    break;

                default:
                    boundArgs = null;
                    return false;
            }
        }

        boundArgs = args;
        return true;
    }

    private static object ParseOperand(string text, Dictionary<string, Label> labels)
    {
        string token = text.Trim();

        if (Registers.TryGetValue(token, out object? register))
            return register;

        if (labels.TryGetValue(token, out Label label))
            return label;

        if (TryParseImmediate(token, out object? value))
            return new Immediate(value!);

        throw new NotSupportedException(
            $"Unrecognized operand '{token}' (expected a register, an immediate, or a label).");
    }

    // Parses an immediate as the narrowest of long/ulong that holds it, boxed. Storing the
    // widest representation lets TryChangeType later narrow it to whatever integer parameter
    // the chosen overload expects — and reject (rather than crash on) values that do not fit.
    private static bool TryParseImmediate(string token, out object? value)
    {
        value = null;
        bool negative = token.StartsWith('-');
        string body = negative ? token[1..] : token;

        if (body.StartsWith("0x", StringComparison.OrdinalIgnoreCase))
        {
            if (!ulong.TryParse(body[2..], NumberStyles.HexNumber, CultureInfo.InvariantCulture, out ulong hex))
                return false;

            value = negative ? -(long)hex : hex;
            return true;
        }

        if (negative)
        {
            if (!long.TryParse(token, NumberStyles.Integer, CultureInfo.InvariantCulture, out long signed))
                return false;

            value = signed;
            return true;
        }

        // Non-negative decimal: prefer long, fall back to ulong for values above long.MaxValue.
        if (long.TryParse(body, NumberStyles.Integer, CultureInfo.InvariantCulture, out long asLong))
            value = asLong;
        else if (ulong.TryParse(body, NumberStyles.Integer, CultureInfo.InvariantCulture, out ulong asULong))
            value = asULong;
        else
            return false;

        return true;
    }

    private static bool TryChangeType(object value, Type targetType, out object? result)
    {
        try
        {
            result = Convert.ChangeType(value, targetType, CultureInfo.InvariantCulture);
            return true;
        }
        catch (Exception ex) when (ex is OverflowException or InvalidCastException or FormatException)
        {
            result = null;
            return false;
        }
    }

    private static bool IsNumeric(Type type) => Type.GetTypeCode(type) is
        TypeCode.SByte or TypeCode.Byte or TypeCode.Int16 or TypeCode.UInt16 or
        TypeCode.Int32 or TypeCode.UInt32 or TypeCode.Int64 or TypeCode.UInt64;

    private static List<(string Mnemonic, string[] Operands)> Tokenize(string text, out List<string> labelNames)
    {
        var result = new List<(string, string[])>();
        labelNames = new List<string>();

        foreach (string rawLine in text.Split('\n'))
        {
            string line = rawLine;

            int comment = line.IndexOf(';');
            if (comment >= 0)
                line = line[..comment];

            line = line.Trim();
            if (line.Length == 0)
                continue;

            // A "name:" prefix is a label definition; keep any instruction that follows it
            // on the same line as a separate entry.
            int colon = line.IndexOf(':');
            if (colon >= 0)
            {
                string labelName = line[..colon].Trim();
                labelNames.Add(labelName);
                result.Add((labelName + ":", Array.Empty<string>()));

                line = line[(colon + 1)..].Trim();
                if (line.Length == 0)
                    continue;
            }

            int space = line.IndexOfAny([' ', '\t']);
            if (space < 0)
            {
                result.Add((line, Array.Empty<string>()));
                continue;
            }

            string mnemonic = line[..space];
            string[] operands = line[(space + 1)..]
                .Split(',', StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries);
            result.Add((mnemonic, operands));
        }

        return result;
    }

    private static Dictionary<string, object> BuildRegisterMap()
    {
        var map = new Dictionary<string, object>(StringComparer.OrdinalIgnoreCase);
        foreach (FieldInfo field in typeof(AssemblerRegisters).GetFields(BindingFlags.Public | BindingFlags.Static))
        {
            object? value = field.GetValue(null);
            if (value is not null)
                map[field.Name] = value;
        }

        return map;
    }

    // A parsed immediate (boxed long or ulong), distinguished from register/label operands
    // so binding can narrow it to whichever integer parameter type the chosen overload
    // expects — or reject it when it does not fit.
    private readonly record struct Immediate(object Value);

    private sealed class ByteListCodeWriter : CodeWriter
    {
        public List<byte> Bytes { get; } = new();

        public override void WriteByte(byte value) => Bytes.Add(value);
    }
}