XFEWarframe

using System.ComponentModel;
using System.Runtime.InteropServices;
using Iced.Intel;

namespace XFEWarframe.Bridge;

internal sealed unsafe class NativeHook
{
    private readonly nint address;
    private readonly byte[] original;
    private readonly byte[] replacement;
    public nint Trampoline { get; }
    private readonly nint allocation;
    private bool enabled;
    public NativeHook(nint address, nint destination) : this(address, destination, false, false) { }
    // Observe arguments, then tail-call the original native function. Its Lua error
    // handler may longjmp: the managed observer must have returned before it runs.
    public static NativeHook Observe(nint address, nint observer) => new(address, observer, true, false);
    private NativeHook(nint address, nint destination, bool observer, bool unused)
    {
        this.address = address;
        var decoder = Decoder.Create(64, new ByteArrayCodeReader(new ReadOnlySpan<byte>((void*)address, 64).ToArray()));
        decoder.IP = (ulong)address;
        var instructions = new List<Instruction>();
        var length = 0;
        while (length < 14)
        {
            decoder.Decode(out var instruction);
            if (instruction.IsInvalid) throw new InvalidDataException("Invalid x64 instruction at hook boundary.");
            instructions.Add(instruction);
            length += instruction.Length;
        }
        original = new ReadOnlySpan<byte>((void*)address, length).ToArray();
        replacement = Enumerable.Repeat((byte)0x90, length).ToArray();
        Trampoline = allocation = AllocateNear(address);
        try
        {
            // Include the return jump in BlockEncoder: it may place branch pointer data after the instructions.
            instructions.Add(Instruction.CreateBranch(Code.Jmp_rel32_64, (ulong)(address + length)));
            var writer = new BufferWriter();
            if (!BlockEncoder.TryEncode(64, new InstructionBlock(writer, instructions, (ulong)Trampoline), out var error, out _))
                throw new InvalidDataException("Cannot relocate x64 prologue: " + error);
            if (writer.Bytes.Count > 2048) throw new InvalidDataException("Trampoline is too large.");
            Marshal.Copy(writer.Bytes.ToArray(), 0, Trampoline, writer.Bytes.Count);
            var target = destination;
            if (observer)
            {
                target = allocation + 2048;
                var stub = ObserverStub(destination, Trampoline);
                Marshal.Copy(stub, 0, target, stub.Length);
            }
            Jump(target).CopyTo(replacement, 0);
            Protect(Trampoline, 4096, 0x20);
            Native.FlushInstructionCache(Native.GetCurrentProcess(), Trampoline, 4096);
        }
        catch { Native.VirtualFree(Trampoline, 0, 0x8000); throw; }
    }
    private static byte[] ObserverStub(nint callback, nint continuation)
    {
        // Entry RSP is 8 mod 16. Save flags and all volatile registers, reserve
        // shadow space plus XMM0..5, and align the callback's stack to 16 bytes.
        var bytes = new List<byte> { 0x9C,0x50,0x51,0x52,0x41,0x50,0x41,0x51,0x41,0x52,0x41,0x53,
            0x48,0x81,0xEC,0x88,0,0,0 };
        for (var i = 0; i < 6; i++) bytes.AddRange([0xF3,0x0F,0x7F,(byte)(0x44 + i * 8),0x24,(byte)(0x20 + i * 16)]);
        bytes.AddRange([0x48,0xB8]); bytes.AddRange(BitConverter.GetBytes((long)callback)); bytes.AddRange([0xFF,0xD0]);
        for (var i = 0; i < 6; i++) bytes.AddRange([0xF3,0x0F,0x6F,(byte)(0x44 + i * 8),0x24,(byte)(0x20 + i * 16)]);
        bytes.AddRange([0x48,0x81,0xC4,0x88,0,0,0,0x41,0x5B,0x41,0x5A,0x41,0x59,0x41,0x58,0x5A,0x59,0x58,0x9D]);
        bytes.AddRange(Jump(continuation));
        return bytes.ToArray();
    }
    // Redirect one audited CALL, retaining the call site's existing Windows x64 ABI.
    // This runs the callback at a precise initialization phase without detouring a
    // shared getter that other ability constructors can call earlier.
    public static NativeHook RedirectCall(nint address, nint destination) => new(address, destination, true);
    private NativeHook(nint address, nint destination, bool callSite)
    {
        this.address = address;
        original = new ReadOnlySpan<byte>((void*)address, 5).ToArray();
        if (original[0] != 0xE8) throw new InvalidDataException("Expected a direct x64 call.");
        Trampoline = address + 5 + BitConverter.ToInt32(original, 1);
        allocation = AllocateNear(address);
        try
        {
            Marshal.Copy(Jump(destination), 0, allocation, 14);
            Protect(allocation, 4096, 0x20);
            Native.FlushInstructionCache(Native.GetCurrentProcess(), allocation, 4096);
            replacement = new byte[5]; replacement[0] = 0xE8;
            BitConverter.TryWriteBytes(replacement.AsSpan(1), checked((int)(allocation - address - 5)));
        }
        catch { Native.VirtualFree(allocation, 0, 0x8000); throw; }
    }
    public void Enable() { Write(address, replacement); enabled = true; }
    public void Rollback()
    {
        if (enabled) Write(address, original);
        Native.VirtualFree(allocation, 0, 0x8000);
    }
    public static void Write(nint address, ReadOnlySpan<byte> bytes)
    {
        var old = Protect(address, (nuint)bytes.Length, 0x40);
        try { bytes.CopyTo(new Span<byte>((void*)address, bytes.Length)); }
        finally { Protect(address, (nuint)bytes.Length, old); }
        if (!Native.FlushInstructionCache(Native.GetCurrentProcess(), address, (nuint)bytes.Length)) throw new Win32Exception();
    }
    private static byte[] Jump(nint target)
    {
        byte[] bytes = [0xFF, 0x25, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        BitConverter.TryWriteBytes(bytes.AsSpan(6), (long)target);
        return bytes;
    }
    private static uint Protect(nint address, nuint size, uint flags)
    {
        if (!Native.VirtualProtect(address, size, flags, out var old)) throw new Win32Exception();
        return old;
    }
    private static nint AllocateNear(nint address)
    {
        var aligned = (long)address & ~0xFFFFL;
        for (long distance = 0x10000; distance < 0x70000000; distance += 0x10000)
        {
            foreach (var candidate in new[] { aligned + distance, aligned - distance })
            {
                if (candidate < 0x10000) continue;
                var result = Native.VirtualAlloc((nint)candidate, 4096, 0x3000, 0x04);
                if (result != 0) return result;
            }
        }
        throw new InvalidOperationException("Cannot allocate a nearby x64 trampoline.");
    }
    private sealed class BufferWriter : CodeWriter
    {
        public List<byte> Bytes { get; } = [];
        public override void WriteByte(byte value) => Bytes.Add(value);
    }
}

internal static partial class Native
{
    [LibraryImport("kernel32", SetLastError = true)] internal static partial nint VirtualAlloc(nint address, nuint size, uint type, uint protect);
    [LibraryImport("kernel32", SetLastError = true)] [return: MarshalAs(UnmanagedType.Bool)] internal static partial bool VirtualFree(nint address, nuint size, uint type);
    [LibraryImport("kernel32", SetLastError = true)] [return: MarshalAs(UnmanagedType.Bool)] internal static partial bool VirtualProtect(nint address, nuint size, uint protect, out uint old);
    [LibraryImport("kernel32", SetLastError = true)] [return: MarshalAs(UnmanagedType.Bool)] internal static partial bool FlushInstructionCache(nint process, nint address, nuint size);
    [LibraryImport("kernel32")] internal static partial nint GetCurrentProcess();
    [LibraryImport("kernel32", EntryPoint = "GetModuleHandleW", StringMarshalling = StringMarshalling.Utf16)] internal static partial nint GetModuleHandle(string? name);
    [LibraryImport("kernel32")] internal static partial void ExitProcess(uint exitCode);
}